Merge pull request #883 from aws/758-reuse-asyncinit

fix: reuse AsyncInitializationWrapper in SpringDelegatingLambdaContainerHandler
aws · May 16, 2024 · 6d4210f · 6d4210f
2 parents 75026f2 + 441fd23
commit 6d4210f
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/...in/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java b/...in/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java
@@ -4,7 +4,10 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.AsyncInitializationWrapper;
 import com.amazonaws.serverless.proxy.InitializationTypeHelper;
+import com.amazonaws.serverless.proxy.internal.InitializableLambdaContainerHandler;
 import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
 import org.springframework.cloud.function.serverless.web.FunctionClassUtils;
 import org.springframework.cloud.function.serverless.web.ServerlessMVC;
@@ -38,25 +41,23 @@
  *
  */
 public class SpringDelegatingLambdaContainerHandler implements RequestStreamHandler {
-
-    private final Class<?>[] startupClasses;
-
     private final ServerlessMVC mvc;
-
     private final ObjectMapper mapper;
-
     private final AwsProxyHttpServletResponseWriter responseWriter;
 
-    public SpringDelegatingLambdaContainerHandler() {
+    public SpringDelegatingLambdaContainerHandler() throws ContainerInitializationException {
         this(new Class[] {FunctionClassUtils.getStartClass()});
     }
 
-    public SpringDelegatingLambdaContainerHandler(Class<?>... startupClasses) {
-        this.startupClasses = startupClasses;
-        this.mvc = ServerlessMVC.INSTANCE(this.startupClasses);
+    public SpringDelegatingLambdaContainerHandler(final Class<?>... startupClasses) throws ContainerInitializationException {
+        SpringDelegatingInitHandler initHandler = new SpringDelegatingInitHandler(startupClasses);
         if (InitializationTypeHelper.isAsyncInitializationDisabled()) {
-    		mvc.waitForContext();
-    	}
+            initHandler.initialize();
+    	} else {
+            AsyncInitializationWrapper asyncInitWrapper = new AsyncInitializationWrapper();
+            asyncInitWrapper.start(initHandler);
+        }
+        this.mvc = initHandler.getMvc();
         this.mapper = new ObjectMapper();
         this.responseWriter = new AwsProxyHttpServletResponseWriter();
     }
@@ -68,4 +69,23 @@ public void handleRequest(InputStream input, OutputStream output, Context lambda
         AwsProxyResponse awsProxyResponse = AwsSpringHttpProcessingUtils.processRequest(httpServletRequest, mvc, responseWriter);
         this.mapper.writeValue(output, awsProxyResponse);
     }
+
+    private static final class SpringDelegatingInitHandler implements InitializableLambdaContainerHandler {
+        private ServerlessMVC mvc;
+        private final Class<?>[] startupClasses;
+
+        public SpringDelegatingInitHandler(final Class<?>... startupClasses) {
+            this.startupClasses = startupClasses;
+        }
+
+        @Override
+        public void initialize() throws ContainerInitializationException {
+            this.mvc = ServerlessMVC.INSTANCE(this.startupClasses);
+            this.mvc.waitForContext();
+        }
+
+        public ServerlessMVC getMvc() {
+            return this.mvc;
+        }
+    }
 }
diff --git a/...va/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandlerTests.java b/...va/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandlerTests.java
@@ -11,6 +11,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.springframework.cloud.function.serverless.web.ServerlessServletContext;
@@ -200,7 +201,7 @@ public class SpringDelegatingLambdaContainerHandlerTests {
 
     private ObjectMapper mapper = new ObjectMapper();
 
-    public void initServletAppTest()  {
+    public void initServletAppTest() throws ContainerInitializationException {
         this.handler = new SpringDelegatingLambdaContainerHandler(ServletApplication.class);
     }