com.amazonaws.serverless.proxy.AsyncInitializationWrapper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of aws-serverless-java-container-core Show documentation
Allows Java applications written for a servlet container to run in AWS Lambda
There is a newer version: 2.0.3
Show newest version
/*
 * Copyright 2020 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;

import com.amazonaws.serverless.exceptions.ContainerInitializationException;
import com.amazonaws.serverless.proxy.internal.InitializableLambdaContainerHandler;
import com.amazonaws.services.lambda.runtime.Context;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.management.ManagementFactory;
import java.time.Instant;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * An async implementation of the InitializationWrapper interface. This initializer calls the
 * {@link InitializableLambdaContainerHandler#initialize()} in a separate thread. Then uses a latch to wait for the maximum Lambda
 * initialization time of 10 seconds, if the initialize method takes longer than 10 seconds to return, the
 * {@link #start(InitializableLambdaContainerHandler)} returns control to the caller and lets the initialization thread continue in
 * the background. The {@link com.amazonaws.serverless.proxy.internal.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 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);


    /**
     * 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;
    }

    /**
     * Creates a new instance of the async initializer using the actual JVM start time as the starting point to measure
     * the 10 seconds timeout.
     */
    public AsyncInitializationWrapper() {
        actualStartTime = ManagementFactory.getRuntimeMXBean().getStartTime();
    }

    @Override
    public void start(InitializableLambdaContainerHandler handler) throws ContainerInitializationException {
        if(ASYNC_INIT_DISABLED){
            log.info("Async init disabled due to \"{}\" initialization", INITIALIZATION_TYPE);
            super.start(handler);
            return;
        }
        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 time
            long awaitTime = (actualStartTime + LAMBDA_MAX_INIT_TIME_MS) - curTime - INIT_GRACE_TIME_MS;
            log.info("Async initialization will wait for {}ms (init grace time is configured to {})",
                    awaitTime, INIT_GRACE_TIME_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);
        }
    }

    public long getActualStartTimeMs() {
        return actualStartTime;
    }

    @Override
    public CountDownLatch getInitializationLatch() {
        if (ASYNC_INIT_DISABLED){
            return super.getInitializationLatch();
        }
        return initializationLatch;
    }

    private static class AsyncInitializer implements Runnable {
        private final InitializableLambdaContainerHandler handler;
        private CountDownLatch initLatch;
        private Logger log = LoggerFactory.getLogger(AsyncInitializationWrapper.class);

        AsyncInitializer(CountDownLatch latch, InitializableLambdaContainerHandler 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();
            }
        }
    }
}