feat(crac): Support CRaC and priming of powertools tracing (#2345)

* fix: cold start only on on-demand invocation

* fix: cold start only on on-demand invocation

* fix: cold start only on on-demand invocation

* mvn clean test is successful

* reversed the sqllite changes

* fix: priming for powertools-tracing

* reverted OS specific changes

* mvn clean test is successful

* reversed the sqllite changes

* fix: priming for powertools-tracing

* reverted OS specific changes

* resolved testing issues

* resolved testing issues

* resolved testing issues

* resolved testing issues

* resolved testing issues

* incorporated comments

* Update pom.xml

Author: Attyuttam

docs/core/tracing.md

@@ -447,3 +447,52 @@ Below is an example configuration needed for each test case.
             // test logic
         }
     ```
+## 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"}.
+
+If you don't set a custom `TracingUtils` in your code yet, 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.validation.Validation;
+    import software.amazon.lambda.powertools.validation.ValidationConfig;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        public MyFunctionHandler() {
+            TracingUtils.init(); // Ensure TracingUtils is loaded for SnapStart
+        }
+
+        @Override
+        @Validation(inboundSchema = "classpath:/schema_in.json", outboundSchema = "classpath:/schema_out.json")
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```
+
+=== "Static Initializer"
+
+    ```java hl_lines="7"
+    import software.amazon.lambda.powertools.validation.Validation;
+    import software.amazon.lambda.powertools.validation.ValidationConfig;
+
+    public class MyFunctionHandler implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+        static {
+            TracingUtils.init(); // Ensure TracingUtils is loaded for SnapStart
+        }
+
+        @Override
+        @Validation(inboundSchema = "classpath:/schema_in.json", outboundSchema = "classpath:/schema_out.json")
+        public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent input, Context context) {
+            // ...
+            return something;
+        }
+    }
+    ```

powertools-tracing/pom.xml

@@ -33,6 +33,10 @@
     </description>
 
     <dependencies>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.aspectj</groupId>
             <artifactId>aspectjrt</artifactId>
@@ -80,6 +84,11 @@
         </dependency>
 
         <!-- Test dependencies -->
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-api</artifactId>
@@ -144,8 +153,33 @@
     </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>
+            <dependencies>
+                <dependency>
+                    <groupId>org.mockito</groupId>
+                    <artifactId>mockito-subclass</artifactId>
+                    <scope>test</scope>
+                </dependency>
+            </dependencies>
             <build>
                 <plugins>
                     <plugin>
@@ -165,6 +199,13 @@
         </profile>
         <profile>
             <id>graalvm-native</id>
+            <dependencies>
+                <dependency>
+                    <groupId>org.mockito</groupId>
+                    <artifactId>mockito-subclass</artifactId>
+                    <scope>test</scope>
+                </dependency>
+            </dependencies>
             <build>
                 <plugins>
                     <plugin>

powertools-tracing/src/main/java/software/amazon/lambda/powertools/tracing/TracingUtils.java

@@ -16,22 +16,37 @@
 
 import static software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor.serviceName;
 
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
 import com.amazonaws.xray.AWSXRay;
 import com.amazonaws.xray.entities.Entity;
 import com.amazonaws.xray.entities.Subsegment;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import java.util.function.Consumer;
 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;
 
+    private static final TracingUtils INSTANCE = new TracingUtils();
+
+    static {
+        Core.getGlobalContext().register(INSTANCE);
+    }
+
+    public static void init() {
+        // Placeholder method used to enable SnapStart priming. Users need a direct reference to this class in order for the CRaC hooks to execute.
+        new TracingUtils();
+    }
+
     /**
      * Put an annotation to the current subsegment with a String value.
      *
@@ -192,4 +207,15 @@ public static void defaultObjectMapper(ObjectMapper objectMapper) {
     public static ObjectMapper objectMapper() {
         return objectMapper;
     }
+
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        init();
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        // This is a no-op, as we don't need to do anything after restore
+    }
 }