feat: Support CRaC priming for powertools-tracing and powertools-serialization

docs/core/tracing.md

@@ -419,5 +419,55 @@ Below is an example configuration needed for each test case.
         }
     ```
 
+## Advanced
 
+### Lambda SnapStart priming
 
+The Tracing utility integrates with AWS Lambda SnapStart to improve restore durations. To make sure the SnapStart priming logic of this utility runs correctly, you need an explicit reference to `TracingUtils` in your code to allow the library to register before SnapStart takes a memory snapshot. Learn more about what priming is in this [blog post](https://aws.amazon.com/blogs/compute/optimizing-cold-start-performance-of-aws-lambda-using-advanced-priming-strategies-with-snapstart/){target="_blank"}.
+
+Make sure to reference `TracingUtils` in your Lambda handler initialization code. This can be done by adding one of the following lines to your handler class:
+
+=== "Constructor"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.tracing.Tracing;
+    import software.amazon.lambda.powertools.tracing.TracingUtils;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        public MyFunctionHandler() {
+            TracingUtils.putAnnotation("init", "priming"); // Ensure TracingUtils is loaded for SnapStart
+        }
+
+        @Override
+        @Tracing
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```
+
+=== "Static Initializer"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.tracing.Tracing;
+    import software.amazon.lambda.powertools.tracing.TracingUtils;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        static {
+            TracingUtils.putAnnotation("init", "priming"); // Ensure TracingUtils is loaded for SnapStart
+        }
+
+        @Override
+        @Tracing
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```
+
+!!! note "Important: Direct TracingUtils reference required"
+    Using only the `@Tracing` annotation is not sufficient to trigger SnapStart priming. You must have a direct reference to `TracingUtils` in your code (as shown in the examples above) to ensure the CRaC hooks are properly registered.

docs/utilities/serialization.md

@@ -472,3 +472,53 @@ to powertools.You can then use it to do your validation or in idempotency module
        }
     }
     ```
+
+## Advanced
+
+### Lambda SnapStart priming
+
+The Serialization utility integrates with AWS Lambda SnapStart to improve restore durations. To make sure the SnapStart priming logic of this utility runs correctly, you need an explicit reference to `JsonConfig` in your code to allow the library to register before SnapStart takes a memory snapshot. Learn more about what priming is in this [blog post](https://aws.amazon.com/blogs/compute/optimizing-cold-start-performance-of-aws-lambda-using-advanced-priming-strategies-with-snapstart/){target="_blank"}.
+
+If you don't set a custom `JsonConfig` in your code yet, make sure to reference `JsonConfig` in your Lambda handler initialization code. This can be done by adding one of the following lines to your handler class:
+
+=== "Constructor"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.utilities.JsonConfig;
+    import static software.amazon.lambda.powertools.utilities.EventDeserializer.extractDataFrom;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        public MyFunctionHandler() {
+            JsonConfig.get(); // Ensure JsonConfig is loaded for SnapStart
+        }
+
+        @Override
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            Product product = extractDataFrom(input).as(Product.class);
+            // ...
+            return something;
+        }
+    }
+    ```
+
+=== "Static Initializer"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.utilities.JsonConfig;
+    import static software.amazon.lambda.powertools.utilities.EventDeserializer.extractDataFrom;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        static {
+            JsonConfig.get(); // Ensure JsonConfig is loaded for SnapStart
+        }
+
+        @Override
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            Product product = extractDataFrom(input).as(Product.class);
+            // ...
+            return something;
+        }
+    }
+    ```

powertools-serialization/pom.xml

@@ -47,6 +47,14 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.lambda</groupId>
+            <artifactId>powertools-common</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -74,6 +82,11 @@
             <artifactId>aws-lambda-java-tests</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
@@ -96,6 +109,24 @@
     </build>
 
     <profiles>
+        <profile>
+            <id>generate-classesloaded-file</id>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.apache.maven.plugins</groupId>
+                        <artifactId>maven-surefire-plugin</artifactId>
+                        <configuration>
+                            <argLine>
+                                -Xlog:class+load=info:classesloaded.txt
+                                --add-opens java.base/java.util=ALL-UNNAMED
+                                --add-opens java.base/java.lang=ALL-UNNAMED
+                            </argLine>
+                        </configuration>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
         <profile>
             <id>generate-graalvm-files</id>
             <build>

powertools-serialization/src/main/java/software/amazon/lambda/powertools/utilities/JsonConfig.java

@@ -27,11 +27,15 @@
 import io.burt.jmespath.function.FunctionRegistry;
 import io.burt.jmespath.jackson.JacksonRuntime;
 import java.util.function.Supplier;
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.utilities.jmespath.Base64Function;
 import software.amazon.lambda.powertools.utilities.jmespath.Base64GZipFunction;
 import software.amazon.lambda.powertools.utilities.jmespath.JsonFunction;
 
-public final class JsonConfig {
+public final class JsonConfig implements Resource {
 
     private static final Supplier<ObjectMapper> objectMapperSupplier = () -> JsonMapper.builder()
             // Don't throw an exception when json has extra fields you are not serializing on.
@@ -61,6 +65,11 @@
 
     private JmesPath<JsonNode> jmesPath = new JacksonRuntime(configuration, getObjectMapper());
 
+    // Static block to ensure CRaC registration happens at class loading time
+    static {
+        Core.getGlobalContext().register(get());
+    }
+
     private JsonConfig() {
     }
 
@@ -103,7 +112,59 @@
         jmesPath = new JacksonRuntime(updatedConfig, getObjectMapper());
     }
 
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        // Initialize key components
+        getObjectMapper();
+        getJmesPath();
+
+        // Perform dummy serialization/deserialization operations to warm up Jackson ObjectMapper
+        try {
+            ObjectMapper mapper = getObjectMapper();
+
+            // Prime common AWS Lambda event types as suggested by Philipp
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent", 
+                "{\"httpMethod\":\"GET\",\"path\":\"/test\",\"headers\":{\"Content-Type\":\"application/json\"}}");
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPEvent", 
+                "{\"version\":\"2.0\",\"routeKey\":\"GET /test\",\"requestContext\":{\"http\":{\"method\":\"GET\"}}}");
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.SQSEvent", 
+                "{\"Records\":[{\"messageId\":\"test\",\"body\":\"test message\"}]}");
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.SNSEvent", 
+                "{\"Records\":[{\"Sns\":{\"Message\":\"test message\",\"TopicArn\":\"arn:aws:sns:us-east-1:123456789012:test\"}}]}");
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.KinesisEvent", 
+                "{\"Records\":[{\"kinesis\":{\"data\":\"dGVzdA==\",\"partitionKey\":\"test\"}}]}");
+            primeEventType(mapper, "com.amazonaws.services.lambda.runtime.events.ScheduledEvent", 
+                "{\"source\":\"aws.events\",\"detail-type\":\"Scheduled Event\"}");
+
+            // Warm up JMESPath function registry
+            getJmesPath().compile("@").search(mapper.readTree("{\"test\":\"value\"}"));
+
+        } catch (Exception e) {
+            // Ignore exceptions during priming as they're expected in some environments
+        }
+
+        // Preload classes
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
+
+    private void primeEventType(ObjectMapper mapper, String className, String sampleJson) {
+        try {
+            Class<?> eventClass = Class.forName(className);
+            // Deserialize sample JSON to the event class
+            Object event = mapper.readValue(sampleJson, eventClass);
+            // Serialize back to JSON to warm up both directions
+            mapper.writeValueAsString(event);
+        } catch (Exception e) {
+            // Ignore exceptions for event types that might not be available
+        }
+    }
+
     private static class ConfigHolder {
         private static final JsonConfig instance = new JsonConfig();
     }
 }

powertools-serialization/src/main/resources/classesloaded.txt

@@ -0,0 +1,74 @@
+java.lang.Object
+java.io.Serializable
+java.lang.Comparable
+java.lang.CharSequence
+java.lang.String
+java.lang.Class
+java.lang.Cloneable
+java.lang.ClassLoader
+java.lang.System
+java.lang.Throwable
+java.lang.Error
+java.lang.Exception
+java.lang.RuntimeException
+com.fasterxml.jackson.databind.ObjectMapper
+com.fasterxml.jackson.databind.JsonNode
+com.fasterxml.jackson.databind.node.ObjectNode
+com.fasterxml.jackson.databind.node.ArrayNode
+com.fasterxml.jackson.databind.node.TextNode
+com.fasterxml.jackson.databind.node.NumericNode
+com.fasterxml.jackson.databind.node.BooleanNode
+com.fasterxml.jackson.databind.node.NullNode
+com.fasterxml.jackson.databind.json.JsonMapper
+com.fasterxml.jackson.core.JsonFactory
+com.fasterxml.jackson.core.JsonGenerator
+com.fasterxml.jackson.core.JsonParser
+com.fasterxml.jackson.core.JsonToken
+com.fasterxml.jackson.databind.DeserializationFeature
+com.fasterxml.jackson.databind.SerializationFeature
+com.fasterxml.jackson.databind.MapperFeature
+com.fasterxml.jackson.databind.JsonSerializer
+com.fasterxml.jackson.databind.JsonDeserializer
+com.fasterxml.jackson.databind.SerializerProvider
+com.fasterxml.jackson.databind.DeserializationContext
+com.fasterxml.jackson.annotation.JsonInclude
+com.fasterxml.jackson.annotation.JsonInclude$Include
+io.burt.jmespath.JmesPath
+io.burt.jmespath.RuntimeConfiguration
+io.burt.jmespath.RuntimeConfiguration$Builder
+io.burt.jmespath.function.BaseFunction
+io.burt.jmespath.function.FunctionRegistry
+io.burt.jmespath.jackson.JacksonRuntime
+software.amazon.lambda.powertools.utilities.JsonConfig
+software.amazon.lambda.powertools.utilities.EventDeserializer
+software.amazon.lambda.powertools.utilities.EventDeserializationException
+software.amazon.lambda.powertools.utilities.jmespath.Base64Function
+software.amazon.lambda.powertools.utilities.jmespath.Base64GZipFunction
+software.amazon.lambda.powertools.utilities.jmespath.JsonFunction
+com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent
+com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent
+com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPEvent
+com.amazonaws.services.lambda.runtime.events.APIGatewayV2HTTPResponse
+com.amazonaws.services.lambda.runtime.events.ActiveMQEvent
+com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerRequestEvent
+com.amazonaws.services.lambda.runtime.events.ApplicationLoadBalancerResponseEvent
+com.amazonaws.services.lambda.runtime.events.CloudFormationCustomResourceEvent
+com.amazonaws.services.lambda.runtime.events.CloudWatchLogsEvent
+com.amazonaws.services.lambda.runtime.events.KafkaEvent
+com.amazonaws.services.lambda.runtime.events.KinesisAnalyticsFirehoseInputPreprocessingEvent
+com.amazonaws.services.lambda.runtime.events.KinesisAnalyticsStreamsInputPreprocessingEvent
+com.amazonaws.services.lambda.runtime.events.KinesisEvent
+com.amazonaws.services.lambda.runtime.events.KinesisFirehoseEvent
+com.amazonaws.services.lambda.runtime.events.RabbitMQEvent
+com.amazonaws.services.lambda.runtime.events.SNSEvent
+com.amazonaws.services.lambda.runtime.events.SQSEvent
+com.amazonaws.services.lambda.runtime.events.ScheduledEvent
+org.slf4j.Logger
+org.slf4j.LoggerFactory
+java.util.function.Supplier
+java.lang.ThreadLocal
+java.util.Map
+java.util.HashMap
+java.util.List
+java.util.ArrayList
+java.util.concurrent.ConcurrentHashMap

powertools-serialization/src/test/java/software/amazon/lambda/powertools/utilities/JsonConfigCracTest.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates.
+ * Licensed under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 software.amazon.lambda.powertools.utilities;
+
+import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.mockito.Mockito.mock;
+
+import org.crac.Context;
+import org.crac.Resource;
+import org.junit.jupiter.api.Test;
+
+class JsonConfigCracTest {
+
+    JsonConfig config = JsonConfig.get();
+    Context<Resource> context = mock(Context.class);
+
+    @Test
+    void testBeforeCheckpointDoesNotThrowException() {
+        assertThatNoException().isThrownBy(() -> config.beforeCheckpoint(context));
+    }
+
+    @Test
+    void testAfterRestoreDoesNotThrowException() {
+        assertThatNoException().isThrownBy(() -> config.afterRestore(context));
+    }
+}