Added new InitializationWrapper and Async implementation to address #210 and #234. Also added a new builder object that makes it easier to construct container handlers

Author: sapessi

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

@@ -0,0 +1,103 @@
+package com.amazonaws.serverless.proxy;
+
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.internal.LambdaContainerHandler;
+import com.amazonaws.services.lambda.runtime.Context;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Instant;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * An async implementation of the <code>InitializationWrapper</code> interface. This initializer calls the
+ * {@link LambdaContainerHandler#initialize()} in a separate thread. Then uses a latch to wait for the maximum Lambda
+ * initialization time of 10 seconds, if the <code>initialize</code> method takes longer than 10 seconds to return, the
+ * {@link #start(LambdaContainerHandler)} returns control to the caller and lets the initialization thread continue in
+ * the background. The {@link LambdaContainerHandler#proxy(Object, Context)} automatically waits for the latch of the
+ * initializer to be released.
+ *
+ * The constructor of this class expects an epoch long. This is meant to be as close as possible to the time the Lambda
+ * function actually started. In most cases, the first action in the constructor of the handler class should be to populate
+ * this long value ({@code Instant.now().toEpochMs();}). This class uses the value to estimate how much of the init 10
+ * seconds has already been used up.
+ */
+public class AsyncInitializationWrapper extends InitializationWrapper {
+    private static final int INIT_GRACE_TIME_MS = 250;
+    private static final int LAMBDA_MAX_INIT_TIME_MS = 10_000;
+
+    private CountDownLatch initializationLatch;
+    private long actualStartTime;
+    private Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
+
+    /**
+     * Creates a new instance of the async initializer.
+     * @param startTime The epoch ms start time of the Lambda function, this should be measured as close as possible to
+     *                  the initialization of the function.
+     */
+    public AsyncInitializationWrapper(long startTime) {
+        actualStartTime = startTime;
+    }
+
+    @Override
+    public void start(LambdaContainerHandler handler) throws ContainerInitializationException {
+        initializationLatch = new CountDownLatch(1);
+        AsyncInitializer initializer = new AsyncInitializer(initializationLatch, handler);
+        Thread initThread = new Thread(initializer);
+        initThread.start();
+        try {
+            long curTime = Instant.now().toEpochMilli();
+            // account for the time it took to call the various constructors with the actual start time + a grace of 500ms
+            long awaitTime = LAMBDA_MAX_INIT_TIME_MS - (curTime - actualStartTime) - INIT_GRACE_TIME_MS;
+            log.info("Async initialization will wait for " + awaitTime + "ms");
+            if (!initializationLatch.await(awaitTime, TimeUnit.MILLISECONDS)) {
+                log.info("Initialization took longer than " + LAMBDA_MAX_INIT_TIME_MS + ", setting new CountDownLatch and " +
+                        "continuing in event handler");
+                initializationLatch = new CountDownLatch(1);
+                initializer.replaceLatch(initializationLatch);
+            }
+        } catch (InterruptedException e) {
+            // at the moment we assume that this happened because of a timeout since the init thread calls System.exit
+            // when an exception is thrown.
+            throw new ContainerInitializationException("Container initialization interrupted", e);
+        }
+    }
+
+    @Override
+    public CountDownLatch getInitializationLatch() {
+        return initializationLatch;
+    }
+
+    private static class AsyncInitializer implements Runnable {
+        private LambdaContainerHandler handler;
+        private CountDownLatch initLatch;
+        private Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);
+
+        AsyncInitializer(CountDownLatch latch, LambdaContainerHandler h) {
+            initLatch = latch;
+            handler = h;
+        }
+
+        synchronized void replaceLatch(CountDownLatch newLatch) {
+            initLatch = newLatch;
+        }
+
+        @Override
+        @SuppressFBWarnings("DM_EXIT")
+        public void run() {
+            log.info("Starting async initializer");
+            try {
+                handler.initialize();
+            } catch (ContainerInitializationException e) {
+                log.error("Failed to initialize container handler", e);
+                // we cannot return the exception so we crash the whole kaboodle here
+                System.exit(1);
+            }
+            synchronized(this) {
+                initLatch.countDown();
+            }
+        }
+    }
+}

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

@@ -0,0 +1,34 @@
+package com.amazonaws.serverless.proxy;
+
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.internal.LambdaContainerHandler;
+
+import java.util.concurrent.CountDownLatch;
+
+/**
+ * This class is in charge of initializing a {@link LambdaContainerHandler}.
+ * In most cases, this means calling the {@link LambdaContainerHandler#initialize()} method. Some implementations may
+ * require additional initialization steps, in this case implementations should provide their own
+ * <code>InitializationWrapper</code>. This library includes an async implementation of this class
+ * {@link AsyncInitializationWrapper} for frameworks that are likely to take longer than 10 seconds to start.
+ */
+public class InitializationWrapper {
+    /**
+     * This is the main entry point. Container handler builder and the static <code>getAwsProxyHandler()</code> methods
+     * of the various implementations will call this to initialize the underlying framework
+     * @param handler The container handler to be initializer
+     * @throws ContainerInitializationException If anything goes wrong during container initialization.
+     */
+    public void start(LambdaContainerHandler handler) throws ContainerInitializationException {
+        handler.initialize();
+    }
+
+    /**
+     * Asynchronous implementations of the framework should return a latch that the container handler can use to decide
+     * whether it can start handling events. Synchronous implementations of this interface should return <code>null</code>.
+     * @return An initialized latch if the underlying container is starting in a separate thread, null otherwise.
+     */
+    public CountDownLatch getInitializationLatch() {
+        return null;
+    }
+}

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

@@ -14,13 +14,9 @@
 
 
 import com.amazonaws.serverless.exceptions.ContainerInitializationException;
-import com.amazonaws.serverless.proxy.LogFormatter;
+import com.amazonaws.serverless.proxy.*;
 import com.amazonaws.serverless.proxy.internal.servlet.ApacheCombinedServletLogFormatter;
 import com.amazonaws.serverless.proxy.model.ContainerConfig;
-import com.amazonaws.serverless.proxy.ExceptionHandler;
-import com.amazonaws.serverless.proxy.RequestReader;
-import com.amazonaws.serverless.proxy.ResponseWriter;
-import com.amazonaws.serverless.proxy.SecurityContextWriter;
 import com.amazonaws.services.lambda.runtime.Context;
 
 import com.fasterxml.jackson.core.JsonParseException;
@@ -38,6 +34,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 
 /**
@@ -67,6 +64,7 @@ public abstract class LambdaContainerHandler<RequestType, ResponseType, Containe
     private ExceptionHandler<ResponseType> exceptionHandler;
     private Class<RequestType> requestTypeClass;
     private Class<ResponseType> responseTypeClass;
+    private InitializationWrapper initializationWrapper;
 
     protected Context lambdaContext;
     private LogFormatter<ContainerRequestType, ContainerResponseType> logFormatter;
@@ -97,16 +95,28 @@ protected LambdaContainerHandler(Class<RequestType> requestClass,
                                      RequestReader<RequestType, ContainerRequestType> requestReader,
                                      ResponseWriter<ContainerResponseType, ResponseType> responseWriter,
                                      SecurityContextWriter<RequestType> securityContextWriter,
-                                     ExceptionHandler<ResponseType> exceptionHandler) {
+                                     ExceptionHandler<ResponseType> exceptionHandler,
+                                     InitializationWrapper init) {
         log.info("Starting Lambda Container Handler");
         requestTypeClass = requestClass;
         responseTypeClass = responseClass;
         this.requestReader = requestReader;
         this.responseWriter = responseWriter;
         this.securityContextWriter = securityContextWriter;
         this.exceptionHandler = exceptionHandler;
+        initializationWrapper = init;
         objectReader = getObjectMapper().readerFor(requestTypeClass);
         objectWriter = getObjectMapper().writerFor(responseTypeClass);
+
+    }
+
+    protected LambdaContainerHandler(Class<RequestType> requestClass,
+                                     Class<ResponseType> responseClass,
+                                     RequestReader<RequestType, ContainerRequestType> requestReader,
+                                     ResponseWriter<ContainerResponseType, ResponseType> responseWriter,
+                                     SecurityContextWriter<RequestType> securityContextWriter,
+                                     ExceptionHandler<ResponseType> exceptionHandler) {
+        this(requestClass, responseClass, requestReader, responseWriter, securityContextWriter, exceptionHandler, new InitializationWrapper());
     }
 
 
@@ -131,6 +141,23 @@ public static ObjectMapper getObjectMapper() {
         return objectMapper;
     }
 
+    /**
+     * Returns the initialization wrapper this container handler will monitor to handle events
+     * @return The initialization wrapper that was passed to the constructor and this instance will use to decide
+     *         whether it can start handling events.
+     */
+    public InitializationWrapper getInitializationWrapper() {
+        return initializationWrapper;
+    }
+
+    /**
+     * Sets a new initialization wrapper.
+     * @param wrapper The wrapper this instance will use to decide whether it can start handling events.
+     */
+    public void setInitializationWrapper(InitializationWrapper wrapper) {
+        initializationWrapper = wrapper;
+    }
+
     /**
      * Configures the library to strip a base path from incoming requests before passing them on to the wrapped
      * framework. This was added in response to issue #34 (https://github.com/awslabs/aws-serverless-java-container/issues/34).
@@ -174,6 +201,13 @@ public ResponseType proxy(RequestType request, Context context) {
             ContainerRequestType containerRequest = requestReader.readRequest(request, securityContext, context, config);
             ContainerResponseType containerResponse = getContainerResponse(containerRequest, latch);
 
+            if (initializationWrapper != null && initializationWrapper.getInitializationLatch() != null) {
+                // we let the potential InterruptedException bubble up
+                if (!initializationWrapper.getInitializationLatch().await(config.getInitializationTimeout(), TimeUnit.MILLISECONDS)) {
+                    throw new ContainerInitializationException("Could not initialize framework within the " + config.getInitializationTimeout() + "ms timeout", null);
+                }
+            }
+
             handleRequest(containerRequest, containerResponse, context);
 
             latch.await();

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

@@ -0,0 +1,148 @@
+package com.amazonaws.serverless.proxy.internal.servlet;
+
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.*;
+import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
+import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
+
+import javax.servlet.http.HttpServletRequest;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Base builder class for {@link AwsLambdaServletContainerHandler}. Implmentations can extend this class to have setters
+ * for the basic parameters.
+ * @param <RequestType> The event object class
+ * @param <ResponseType> The output object class
+ * @param <ContainerRequestType> The container request type. For proxy implementations, this is {@link AwsProxyHttpServletRequest}.
+ *                              The response type is hardcoded to {@link AwsHttpServletResponse} since it is a generic
+ *                              servlet response implementation.
+ * @param <HandlerType> The type of the handler we are building
+ * @param <Builder> The builder object itself. This is used to allow implementations to re-use the setter method from this
+ *                 abstract class through <a href="http://www.artima.com/weblogs/viewpost.jsp?thread=133275">
+ *                     "curiously recurring generic patterns"</a>
+ */
+public abstract class ServletLambdaContainerHandlerBuilder<
+        RequestType,
+        ResponseType,
+        ContainerRequestType extends HttpServletRequest,
+        HandlerType extends AwsLambdaServletContainerHandler<RequestType, ResponseType, ContainerRequestType, AwsHttpServletResponse>,
+        Builder extends ServletLambdaContainerHandlerBuilder<RequestType, ResponseType, ContainerRequestType, HandlerType, Builder>>
+{
+    private static final String MISSING_FIELD_ERROR = "Missing %s in lambda container handler builder";
+
+    protected InitializationWrapper initializationWrapper;
+    protected RequestReader<RequestType, ContainerRequestType> requestReader;
+    protected ResponseWriter<AwsHttpServletResponse, ResponseType> responseWriter;
+    protected SecurityContextWriter<RequestType> securityContextWriter;
+    protected ExceptionHandler<ResponseType> exceptionHandler;
+    protected Class<RequestType> requestTypeClass;
+    protected Class<ResponseType> responseTypeClass;
+
+    /**
+     * Validates that all of the required fields are populated.
+     * @throws ContainerInitializationException If values have not been set on the builder. The message in the exception
+     *  contains a standard error message {@link ServletLambdaContainerHandlerBuilder#MISSING_FIELD_ERROR} populated with
+     *  the list of missing fields.
+     */
+    protected void validate() throws ContainerInitializationException {
+        List<String> errFields = new ArrayList<>();
+        if (requestTypeClass == null) {
+            errFields.add("request type class");
+        }
+        if (responseTypeClass == null) {
+            errFields.add("response type class");
+        }
+        if (requestReader == null) {
+            errFields.add("request reader");
+        }
+        if (responseWriter == null) {
+            errFields.add("response writer");
+        }
+        if (securityContextWriter == null) {
+            errFields.add("security context writer");
+        }
+        if (exceptionHandler == null) {
+            errFields.add("exception handler");
+        }
+        if (initializationWrapper == null) {
+            errFields.add("initialization wrapper");
+        }
+        if (!errFields.isEmpty()) {
+            throw new ContainerInitializationException(String.format(MISSING_FIELD_ERROR, String.join(", ", errFields)), null);
+        }
+    }
+
+    /**
+     * Sets all of the required fields in the builder to the default settings for a Servlet-compatible framework that wants
+     * to support AWS proxy event and output types.
+     * @return A populated builder
+     */
+    public Builder defaultProxy() {
+        initializationWrapper(new InitializationWrapper())
+                .requestReader((RequestReader<RequestType, ContainerRequestType>) new AwsProxyHttpServletRequestReader())
+                .responseWriter((ResponseWriter<AwsHttpServletResponse, ResponseType>) new AwsProxyHttpServletResponseWriter())
+                .securityContextWriter((SecurityContextWriter<RequestType>) new AwsProxySecurityContextWriter())
+                .exceptionHandler((ExceptionHandler<ResponseType>) new AwsProxyExceptionHandler())
+                .requestTypeClass((Class<RequestType>) AwsProxyRequest.class)
+                .responseTypeClass((Class<ResponseType>) AwsProxyResponse.class);
+        return self();
+    }
+
+    /**
+     * Sets the initialization wrapper to be used by the {@link ServletLambdaContainerHandlerBuilder#buildAndInitialize()}
+     * method to start the framework implementations
+     * @param initializationWrapper An implementation of <code>InitializationWrapper</code>. In most cases, this will be
+     *                              set to {@link InitializationWrapper}. The {@link ServletLambdaContainerHandlerBuilder#asyncInit(long)}
+     *                              method sets this to {@link AsyncInitializationWrapper}.
+     * @return This builder object
+     */
+    public Builder initializationWrapper(InitializationWrapper initializationWrapper) {
+        this.initializationWrapper = initializationWrapper;
+        return self();
+    }
+
+    public Builder requestReader(RequestReader<RequestType, ContainerRequestType> requestReader) {
+        this.requestReader = requestReader;
+        return self();
+    }
+
+    public Builder responseWriter(ResponseWriter<AwsHttpServletResponse, ResponseType> responseWriter) {
+        this.responseWriter = responseWriter;
+        return self();
+    }
+
+    public Builder securityContextWriter(SecurityContextWriter<RequestType> securityContextWriter) {
+        this.securityContextWriter = securityContextWriter;
+        return self();
+    }
+
+    public Builder exceptionHandler(ExceptionHandler<ResponseType> exceptionHandler) {
+        this.exceptionHandler = exceptionHandler;
+        return self();
+    }
+
+    public Builder requestTypeClass(Class<RequestType> requestType) {
+        this.requestTypeClass = requestType;
+        return self();
+    }
+
+    public Builder responseTypeClass(Class<ResponseType> responseType) {
+        this.responseTypeClass = responseType;
+        return self();
+    }
+
+    public Builder asyncInit(long actualStartTime) {
+        this.initializationWrapper = new AsyncInitializationWrapper(actualStartTime);
+        return self();
+    }
+
+    /**
+     * Implementations should implement this method to return their type. All of the builder methods in this abstract
+     * class use this method to return the correct builder type.
+     * @return The current builder.
+     */
+    protected abstract Builder self();
+    public abstract HandlerType build() throws ContainerInitializationException;
+    public abstract HandlerType buildAndInitialize() throws ContainerInitializationException;
+}