Copybara import of the project:

--
ca5ae88 by Lachlan Roberts <[email protected]>:
Fixes to completion of RequestManager in HttpMode
Signed-off-by: Lachlan Roberts <[email protected]>
--
e4ac138 by Lachlan Roberts <[email protected]>:
Fixes to completion of RequestManager in HttpMode
Signed-off-by: Lachlan Roberts <[email protected]>
--
668201b by Lachlan Roberts <[email protected]>:
ensure the environment is set for finishRequest
Signed-off-by: Lachlan Roberts <[email protected]>
--
e3c930f by Lachlan Roberts <[email protected]>:
changes from review
Signed-off-by: Lachlan Roberts <[email protected]>
--
c669f96 by Lachlan Roberts <[email protected]>:
resolve conflicts
Signed-off-by: Lachlan Roberts <[email protected]>
PiperOrigin-RevId: 699218224
Change-Id: I15a31a4519057e9af342edd00227a9042b215d05

Author: ludoch

runtime/runtime_impl_jetty12/src/main/java/com/google/apphosting/runtime/jetty/http/JettyHttpHandler.java

@@ -33,13 +33,15 @@
 import com.google.apphosting.runtime.RequestRunner.EagerRunner;
 import com.google.apphosting.runtime.ResponseAPIData;
 import com.google.apphosting.runtime.ServletEngineAdapter;
+import com.google.apphosting.runtime.anyrpc.AnyRpcServerContext;
 import com.google.apphosting.runtime.jetty.AppInfoFactory;
 import com.google.common.flogger.GoogleLogger;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.time.Duration;
 import java.util.concurrent.TimeoutException;
 import org.eclipse.jetty.server.Handler;
+import org.eclipse.jetty.server.HttpStream;
 import org.eclipse.jetty.server.Request;
 import org.eclipse.jetty.server.Response;
 import org.eclipse.jetty.server.Server;
@@ -108,8 +110,12 @@ public boolean handle(Request request, Response response, Callback callback) thr
     ApiProxy.Environment currentEnvironment = ApiProxy.getCurrentEnvironment();
     request.setAttribute(AppEngineConstants.ENVIRONMENT_ATTR, currentEnvironment);
 
+    // Only run code to finish request with the RequestManager once the stream is complete.
+    Request.addCompletionListener(
+        request, t -> finishRequest(currentEnvironment, requestToken, genericResponse, context));
+
     try {
-      handled = dispatchRequest(requestToken, genericRequest, genericResponse, callback);
+      handled = dispatchRequest(requestToken, genericRequest, genericResponse);
       if (handled) {
         callback.succeeded();
       }
@@ -123,27 +129,41 @@ public boolean handle(Request request, Response response, Callback callback) thr
       handled = handleException(ex, requestToken, genericResponse);
       Response.writeError(request, response, callback, ex);
     } finally {
-      requestManager.finishRequest(requestToken);
-    }
-    // Do not put this in a final block.  If we propagate an
-    // exception the callback will be invoked automatically.
-    genericResponse.finishWithResponse(context);
-    // We don't want threads used for background requests to go back
-    // in the thread pool, because users may have stashed references
-    // to them or may be expecting them to exit.  Setting the
-    // interrupt bit causes the pool to drop them.
-    if (genericRequest.getRequestType() == RuntimePb.UPRequest.RequestType.BACKGROUND) {
-      Thread.currentThread().interrupt();
+      // We don't want threads used for background requests to go back
+      // in the thread pool, because users may have stashed references
+      // to them or may be expecting them to exit.  Setting the
+      // interrupt bit causes the pool to drop them.
+      if (genericRequest.getRequestType() == RuntimePb.UPRequest.RequestType.BACKGROUND) {
+        Thread.currentThread().interrupt();
+      }
     }
 
     return handled;
   }
 
+  private void finishRequest(
+      ApiProxy.Environment env,
+      RequestManager.RequestToken requestToken,
+      JettyResponseAPIData response,
+      AnyRpcServerContext context) {
+
+    ApiProxy.Environment oldEnv = ApiProxy.getCurrentEnvironment();
+    try {
+      ApiProxy.setEnvironmentForCurrentThread(env);
+      requestManager.finishRequest(requestToken);
+
+      // Do not put this in a final block.  If we propagate an
+      // exception the callback will be invoked automatically.
+      response.finishWithResponse(context);
+    } finally {
+      ApiProxy.setEnvironmentForCurrentThread(oldEnv);
+    }
+  }
+
   private boolean dispatchRequest(
       RequestManager.RequestToken requestToken,
       JettyRequestAPIData request,
-      JettyResponseAPIData response,
-      Callback callback)
+      JettyResponseAPIData response)
       throws Throwable {
     switch (request.getRequestType()) {
       case SHUTDOWN:
@@ -154,14 +174,13 @@ private boolean dispatchRequest(
         dispatchBackgroundRequest(request, response);
         return true;
       case OTHER:
-        return dispatchServletRequest(request, response, callback);
+        return dispatchServletRequest(request, response);
       default:
         throw new IllegalStateException(request.getRequestType().toString());
     }
   }
 
-  private boolean dispatchServletRequest(
-      JettyRequestAPIData request, JettyResponseAPIData response, Callback callback)
+  private boolean dispatchServletRequest(JettyRequestAPIData request, JettyResponseAPIData response)
       throws Throwable {
     Request jettyRequest = request.getWrappedRequest();
     Response jettyResponse = response.getWrappedResponse();

runtime/runtime_impl_jetty9/src/main/java/com/google/apphosting/runtime/jetty9/JettyHttpHandler.java

@@ -34,6 +34,7 @@
 import com.google.apphosting.runtime.RequestRunner.EagerRunner;
 import com.google.apphosting.runtime.ResponseAPIData;
 import com.google.apphosting.runtime.ServletEngineAdapter;
+import com.google.apphosting.runtime.anyrpc.AnyRpcServerContext;
 import com.google.common.flogger.GoogleLogger;
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -44,8 +45,10 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import org.eclipse.jetty.server.Handler;
+import org.eclipse.jetty.server.HttpChannel;
 import org.eclipse.jetty.server.Request;
 import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.server.handler.HandlerWrapper;
 
 /**
@@ -62,6 +65,9 @@
 public class JettyHttpHandler extends HandlerWrapper {
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
+  public static final String FINISH_REQUEST_ATTRIBUTE =
+      "com.google.apphosting.runtime.jetty9.finishRequestAttribute";
+
   private final boolean passThroughPrivateHeaders;
   private final AppInfoFactory appInfoFactory;
   private final AppVersionKey appVersionKey;
@@ -73,7 +79,8 @@ public JettyHttpHandler(
       ServletEngineAdapter.Config runtimeOptions,
       AppVersion appVersion,
       AppVersionKey appVersionKey,
-      AppInfoFactory appInfoFactory) {
+      AppInfoFactory appInfoFactory,
+      ServerConnector connector) {
     this.passThroughPrivateHeaders = runtimeOptions.passThroughPrivateHeaders();
     this.appInfoFactory = appInfoFactory;
     this.appVersionKey = appVersionKey;
@@ -82,6 +89,7 @@ public JettyHttpHandler(
     ApiProxyImpl apiProxyImpl = (ApiProxyImpl) ApiProxy.getDelegate();
     coordinator = apiProxyImpl.getBackgroundRequestCoordinator();
     requestManager = (RequestManager) apiProxyImpl.getRequestThreadManager();
+    connector.addBean(new CompletionListener());
   }
 
   @Override
@@ -113,6 +121,10 @@ public void handle(
     ApiProxy.Environment currentEnvironment = ApiProxy.getCurrentEnvironment();
     request.setAttribute(AppEngineConstants.ENVIRONMENT_ATTR, currentEnvironment);
 
+    Runnable finishRequest =
+        () -> finishRequest(currentEnvironment, requestToken, genericResponse, context);
+    baseRequest.setAttribute(FINISH_REQUEST_ATTRIBUTE, finishRequest);
+
     try {
       dispatchRequest(target, requestToken, genericRequest, genericResponse);
       if (!baseRequest.isHandled()) {
@@ -128,17 +140,32 @@ public void handle(
       handleException(ex, requestToken, genericResponse);
       response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, ex.getMessage());
     } finally {
-      requestManager.finishRequest(requestToken);
+      // We don't want threads used for background requests to go back
+      // in the thread pool, because users may have stashed references
+      // to them or may be expecting them to exit.  Setting the
+      // interrupt bit causes the pool to drop them.
+      if (genericRequest.getRequestType() == RuntimePb.UPRequest.RequestType.BACKGROUND) {
+        Thread.currentThread().interrupt();
+      }
     }
-    // Do not put this in a final block.  If we propagate an
-    // exception the callback will be invoked automatically.
-    genericResponse.finishWithResponse(context);
-    // We don't want threads used for background requests to go back
-    // in the thread pool, because users may have stashed references
-    // to them or may be expecting them to exit.  Setting the
-    // interrupt bit causes the pool to drop them.
-    if (genericRequest.getRequestType() == RuntimePb.UPRequest.RequestType.BACKGROUND) {
-      Thread.currentThread().interrupt();
+  }
+
+  private void finishRequest(
+      ApiProxy.Environment env,
+      RequestManager.RequestToken requestToken,
+      JettyResponseAPIData response,
+      AnyRpcServerContext context) {
+
+    ApiProxy.Environment oldEnv = ApiProxy.getCurrentEnvironment();
+    try {
+      ApiProxy.setEnvironmentForCurrentThread(env);
+      requestManager.finishRequest(requestToken);
+
+      // Do not put this in a final block.  If we propagate an
+      // exception the callback will be invoked automatically.
+      response.finishWithResponse(context);
+    } finally {
+      ApiProxy.setEnvironmentForCurrentThread(oldEnv);
     }
   }
 
@@ -288,4 +315,15 @@ private String getBackgroundRequestId(JettyRequestAPIData upRequest) {
     }
     return backgroundRequestId;
   }
+
+  private static class CompletionListener implements HttpChannel.Listener {
+    @Override
+    public void onComplete(Request request) {
+      Runnable finishRequest =
+          (Runnable) request.getAttribute(JettyHttpHandler.FINISH_REQUEST_ATTRIBUTE);
+      if (finishRequest != null) {
+        finishRequest.run();
+      }
+    }
+  }
 }

runtime/runtime_impl_jetty9/src/main/java/com/google/apphosting/runtime/jetty9/JettyServletEngineAdapter.java

@@ -43,6 +43,7 @@
 import java.util.Optional;
 import javax.servlet.ServletException;
 import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.server.handler.SizeLimitHandler;
 import org.eclipse.jetty.util.thread.QueuedThreadPool;
 
@@ -143,10 +144,11 @@ public void start(String serverInfo, ServletEngineAdapter.Config runtimeOptions)
           logger.atInfo().log("Using HTTP_CONNECTOR_MODE to bypass RPC");
           appVersionKey = AppVersionKey.fromAppInfo(appinfo);
           AppVersion appVersion = appVersionHandlerMap.getAppVersion(appVersionKey);
+          ServerConnector connector = JettyHttpProxy.newConnector(server, runtimeOptions);
+          server.addConnector(connector);
           server.insertHandler(
-              new JettyHttpHandler(runtimeOptions, appVersion, appVersionKey, appInfoFactory));
+              new JettyHttpHandler(runtimeOptions, appVersion, appVersionKey, appInfoFactory, connector));
           JettyHttpProxy.insertHandlers(server, ignoreResponseSizeLimit);
-          server.addConnector(JettyHttpProxy.newConnector(server, runtimeOptions));
         } else {
           server.setAttribute(
               "com.google.apphosting.runtime.jetty9.appYaml",

runtime/testapps/src/main/resources/com/google/apphosting/runtime/jetty9/gzipapp/jetty94/WEB-INF/appengine-web.xml

@@ -20,6 +20,6 @@
   <application>gzip</application>
   <threadsafe>true</threadsafe>
   <system-properties>
-    <property name="appengine.use.EE8" value="true"/>
+    <property name="appengine.use.EE8" value="false"/>
   </system-properties>
 </appengine-web-app>