From 17afb9a418f9a66b2a91e3483fa84eed626aaa2f Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Tue, 26 Mar 2024 11:52:47 +0100
Subject: [PATCH] GH-805 Fix web context initialization during snapstart

Also updated to s-c-function-serverless-webb 4.1.1-SNAPSHOT
This commit forces wait for full context initialization if context is created during snapstart creation
---
 aws-serverless-java-container-springboot3/pom.xml             | 2 +-
 .../proxy/spring/SpringDelegatingLambdaContainerHandler.java  | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/aws-serverless-java-container-springboot3/pom.xml b/aws-serverless-java-container-springboot3/pom.xml
index 23ceec081..def9c9015 100644
--- a/aws-serverless-java-container-springboot3/pom.xml
+++ b/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-SNAPSHOT</version>
         </dependency>
         <dependency>
             <groupId>com.amazonaws.serverless</groupId>
diff --git a/aws-serverless-java-container-springboot3/src/main/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java b/aws-serverless-java-container-springboot3/src/main/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java
index f1a3b526e..56ebe8305 100644
--- a/aws-serverless-java-container-springboot3/src/main/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java
+++ b/aws-serverless-java-container-springboot3/src/main/java/com/amazonaws/serverless/proxy/spring/SpringDelegatingLambdaContainerHandler.java
@@ -53,6 +53,10 @@ public SpringDelegatingLambdaContainerHandler() {
     public SpringDelegatingLambdaContainerHandler(Class<?>... startupClasses) {
         this.startupClasses = startupClasses;
         this.mvc = ServerlessMVC.INSTANCE(this.startupClasses);
+        if (System.getenv().containsKey("AWS_LAMBDA_INITIALIZATION_TYPE") 
+        		&& System.getenv().get("AWS_LAMBDA_INITIALIZATION_TYPE").equals("snap-start")) {
+    		mvc.waitForContext();
+    	}
         this.mapper = new ObjectMapper();
         this.responseWriter = new AwsProxyHttpServletResponseWriter();
     }