feat(type-safe-api): enable snapstart for java handlers by default and optimise java interceptors

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