feat: Support CRaC priming for powertools-tracing

.sonarcloud.properties

@@ -15,3 +15,9 @@ sonar.exclusions=examples/**/*,powertools-e2e-tests/handlers/**/*
 
 # Ignore code duplicates in the examples
 sonar.cpd.exclusions=examples/**/*,powertools-e2e-tests/**/*
+
+# Ignore singleton pattern detection for CRaC Resource implementations
+# Singleton pattern is required for CRaC (Coordinated Restore at Checkpoint) Resource interface
+sonar.issue.ignore.multicriteria=e1
+sonar.issue.ignore.multicriteria.e1.ruleKey=java:S6548
+sonar.issue.ignore.multicriteria.e1.resourceKey=**/TracingUtils.java,**/JsonConfig.java

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-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaConstants.java

@@ -23,4 +23,6 @@ public class LambdaConstants {
     public static final String ROOT_EQUALS = "Root=";
     public static final String POWERTOOLS_SERVICE_NAME = "POWERTOOLS_SERVICE_NAME";
     public static final String SERVICE_UNDEFINED = "service_undefined";
+    public static final String AWS_LAMBDA_INITIALIZATION_TYPE = "AWS_LAMBDA_INITIALIZATION_TYPE";
+    public static final String PROVISIONED_CONCURRENCY = "provisioned-concurrency";
 }

powertools-common/src/main/java/software/amazon/lambda/powertools/common/internal/LambdaHandlerProcessor.java

@@ -88,7 +88,10 @@ protected static void resetServiceName() {
     }
 
     public static boolean isColdStart() {
-        return IS_COLD_START == null;
+        // If this is not the first invocation, it's definitely not a cold start
+        // Check if this execution environment was pre-warmed via provisioned concurrency
+        // Traditional cold start detection - first invocation without provisioned concurrency
+        return IS_COLD_START == null && !LambdaConstants.PROVISIONED_CONCURRENCY.equals(getenv(LambdaConstants.AWS_LAMBDA_INITIALIZATION_TYPE));
     }
 
     public static void coldStartDone() {

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,32 @@
 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)
+        Core.getGlobalContext().register(this);
+    }
+
     /**
      * Put an annotation to the current subsegment with a String value.
      *
@@ -192,4 +207,44 @@ 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 {
+        // 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");
+        }
+    }
+
+    private static synchronized void initializeObjectMapper() {
+        if (objectMapper == null) {
+            objectMapper = new ObjectMapper();
+        }
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }

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,42 @@
+/*
+ * 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.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>