Fixes for ContextScopeListener when using RPC

Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
GoogleCloudPlatform · May 8, 2024 · 4335097 · 4335097
1 parent 7430571
commit 4335097
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 29 deletions.
diff --git a/..._jetty12/src/main/java/com/google/apphosting/runtime/jetty/JettyServletEngineAdapter.java b/..._jetty12/src/main/java/com/google/apphosting/runtime/jetty/JettyServletEngineAdapter.java
@@ -15,9 +15,7 @@
  */
 package com.google.apphosting.runtime.jetty;
 
-import static com.google.apphosting.runtime.jetty.AppEngineConstants.HTTP_CONNECTOR_MODE;
-import static java.nio.charset.StandardCharsets.UTF_8;
-
+import com.google.apphosting.api.ApiProxy;
 import com.google.apphosting.base.AppVersionKey;
 import com.google.apphosting.base.protos.AppinfoPb;
 import com.google.apphosting.base.protos.EmptyMessage;
@@ -36,12 +34,6 @@
 import com.google.apphosting.utils.config.AppEngineConfigException;
 import com.google.apphosting.utils.config.AppYaml;
 import com.google.common.flogger.GoogleLogger;
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.InputStreamReader;
-import java.util.Objects;
-import java.util.concurrent.ExecutionException;
 import org.eclipse.jetty.http.CookieCompliance;
 import org.eclipse.jetty.http.UriCompliance;
 import org.eclipse.jetty.server.HttpConfiguration;
@@ -52,6 +44,16 @@
 import org.eclipse.jetty.util.resource.ResourceFactory;
 import org.eclipse.jetty.util.thread.QueuedThreadPool;
 
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.InputStreamReader;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+
+import static com.google.apphosting.runtime.jetty.AppEngineConstants.HTTP_CONNECTOR_MODE;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
 /**
  * This is an implementation of ServletEngineAdapter that uses the third-party Jetty servlet engine.
  *
@@ -264,6 +266,7 @@ public void serviceRequest(UPRequest upRequest, MutableUpResponse upResponse) th
     }
     DelegateRpcExchange rpcExchange = new DelegateRpcExchange(upRequest, upResponse);
     rpcExchange.setAttribute(JettyConstants.APP_VERSION_KEY_REQUEST_ATTR, appVersionKey);
+    rpcExchange.setAttribute(AppEngineConstants.ENVIRONMENT_ATTR, ApiProxy.getCurrentEnvironment());
     rpcConnector.service(rpcExchange);
     try {
       rpcExchange.awaitResponse();

diff --git a/.../src/main/java/com/google/apphosting/runtime/jetty/ee10/EE10AppVersionHandlerFactory.java b/.../src/main/java/com/google/apphosting/runtime/jetty/ee10/EE10AppVersionHandlerFactory.java
@@ -29,10 +29,6 @@
 import jakarta.servlet.ServletException;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
-import java.io.File;
-import java.io.IOException;
-import java.io.PrintWriter;
-import javax.servlet.jsp.JspFactory;
 import org.eclipse.jetty.ee10.annotations.AnnotationConfiguration;
 import org.eclipse.jetty.ee10.quickstart.QuickStartConfiguration;
 import org.eclipse.jetty.ee10.servlet.Dispatcher;
@@ -52,6 +48,13 @@
 import org.eclipse.jetty.server.handler.ContextHandler;
 import org.eclipse.jetty.util.Callback;
 
+import javax.servlet.jsp.JspFactory;
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintWriter;
+
+import static com.google.apphosting.runtime.jetty.AppEngineConstants.HTTP_CONNECTOR_MODE;
+
 /**
  * {@code AppVersionHandlerFactory} implements a {@code Handler} for a given {@code AppVersionKey}.
  */
@@ -189,25 +192,28 @@ private org.eclipse.jetty.server.Handler doCreateHandler(AppVersion appVersion)
       EE10SessionManagerHandler.create(builder.build());
       // Pass the AppVersion on to any of our servlets (e.g. ResourceFileServlet).
       context.setAttribute(JettyConstants.APP_VERSION_CONTEXT_ATTR, appVersion);
-      context.addEventListener(
+
+      if (Boolean.getBoolean(HTTP_CONNECTOR_MODE)) {
+        context.addEventListener(
           new ContextHandler.ContextScopeListener() {
             @Override
             public void enterScope(Context context, Request request) {
               if (request != null) {
                 ApiProxy.Environment environment =
-                    (ApiProxy.Environment)
-                        request.getAttribute(AppEngineConstants.ENVIRONMENT_ATTR);
+                        (ApiProxy.Environment)
+                                request.getAttribute(AppEngineConstants.ENVIRONMENT_ATTR);
                 if (environment != null) ApiProxy.setEnvironmentForCurrentThread(environment);
               }
             }
 
             @Override
             public void exitScope(Context context, Request request) {
-//              if (request != null) {
-//                ApiProxy.clearEnvironmentForCurrentThread();
-//              }
+              if (request != null) {
+                ApiProxy.clearEnvironmentForCurrentThread();
+              }
             }
           });
+      }
       return context;
     } catch (Exception ex) {
       throw new ServletException(ex);

diff --git a/...12/src/main/java/com/google/apphosting/runtime/jetty/ee8/EE8AppVersionHandlerFactory.java b/...12/src/main/java/com/google/apphosting/runtime/jetty/ee8/EE8AppVersionHandlerFactory.java
@@ -46,6 +46,8 @@
 import java.io.IOException;
 import java.io.PrintWriter;
 
+import static com.google.apphosting.runtime.jetty.AppEngineConstants.HTTP_CONNECTOR_MODE;
+
 /**
  * {@code AppVersionHandlerFactory} implements a {@code Handler} for a given {@code AppVersionKey}.
  */
@@ -205,17 +207,18 @@ private org.eclipse.jetty.server.Handler doCreateHandler(AppVersion appVersion)
       // Pass the AppVersion on to any of our servlets (e.g. ResourceFileServlet).
       context.setAttribute(JettyConstants.APP_VERSION_CONTEXT_ATTR, appVersion);
 
-      context.addEventListener(
+      if (Boolean.getBoolean(HTTP_CONNECTOR_MODE)) {
+        context.addEventListener(
           new ContextHandler.ContextScopeListener() {
             @Override
             public void enterScope(
-                ContextHandler.APIContext context,
-                org.eclipse.jetty.ee8.nested.Request request,
-                Object reason) {
+                    ContextHandler.APIContext context,
+                    org.eclipse.jetty.ee8.nested.Request request,
+                    Object reason) {
               if (request != null) {
                 ApiProxy.Environment environment =
-                    (ApiProxy.Environment)
-                        request.getAttribute(AppEngineConstants.ENVIRONMENT_ATTR);
+                        (ApiProxy.Environment)
+                                request.getAttribute(AppEngineConstants.ENVIRONMENT_ATTR);
                 if (environment != null) {
                   ApiProxy.setEnvironmentForCurrentThread(environment);
                 }
@@ -224,12 +227,13 @@ public void enterScope(
 
             @Override
             public void exitScope(
-                ContextHandler.APIContext context, org.eclipse.jetty.ee8.nested.Request request) {
-//              if (request != null) {
-//                ApiProxy.clearEnvironmentForCurrentThread();
-//              }
+                    ContextHandler.APIContext context, org.eclipse.jetty.ee8.nested.Request request) {
+              if (request != null) {
+                ApiProxy.clearEnvironmentForCurrentThread();
+              }
             }
           });
+      }
 
       return context.get();
     } catch (Exception ex) {