feat(type-safe-api): enable snapstart for java handlers by default an…

…d optimise java interceptors

Java handlers now have snapstart enabled by default to improve cold start times. Ensure logger,
metrics and tracer perform as much initialisation as possible during the init phase which can
therefore be saved from snapstart invocations.

Fixes #572

packages/type-safe-api/scripts/type-safe-api/generators/java-cdk-infrastructure/templates/functions.handlebars

@@ -15,12 +15,13 @@
 
 import org.jetbrains.annotations.NotNull;
 import software.amazon.awscdk.services.lambda.Function;
+import software.aws.pdk.type_safe_api.SnapStartFunction;
 import software.constructs.Construct;
 
 /**
  * Lambda function construct which points to the {{vendorExtensions.x-handler.language}} implementation of {{operationIdCamelCase}}
  */
-public class {{operationIdCamelCase}}Function extends Function {
+public class {{operationIdCamelCase}}Function extends {{#startsWith vendorExtensions.x-handler.language 'java' ~}}SnapStart{{~/startsWith}}Function {
     public {{operationIdCamelCase}}Function(@NotNull Construct scope, @NotNull String id, @NotNull {{operationIdCamelCase}}FunctionProps props) {
         super(scope, id, props);
     }
@@ -67,6 +68,7 @@ import software.amazon.awscdk.services.lambda.VersionOptions;
 import software.amazon.awscdk.services.logs.RetentionDays;
 import software.amazon.awscdk.services.sns.ITopic;
 import software.amazon.awscdk.services.sqs.IQueue;
+import software.aws.pdk.type_safe_api.SnapStartFunctionProps;
 
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -77,7 +79,7 @@ import java.util.List;
 import java.util.Map;
 
 @lombok.Builder @lombok.Getter
-public class {{operationIdCamelCase}}FunctionProps implements FunctionProps {
+public class {{operationIdCamelCase}}FunctionProps implements {{#startsWith vendorExtensions.x-handler.language 'java' ~}}SnapStart{{~/startsWith}}FunctionProps {
     private static String infraProjectAbsolutePath;
 
     static {

packages/type-safe-api/scripts/type-safe-api/generators/java-lambda-handlers/templates/handlers.handlebars

@@ -31,13 +31,18 @@ import java.util.List;
  * The {{operationIdCamelCase}} class manages marshalling inputs and outputs.
  */
 public class {{operationIdCamelCase}}Handler extends {{operationIdCamelCase}} {
+    /**
+     * Interceptors are initialised once during the lambda "init" phase
+     */
+    private final List<Interceptor<{{operationIdCamelCase}}Input>> interceptors = DefaultInterceptors.all();
+
     /**
      * Return the interceptors for this handler.
      * You can also use the @Interceptors annotation on the class to add interceptors
      */
     @Override
     public List<Interceptor<{{operationIdCamelCase}}Input>> getInterceptors() {
-        return DefaultInterceptors.all();
+        return this.interceptors;
     }
 
     /**

packages/type-safe-api/scripts/type-safe-api/generators/java/templates/handlers.mustache

@@ -1023,6 +1023,11 @@ public abstract class {{operationIdCamelCase}} implements RequestHandler<APIGate
      */
     public abstract {{operationIdCamelCase}}Response handle(final {{operationIdCamelCase}}RequestInput request);
 
+    /**
+     * Interceptors that the handler class has been decorated with
+     */
+    private List<Interceptor<{{operationIdCamelCase}}Input>> annotationInterceptors = Handlers.getAnnotationInterceptors({{operationIdCamelCase}}.class);
+
     /**
      * For more complex interceptors that require instantiation with parameters, you may override this method to
      * return a list of instantiated interceptors. For simple interceptors with no need for constructor arguments,
@@ -1055,9 +1060,6 @@ public abstract class {{operationIdCamelCase}} implements RequestHandler<APIGate
 
         List<Interceptor<{{operationIdCamelCase}}Input>> interceptors = new ArrayList<>();
         interceptors.addAll(additionalInterceptors);
-
-        List<Interceptor<{{operationIdCamelCase}}Input>> annotationInterceptors = Handlers.getAnnotationInterceptors(this.getClass());
-
         interceptors.addAll(annotationInterceptors);
         interceptors.addAll(this.getInterceptors());
 

packages/type-safe-api/scripts/type-safe-api/generators/java/templates/interceptors.mustache

@@ -123,6 +123,7 @@ import software.amazon.lambda.powertools.logging.LoggingUtils;
  * See https://docs.powertools.aws.dev/lambda/java/latest/core/logging/
  */
 public class LoggingInterceptor<Input> implements Interceptor<Input> {
+    private Logger logger = LogManager.getLogger(LoggingInterceptor.class);
 
     /**
      * Return the instance of the logger from the interceptor context
@@ -161,7 +162,6 @@ public class LoggingInterceptor<Input> implements Interceptor<Input> {
         LoggingUtils.appendKey("operationId", operationId);
 
         // Add the logger to the interceptor context
-        Logger logger = LogManager.getLogger(operationId);
         input.getInterceptorContext().put("logger", logger);
 
         Response response = input.getChain().next(input);
@@ -202,6 +202,13 @@ import software.amazon.lambda.powertools.tracing.TracingUtils;
  */
 public class TracingInterceptor<Input> implements Interceptor<Input> {
 
+    {
+        // Create a segment during the lambda init phase to ensure xray emitter
+        // is warmed up prior to invocation phase
+        AWSXRay.beginSubsegment("Tracing Interceptor - Init");
+        AWSXRay.endSubsegment();
+    }
+
     private final boolean captureResponse;
 
     public TracingInterceptor(final boolean captureResponse) {
@@ -284,6 +291,7 @@ import software.amazon.lambda.powertools.metrics.MetricsUtils;
  * See: https://docs.powertools.aws.dev/lambda/typescript/latest/core/metrics
  */
 public class MetricsInterceptor<Input> implements Interceptor<Input> {
+    private MetricsLogger metrics = MetricsUtils.metricsLogger();
 
     /**
      * Return the instance of the metrics logger from the interceptor context
@@ -298,7 +306,6 @@ public class MetricsInterceptor<Input> implements Interceptor<Input> {
 
     @Override
     public Response handle(final ChainedRequestInput<Input> input) {
-        MetricsLogger metrics = MetricsUtils.metricsLogger();
         metrics.putDimensions(DimensionSet.of("operationId", (String) input.getInterceptorContext().get("operationId")));
 
         input.getInterceptorContext().put("metrics", metrics);

packages/type-safe-api/scripts/type-safe-api/generators/python-cdk-infrastructure/templates/functions.handlebars

@@ -10,6 +10,7 @@
 from aws_cdk.aws_lambda import (
     Function, Runtime, Tracing, Code
 )
+from aws_pdk.type_safe_api import SnapStartFunction
 from os import path
 from pathlib import Path
 
@@ -19,7 +20,7 @@ from pathlib import Path
 {{#operation ~}}
 {{#if vendorExtensions.x-handler}}
 
-class {{operationIdCamelCase}}Function(Function):
+class {{operationIdCamelCase}}Function({{#startsWith vendorExtensions.x-handler.language 'java' ~}}SnapStart{{~/startsWith}}Function):
     """
     Lambda function construct which points to the {{vendorExtensions.x-handler.language}} implementation of {{operationIdCamelCase}}
     """

packages/type-safe-api/scripts/type-safe-api/generators/typescript-cdk-infrastructure/templates/functions.handlebars

@@ -9,6 +9,7 @@
 }
 ###/TSAPI_WRITE_FILE###import { Construct } from "constructs";
 import { Duration } from "aws-cdk-lib";
+import { SnapStartFunction, SnapStartFunctionProps } from "@aws/pdk/type-safe-api";
 import { Code, Function, Runtime, Tracing, FunctionProps } from "aws-cdk-lib/aws-lambda";
 import * as path from "path";
 {{#apiInfo ~}}
@@ -20,12 +21,12 @@ import * as path from "path";
 /**
  * Options for the {{operationIdCamelCase}}Function construct
  */
-export interface {{operationIdCamelCase}}FunctionProps extends Omit<FunctionProps, 'code' | 'handler' | 'runtime'> {}
+export interface {{operationIdCamelCase}}FunctionProps extends Omit<{{#startsWith vendorExtensions.x-handler.language 'java' ~}}SnapStart{{~/startsWith}}FunctionProps, 'code' | 'handler' | 'runtime'> {}
 
 /**
  * Lambda function construct which points to the {{vendorExtensions.x-handler.language}} implementation of {{operationIdCamelCase}}
  */
-export class {{operationIdCamelCase}}Function extends Function {
+export class {{operationIdCamelCase}}Function extends {{#startsWith vendorExtensions.x-handler.language 'java' ~}}SnapStart{{~/startsWith}}Function {
   constructor(scope: Construct, id: string, props?: {{operationIdCamelCase}}FunctionProps) {
     super(scope, id, {
       {{#startsWith vendorExtensions.x-handler.language 'typescript' ~}}

packages/type-safe-api/src/construct/functions/index.ts

@@ -0,0 +1,3 @@
+/*! Copyright [Amazon.com](http://amazon.com/), Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: Apache-2.0 */
+export * from "./snap-start-java-function";

packages/type-safe-api/src/construct/functions/snap-start-java-function.ts

@@ -0,0 +1,30 @@
+/*! Copyright [Amazon.com](http://amazon.com/), Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: Apache-2.0 */
+import { CfnFunction, Function, FunctionProps } from "aws-cdk-lib/aws-lambda";
+import { Construct } from "constructs";
+
+/**
+ * Options for the SnapStartFunction construct
+ */
+export interface SnapStartFunctionProps extends FunctionProps {
+  /**
+   * When true, disable snap start
+   * @default false
+   */
+  readonly disableSnapStart?: boolean;
+}
+
+/**
+ * A lambda function which enables SnapStart on published versions by default
+ */
+export class SnapStartFunction extends Function {
+  constructor(scope: Construct, id: string, props: SnapStartFunctionProps) {
+    super(scope, id, props);
+
+    if (!props.disableSnapStart) {
+      (this.node.defaultChild as CfnFunction).addPropertyOverride("SnapStart", {
+        ApplyOn: "PublishedVersions",
+      });
+    }
+  }
+}

packages/type-safe-api/src/construct/index.ts

@@ -5,3 +5,4 @@ export * from "./waf/types";
 export * from "./authorizers";
 export * from "./integrations";
 export * from "./spec";
+export * from "./functions";

packages/type-safe-api/src/construct/integrations/lambda.ts

@@ -8,6 +8,7 @@ import {
   IntegrationGrantProps,
   IntegrationRenderProps,
 } from "./integration";
+import { SnapStartFunction } from "../functions/snap-start-java-function";
 import { functionInvocationUri } from "../spec/utils";
 
 /**
@@ -18,7 +19,12 @@ export class LambdaIntegration extends Integration {
 
   constructor(lambdaFunction: IFunction) {
     super();
-    this.lambdaFunction = lambdaFunction;
+    // Snap Start applies only to versions, so if the function is a SnapStartFunction, we'll reference the current version
+    if (lambdaFunction instanceof SnapStartFunction) {
+      this.lambdaFunction = lambdaFunction.currentVersion;
+    } else {
+      this.lambdaFunction = lambdaFunction;
+    }
   }
 
   /**