feat: Support CRaC priming for powertools-tracing

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.prime(); // 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.prime(); // 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.

powertools-tracing/pom.xml

@@ -74,6 +74,10 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -118,9 +122,32 @@
             <artifactId>assertj-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <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-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java

@@ -21,17 +21,38 @@
 import com.amazonaws.xray.entities.Subsegment;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.util.function.Consumer;
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 
 /**
  * A class of helper functions to add additional functionality and ease
  * of use.
  */
-public final class TracingUtils {
+public final class TracingUtils implements Resource {
     private static final Logger LOG = LoggerFactory.getLogger(TracingUtils.class);
     private static ObjectMapper objectMapper;
 
+    // Static block to ensure CRaC registration happens at class loading time
+    static {
+        // Use constructor registration approach like DynamoDBPersistenceStore
+        new TracingUtils();
+    }
+
+    private TracingUtils() {
+        // Register this instance with CRaC (same pattern as DynamoDBPersistenceStore)
+        // Wrap in try-catch for GraalVM compatibility
+        try {
+            Core.getGlobalContext().register(this);
+        } catch (Exception e) {
+            // CRaC registration failed - likely in GraalVM native image or unsupported environment
+            LOG.debug("CRaC registration failed, CRaC priming will be disabled: {}", e.getMessage());
+        }
+    }
+
     /**
      * Put an annotation to the current subsegment with a String value.
      *
@@ -192,4 +213,53 @@ public static void defaultObjectMapper(ObjectMapper objectMapper) {
     public static ObjectMapper objectMapper() {
         return objectMapper;
     }
+
+    /**
+     * Prime TracingUtils for AWS Lambda SnapStart.
+     * This method has no side-effects and can be safely called to trigger SnapStart priming.
+     */
+    public static void prime() {
+        // This method intentionally does nothing but ensures TracingUtils is loaded
+        // The actual priming happens in the beforeCheckpoint() method via CRaC hooks
+    }
+
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        try {
+            // Preload classes first to ensure this always runs
+            ClassPreLoader.preloadClasses();
+
+            // Initialize key components
+            initializeObjectMapper();
+
+            // Initialize X-Ray components by accessing them
+            AWSXRay.getGlobalRecorder();
+
+            // Warm up tracing utilities by calling key methods
+            serviceName();
+
+            // Initialize ObjectMapper for JSON serialization
+            if (objectMapper != null) {
+                objectMapper.writeValueAsString("dummy");
+            }
+        } catch (Exception e) {
+            // Log but don't fail - GraalVM environments may not support all priming operations
+            LOG.debug("CRaC beforeCheckpoint priming encountered an issue: {}", e.getMessage());
+        }
+    }
+
+    private static synchronized void initializeObjectMapper() {
+        if (objectMapper == null) {
+            objectMapper = new ObjectMapper();
+        }
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        try {
+            // No action needed after restore for now, but wrapped for safety
+        } catch (Exception e) {
+            LOG.debug("CRaC afterRestore encountered an issue: {}", e.getMessage());
+        }
+    }
 }

powertools-tracing/src/main/resources/META-INF/native-image/software.amazon.lambda/powertools-tracing/reflect-config.json

@@ -361,5 +361,26 @@
 {
   "name":"sun.security.provider.SHA",
   "methods":[{"name":"<init>","parameterTypes":[] }]
+},
+{
+  "name":"org.crac.Context",
+  "queryAllDeclaredMethods":true,
+  "queryAllDeclaredConstructors":true
+},
+{
+  "name":"org.crac.Core",
+  "queryAllDeclaredMethods":true,
+  "queryAllDeclaredConstructors":true
+},
+{
+  "name":"org.crac.Resource",
+  "queryAllDeclaredMethods":true,
+  "queryAllDeclaredConstructors":true
+},
+{
+  "name":"software.amazon.lambda.powertools.tracing.TracingUtils",
+  "allDeclaredFields":true,
+  "queryAllDeclaredMethods":true,
+  "queryAllDeclaredConstructors":true
 }
 ]

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

@@ -0,0 +1,66 @@
+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.amazonaws.xray.AWSXRay
+com.amazonaws.xray.entities.Entity
+com.amazonaws.xray.entities.Subsegment
+com.amazonaws.xray.entities.Segment
+com.amazonaws.xray.entities.TraceID
+com.amazonaws.xray.entities.TraceHeader
+com.amazonaws.xray.strategy.sampling.SamplingStrategy
+com.amazonaws.xray.strategy.sampling.LocalizedSamplingStrategy
+com.amazonaws.xray.strategy.sampling.NoSamplingStrategy
+com.amazonaws.xray.strategy.sampling.AllSamplingStrategy
+com.amazonaws.xray.strategy.sampling.CentralizedSamplingStrategy
+com.amazonaws.xray.strategy.ContextMissingStrategy
+com.amazonaws.xray.strategy.LogErrorContextMissingStrategy
+com.amazonaws.xray.strategy.RuntimeErrorContextMissingStrategy
+com.amazonaws.xray.strategy.IgnoreErrorContextMissingStrategy
+com.amazonaws.xray.contexts.LambdaSegmentContext
+com.amazonaws.xray.contexts.SegmentContext
+com.amazonaws.xray.contexts.ThreadLocalSegmentContext
+com.amazonaws.xray.emitters.Emitter
+com.amazonaws.xray.emitters.UDPEmitter
+com.amazonaws.xray.listeners.SegmentListener
+com.amazonaws.xray.plugins.Plugin
+com.amazonaws.xray.plugins.ECSPlugin
+com.amazonaws.xray.plugins.EC2Plugin
+com.amazonaws.xray.plugins.EKSPlugin
+software.amazon.lambda.powertools.tracing.TracingUtils
+software.amazon.lambda.powertools.tracing.Tracing
+software.amazon.lambda.powertools.tracing.CaptureMode
+software.amazon.lambda.powertools.tracing.internal.LambdaTracingAspect
+software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor
+software.amazon.lambda.powertools.common.internal.LambdaConstants
+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.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
+org.slf4j.Logger
+org.slf4j.LoggerFactory
+org.slf4j.MDC

powertools-tracing/src/test/java/software/amazon/lambda/powertools/tracing/TracingUtilsCracTest.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to
+ * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+ * the Software, and to permit persons to whom the Software is furnished to do so.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+ * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+ * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+package software.amazon.lambda.powertools.tracing;
+
+import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.mockito.Mockito.mock;
+
+import java.lang.reflect.Field;
+import org.crac.Context;
+import org.crac.Resource;
+import org.junit.jupiter.api.Test;
+
+class TracingUtilsCracTest {
+
+    Context<Resource> context = mock(Context.class);
+
+    @Test
+    void testPrimeMethodDoesNotThrowException() {
+        assertThatNoException().isThrownBy(() -> TracingUtils.prime());
+    }
+
+    @Test
+    void testTracingUtilsLoadsSuccessfully() {
+        // Simply calling TracingUtils.prime() should trigger CRaC registration
+        assertThatNoException().isThrownBy(() -> TracingUtils.prime());
+
+        // Verify that TracingUtils class is loaded and accessible
+        assertThatNoException().isThrownBy(() -> TracingUtils.objectMapper());
+    }
+}

spotbugs-exclude.xml

@@ -193,8 +193,16 @@
     </Match>
     <Match>
         <Bug pattern="RV_RETURN_VALUE_IGNORED_NO_SIDE_EFFECT"/>
-        <Class name="software.amazon.lambda.powertools.validation.ValidationConfig"/>
-        <Method name="beforeCheckpoint"/>
+        <Or>
+            <And>
+                <Class name="software.amazon.lambda.powertools.validation.ValidationConfig"/>
+                <Method name="beforeCheckpoint"/>
+            </And>
+            <And>
+                <Class name="software.amazon.lambda.powertools.tracing.TracingUtils"/>
+                <Method name="beforeCheckpoint"/>
+            </And>
+        </Or>
     </Match>
     <!--Functionally needed-->
     <Match>