Set the Thread Context ClassLoader while executing function code.

Users justifiably expect that the TCCL will be the loader of their code, and lots of things fail to work if that is not true.

PiperOrigin-RevId: 301824288

Author: eamonnmcmanus

invoker/core/src/main/java/com/google/cloud/functions/invoker/NewBackgroundFunctionExecutor.java

@@ -168,14 +168,18 @@ private static Map<String, Object> httpHeaderMap(HttpServletRequest req) {
    *     unmarshalling this event, if it is a CloudEvent.
    */
   private abstract static class FunctionExecutor<CloudEventDataT> {
-    private final String functionName;
+    private final Class<?> functionClass;
 
-    FunctionExecutor(String functionName) {
-      this.functionName = functionName;
+    FunctionExecutor(Class<?> functionClass) {
+      this.functionClass = functionClass;
     }
 
     final String functionName() {
-      return functionName;
+      return functionClass.getCanonicalName();
+    }
+
+    final ClassLoader functionClassLoader() {
+      return functionClass.getClassLoader();
     }
 
     abstract void serviceLegacyEvent(HttpServletRequest req)
@@ -193,7 +197,7 @@ private static class RawFunctionExecutor extends FunctionExecutor<Map<?, ?>> {
     private final RawBackgroundFunction function;
 
     RawFunctionExecutor(RawBackgroundFunction function) {
-      super(function.getClass().getCanonicalName());
+      super(function.getClass());
       this.function = function;
     }
 
@@ -234,7 +238,7 @@ private static class TypedFunctionExecutor<T> extends FunctionExecutor<T> {
     private final BackgroundFunction<T> function;
 
     private TypedFunctionExecutor(Type type, BackgroundFunction<T> function) {
-      super(function.getClass().getCanonicalName());
+      super(function.getClass());
       this.type = type;
       this.function = function;
     }
@@ -287,7 +291,9 @@ Class<T> cloudEventDataType() {
   @Override
   public void service(HttpServletRequest req, HttpServletResponse res) throws IOException {
     String contentType = req.getContentType();
+    ClassLoader oldContextLoader = Thread.currentThread().getContextClassLoader();
     try {
+      Thread.currentThread().setContextClassLoader(functionExecutor.functionClassLoader());
       if (contentType != null && contentType.startsWith("application/cloudevents+json")) {
         serviceCloudEvent(req, CloudEventKind.STRUCTURED);
       } else if (req.getHeader("ce-specversion") != null) {
@@ -299,6 +305,8 @@ public void service(HttpServletRequest req, HttpServletResponse res) throws IOEx
     } catch (Throwable t) {
       res.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
       logger.log(Level.WARNING, "Failed to execute " + functionExecutor.functionName(), t);
+    } finally {
+      Thread.currentThread().setContextClassLoader(oldContextLoader);
     }
   }
 

invoker/core/src/main/java/com/google/cloud/functions/invoker/NewHttpFunctionExecutor.java

@@ -61,7 +61,9 @@ public static NewHttpFunctionExecutor forClass(Class<?> functionClass) {
   public void service(HttpServletRequest req, HttpServletResponse res) {
     HttpRequestImpl reqImpl = new HttpRequestImpl(req);
     HttpResponseImpl respImpl = new HttpResponseImpl(res);
+    ClassLoader oldContextLoader = Thread.currentThread().getContextClassLoader();
     try {
+      Thread.currentThread().setContextClassLoader(function.getClass().getClassLoader());
       function.service(reqImpl, respImpl);
     } catch (Throwable t) {
       // TODO(b/146510646): this should be logged properly as an exception, but that currently
@@ -71,6 +73,7 @@ public void service(HttpServletRequest req, HttpServletResponse res) {
       t.printStackTrace();
       res.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
     } finally {
+      Thread.currentThread().setContextClassLoader(oldContextLoader);
       try {
         // We can't use HttpServletResponse.flushBuffer() because we wrap the PrintWriter
         // returned by HttpServletResponse in our own BufferedWriter to match our API.

invoker/testfunction/src/test/java/com/example/functionjar/Checker.java

@@ -18,6 +18,15 @@
 
 class Checker {
   void serviceOrAssert(String runtimeClassName) {
+    // Check that the context class loader is the loader that loaded this class.
+    if (getClass().getClassLoader() != Thread.currentThread().getContextClassLoader()) {
+      throw new AssertionError(
+          String.format(
+              "ClassLoader mismatch: mine %s; context %s",
+              getClass().getClassLoader(),
+              Thread.currentThread().getContextClassLoader()));
+    }
+
     ClassLoader myLoader = getClass().getClassLoader();
     Class<Template> templateClass = Template.class;
     ClassLoader templateLoader = templateClass.getClassLoader();