fix: Reinitialization of Spring context when using SnapStart #805, fixes #814

Author: deki

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/AsyncInitializationWrapper.java

@@ -41,17 +41,13 @@ public class AsyncInitializationWrapper extends InitializationWrapper {
 
     private static final int DEFAULT_INIT_GRACE_TIME_MS = 150;
     private static final String INIT_GRACE_TIME_ENVIRONMENT_VARIABLE_NAME = "AWS_SERVERLESS_JAVA_CONTAINER_INIT_GRACE_TIME";
-    private static final String INITIALIZATION_TYPE_ENVIRONMENT_VARIABLE_NAME = "AWS_LAMBDA_INITIALIZATION_TYPE";
-    private static final String INITIALIZATION_TYPE_ON_DEMAND = "on-demand";
-    private static final String INITIALIZATION_TYPE = System.getenv().getOrDefault(INITIALIZATION_TYPE_ENVIRONMENT_VARIABLE_NAME,INITIALIZATION_TYPE_ON_DEMAND);
-    private static final boolean ASYNC_INIT_DISABLED = !INITIALIZATION_TYPE.equals(INITIALIZATION_TYPE_ON_DEMAND);
     private static final int INIT_GRACE_TIME_MS = Integer.parseInt(System.getenv().getOrDefault(
             INIT_GRACE_TIME_ENVIRONMENT_VARIABLE_NAME, Integer.toString(DEFAULT_INIT_GRACE_TIME_MS)));
     private static final int LAMBDA_MAX_INIT_TIME_MS = 10_000;
 
     private CountDownLatch initializationLatch;
     private final long actualStartTime;
-    private Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
+    private final Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
 
 
     /**
@@ -73,8 +69,8 @@ public AsyncInitializationWrapper() {
 
     @Override
     public void start(InitializableLambdaContainerHandler handler) throws ContainerInitializationException {
-        if(ASYNC_INIT_DISABLED){
-            log.info("Async init disabled due to \"{}\" initialization", INITIALIZATION_TYPE);
+        if (InitializationTypeHelper.isAsyncInitializationDisabled()){
+            log.info("Async init disabled due to \"{}\" initialization", InitializationTypeHelper.getInitializationType());
             super.start(handler);
             return;
         }
@@ -107,7 +103,7 @@ public long getActualStartTimeMs() {
 
     @Override
     public CountDownLatch getInitializationLatch() {
-        if (ASYNC_INIT_DISABLED){
+        if (InitializationTypeHelper.isAsyncInitializationDisabled()){
             return super.getInitializationLatch();
         }
         return initializationLatch;
@@ -116,7 +112,7 @@ public CountDownLatch getInitializationLatch() {
     private static class AsyncInitializer implements Runnable {
         private final InitializableLambdaContainerHandler handler;
         private CountDownLatch initLatch;
-        private Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
+        private final Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
 
         AsyncInitializer(CountDownLatch latch, InitializableLambdaContainerHandler h) {
             initLatch = latch;

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/InitializationTypeHelper.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package com.amazonaws.serverless.proxy;
+
+/**
+ * Utility class that helps determine the initialization type
+ */
+public final class InitializationTypeHelper {
+
+    private static final String INITIALIZATION_TYPE_ENVIRONMENT_VARIABLE_NAME = "AWS_LAMBDA_INITIALIZATION_TYPE";
+    private static final String INITIALIZATION_TYPE_ON_DEMAND = "on-demand";
+    private static final String INITIALIZATION_TYPE = System.getenv().getOrDefault(INITIALIZATION_TYPE_ENVIRONMENT_VARIABLE_NAME,
+            INITIALIZATION_TYPE_ON_DEMAND);
+    private static final boolean ASYNC_INIT_DISABLED = !INITIALIZATION_TYPE.equals(INITIALIZATION_TYPE_ON_DEMAND);
+
+    public static boolean isAsyncInitializationDisabled() {
+        return ASYNC_INIT_DISABLED;
+    }
+
+    public static String getInitializationType() {
+        return INITIALIZATION_TYPE;
+    }
+}

aws-serverless-java-container-springboot3/pom.xml

@@ -25,7 +25,7 @@
         <dependency>
             <groupId>org.springframework.cloud</groupId>
             <artifactId>spring-cloud-function-serverless-web</artifactId>
-            <version>4.0.6</version>
+            <version>4.1.1</version>
         </dependency>
         <dependency>
             <groupId>com.amazonaws.serverless</groupId>

aws-serverless-java-container-springboot3/src/main/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java

@@ -4,6 +4,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 
+import com.amazonaws.serverless.proxy.InitializationTypeHelper;
 import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
 import org.springframework.cloud.function.serverless.web.FunctionClassUtils;
 import org.springframework.cloud.function.serverless.web.ServerlessMVC;
@@ -53,6 +54,9 @@ public SpringDelegatingLambdaContainerHandler() {
     public SpringDelegatingLambdaContainerHandler(Class<?>... startupClasses) {
         this.startupClasses = startupClasses;
         this.mvc = ServerlessMVC.INSTANCE(this.startupClasses);
+        if (InitializationTypeHelper.isAsyncInitializationDisabled()) {
+    		mvc.waitForContext();
+    	}
         this.mapper = new ObjectMapper();
         this.responseWriter = new AwsProxyHttpServletResponseWriter();
     }