AWS SnapStart For Java Lambda Using AWS CDK

Question

I am trying to orchestrate Lambda's using the Java 11 (or 17) AWS CDK. Lambda has serious issues with JVM compilation and cold starts, so I have been looking into using the newer SnapStart feature offered by AWS. The CDK pushes the CloudFormation stack with SnapStart components included, a new version of the Lambda, and an Alias for the version. These are all necessary requirements for SnapStart.

CDK Lambda Creation w/ SnapStart Config & Alias

       //Lambda Function
        String lambdaHandlerName = "com.handlers.lambdaHandler::handleRequest";
        IRole role5 = Role.fromRoleName(this, "Lambda", "LambdaRole");

        Function lambdaFunction = new Function(this, "lambdaFunction", getLambdaFunctionProps(lambdaEnvMap, lambdaHandlerName, role5, "lambdaHandler"));
        dynamoTable.grantReadData(lambdaFunction);

        //SnapStart Config
        lambdaFunction.getCurrentVersion().addAlias("snap", AliasOptions.builder()
                .description("Alias version for snap resources")
                .build());
        CfnFunction lambdaCfnFunction = (CfnFunction) lambdaCfnFunction.getNode().getDefaultChild();
        lambdaCfnFunction.setSnapStart(CfnFunction.SnapStartProperty.builder()
                .applyOn("PublishedVersions")
                .build());

Within the AWS Console itself, it appears SnapStart is enabled for the base function set to "PublishedVersions".

SnapStart enabled on base function

And when I look at my most recently published version, I find that it too has SnapStart enabled.

SnapStart enabled on Version

But when I invoke my Lambda function through API Gateway (which is configured to the correct version of Lambda), SnapStart does not seem to be working for the function and my cold start times are still ~2s.

I have also tried implementing runtime hooks, suggested by AWS for SnapStart to the Lambda function handler itself:

Handler Runtime Hooks for SnapStart

   public GetAllQuotesHandler() {
        Core.getGlobalContext().register(this);
    }

    @Override
    public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent event, Context context) {
      ...
    }

    
    @Override
    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
        this.dynamoDb.close();
        System.out.println("Before");
    }

    @Override
    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
        System.out.println("After");
    }
}

I am not sure what is wrong with my configuration, any advice?

I am expecting AWS Lambda to engage SnapStart and invoke my Lambda function without cold starts.

Answer 1

So far it looks as if you have configured SnapStart on the lambda config, which on its own will give you a slight reduction in cold start times. In order to get the full reduction in cold start time you need to ‘prime’ the state of the JVM before the snapshot is taken when the lambda version is created.

The beforeCheckpoint(...) runtime hook is called before the snapshot of the JVM/container is taken. By making calls to units of your code from within this method, the JVM will interpret the byte-code stored .class files for these units and compile them to native machine code using JIT compilation. This will be saved in the snapshot so this bytecode won't need to be interpreted when the lambda is first invoked for a real request, which should result in a fast startup time.

The simplest way to do this 'priming' with your code is to call the handleRequest(…) method from within the beforeCheckpoint(…) method with a testAPIGatewayProxyRequestEvent. Make sure that the test event has enough information in it to get past any business logic you have, and actually call the dynamoDb client. Ideally, you would call as much of you code as possible so that it gets interpreted before the snapshot is taken.

So it will look something like this:

public class GetAllQuotesHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent>, Resource {

   private final DynamoDbEnhancedClient dynamoDb;

   public GetAllQuotesHandler() {

        AwsCredentialsProvider credentialsProvider = ContainerCredentialsProvider.builder().build();
    
        this.dynamoDb = DynamoDbEnhancedClient.builder()
                .dynamoDbClient(DynamoDbClient.builder()
                .credentialsProvider(credentialsProvider)
                .region(region)
                .httpClientBuilder(UrlConnectionHttpClient.builder())
                .build()).build();
  

        Core.getGlobalContext().register(this);
    }

    @Override
    public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent event, Context context) {
      
     // query DB based on request.
     ...
    }

    // Call the handler method with a test event. 
    @Override
    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
        log.info("beforeCheckpoint({}) called.", context);
        
        APIGatewayProxyRequestEvent testEvent = new APIGatewayProxyRequestEvent().withBody("string body").withHttpMethod("POST");
        try {
            handleRequest(testEvent, null);
        } catch (Exception ex) {
            log.info("Error executing beforeCheckpoint({})", testEvent);
        }
    }

    @Override
    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
        // not used
    }
}

You will also need to use a different credentials provider if your function calls other AWS resources using the aws SDK with SnapStart:

ContainerCredentialsProvider.builder().build();

Since these resources need up-to-date credentials. You can't use the environment variable credentials provider, because these credentials will only get loaded in when the version is created and they will eventually expire. This is because the no-args constructor will only get called once, when the version is created, rather than every time the lambda function is invoked in a new execution environment.