Refactor: Move Redispatch Logic from Spring Boot 3 and Spring Packages to Core

Author: leekib

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsLambdaServletContainerHandler.java

@@ -152,13 +152,25 @@ protected void doFilter(HttpServletRequest request, HttpServletResponse response
 
         FilterChain chain = getFilterChain(request, servlet);
         chain.doFilter(request, response);
-
+        if(requiresAsyncReDispatch(request)) {
+            chain = getFilterChain(request, servlet);
+            chain.doFilter(request, response);
+        }
         // if for some reason the response wasn't flushed yet, we force it here unless it's being processed asynchronously (WebFlux)
         if (!response.isCommitted() && request.getDispatcherType() != DispatcherType.ASYNC) {
             response.flushBuffer();
         }
     }
 
+    private boolean requiresAsyncReDispatch(HttpServletRequest request) {
+        if (request.isAsyncStarted()) {
+            AsyncContext asyncContext = request.getAsyncContext();
+            return asyncContext instanceof AwsAsyncContext
+                    && ((AwsAsyncContext) asyncContext).isDispatchStarted();
+        }
+        return false;
+    }
+
     @Override
     public void initialize() throws ContainerInitializationException {
         // we expect all servlets to be wrapped in an AwsServletRegistration

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

@@ -20,7 +20,6 @@
 import com.amazonaws.serverless.proxy.internal.servlet.*;
 import com.amazonaws.serverless.proxy.model.HttpApiV2ProxyRequest;
 import com.amazonaws.services.lambda.runtime.Context;
-import jakarta.servlet.AsyncContext;
 import org.springframework.web.context.ConfigurableWebApplicationContext;
 import org.springframework.web.servlet.DispatcherServlet;
 
@@ -161,21 +160,9 @@ protected void handleRequest(HttpServletRequest containerRequest, AwsHttpServlet
         // process filters
         Servlet reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
         doFilter(containerRequest, containerResponse, reqServlet);
-        if(requiresAsyncReDispatch(containerRequest)) {
-            reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
-            doFilter(containerRequest, containerResponse, reqServlet);
-        }
         Timer.stop("SPRING_HANDLE_REQUEST");
     }
 
-    private boolean requiresAsyncReDispatch(HttpServletRequest request) {
-        if (request.isAsyncStarted()) {
-            AsyncContext asyncContext = request.getAsyncContext();
-            return asyncContext instanceof AwsAsyncContext
-                    && ((AwsAsyncContext) asyncContext).isDispatchStarted();
-        }
-        return false;
-    }
 
     @Override
     public void initialize()

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

@@ -22,7 +22,6 @@
 import com.amazonaws.serverless.proxy.spring.embedded.ServerlessReactiveServletEmbeddedServerFactory;
 import com.amazonaws.serverless.proxy.spring.embedded.ServerlessServletEmbeddedServerFactory;
 import com.amazonaws.services.lambda.runtime.Context;
-import jakarta.servlet.AsyncContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.boot.WebApplicationType;
@@ -173,22 +172,9 @@ protected void handleRequest(HttpServletRequest containerRequest, AwsHttpServlet
             ((AwsHttpServletRequest)containerRequest).setResponse(containerResponse);
         }
         doFilter(containerRequest, containerResponse, reqServlet);
-        if(requiresAsyncReDispatch(containerRequest)) {
-            reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
-            doFilter(containerRequest, containerResponse, reqServlet);
-        }
         Timer.stop("SPRINGBOOT2_HANDLE_REQUEST");
     }
 
-    private boolean requiresAsyncReDispatch(HttpServletRequest request) {
-        if (request.isAsyncStarted()) {
-            AsyncContext asyncContext = request.getAsyncContext();
-            return asyncContext instanceof AwsAsyncContext
-                    && ((AwsAsyncContext) asyncContext).isDispatchStarted();
-        }
-        return false;
-    }
-
 
     @Override
     public void initialize()