feat: Support CRaC and priming of powertools large messages

docs/utilities/large_messages.md

@@ -586,3 +586,58 @@ it reads or deletes from S3. An empty allowlist, the default, applies no restric
 - The new module only provides an annotation, an equivalent to the `SqsUtils` class is not available anymore in this new version.
 
 Finally, if you are still using the `powertools-sqs` library for batch processing, consider moving to `powertools-batch` at the same time to remove the dependency on this library completely; it has been deprecated and will be removed in v2.
+
+## Advanced
+
+### Lambda SnapStart priming
+
+The LargerMessages class integrates with AWS Lambda SnapStart to improve restore durations. To make sure the SnapStart
+priming logic of this class runs correctly, you need an explicit reference to `LargerMessages` 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 `LargerMessages` in your code yet, make sure to reference `LargerMessages` 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() {
+            LargerMessages.init(); // Ensure LargerMessages 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 {
+            LargerMessages.init(); // Ensure LargerMessages 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-large-messages/pom.xml

@@ -14,11 +14,13 @@
   -->
 
 <project xmlns="http://maven.apache.org/POM/4.0.0"
-    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
     <modelVersion>4.0.0</modelVersion>
 
-    <description>A suite of utilities for AWS Lambda Functions that makes handling large messages in SQS and SNS easier.</description>
+    <description>A suite of utilities for AWS Lambda Functions that makes handling large messages in SQS and SNS
+        easier.
+    </description>
 
     <parent>
         <groupId>software.amazon.lambda</groupId>
@@ -87,6 +89,10 @@
             <artifactId>url-connection-client</artifactId>
             <version>${aws.sdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -141,6 +147,27 @@
         </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>
+    </profiles>
+
     <build>
         <plugins>
             <plugin>

powertools-large-messages/src/main/java/software/amazon/lambda/powertools/largemessages/LargeMessageConfig.java

@@ -16,6 +16,9 @@
 
 import static software.amazon.lambda.powertools.common.internal.LambdaConstants.AWS_REGION_ENV;
 
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
@@ -26,9 +29,11 @@
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 
 /**
- * Singleton instance for Large Message Config. We need this to provide a way to customize the S3 client configuration used by the annotation.
+ * Singleton instance for Large Message Config. We need this to provide a way to customize the S3 client
+ * configuration used by the annotation.
  * <br/>
- * Optional: Use it in your Lambda constructor to pass a custom {@link S3Client} to the {@link software.amazon.lambda.powertools.largemessages.internal.LargeMessageProcessor}
+ * Optional: Use it in your Lambda constructor to pass a custom {@link S3Client} to the
+ * {@link software.amazon.lambda.powertools.largemessages.internal.LargeMessageProcessor}
  * <br/>
  * If you don't use this, a default S3Client will be created.
  * <pre>
@@ -37,13 +42,14 @@
  * }
  * </pre>
  */
-public class LargeMessageConfig {
+public class LargeMessageConfig implements Resource {
 
     private static final LargeMessageConfig INSTANCE = new LargeMessageConfig();
     private S3Client s3Client;
     private Set<String> allowedBuckets = Collections.emptySet();
 
     private LargeMessageConfig() {
+        Core.getGlobalContext().register(this);
     }
 
     /**
@@ -120,4 +126,14 @@
         }
         return this.s3Client;
     }
+
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        getS3Client();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }

powertools-large-messages/src/main/java/software/amazon/lambda/powertools/largemessages/LargeMessages.java

@@ -17,23 +17,30 @@
 import java.util.Optional;
 import java.util.function.Function;
 
+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;
 import software.amazon.lambda.powertools.largemessages.internal.LargeMessageProcessor;
 import software.amazon.lambda.powertools.largemessages.internal.LargeMessageProcessorFactory;
 
 /**
  * Functional API for processing large messages without AspectJ.
  * <p>
  * Use this class to handle large messages (> 1 MB) from SQS or SNS.
- * When large messages are sent to an SQS Queue or SNS Topic, they are offloaded to S3 and only a reference is passed in the message/record.
+ * When large messages are sent to an SQS Queue or SNS Topic, they are offloaded to S3 and only a reference is passed
+ * in the message/record.
  * <p>
  * {@code LargeMessages} automatically retrieves and optionally deletes messages
- * which have been offloaded to S3 using the {@code amazon-sqs-java-extended-client-lib} or {@code amazon-sns-java-extended-client-lib}
+ * which have been offloaded to S3 using the {@code amazon-sqs-java-extended-client-lib} or {@code amazon-sns-java
+ * -extended-client-lib}
  * client libraries.
  * <p>
- * This version is compatible with version 1.1.0+ and 2.0.0+ of {@code amazon-sqs-java-extended-client-lib} / {@code amazon-sns-java-extended-client-lib}.
+ * This version is compatible with version 1.1.0+ and 2.0.0+ of {@code amazon-sqs-java-extended-client-lib} / {@code
+ * amazon-sns-java-extended-client-lib}.
  * <p>
  * <u>SQS Example</u>:
  * <pre>
@@ -88,12 +95,22 @@
  *
  * @see LargeMessage
  */
-public final class LargeMessages {
+public final class LargeMessages implements Resource {
 
     private static final Logger LOG = LoggerFactory.getLogger(LargeMessages.class);
 
-    private LargeMessages() {
-        // Utility class
+    // Dummy instance to register LargeMessages with CRaC
+    private static final LargeMessages INSTANCE = new LargeMessages();
+
+    // Static block to ensure CRaC registration happens at class loading time
+    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 LargeMessages();
     }
 
     /**
@@ -108,10 +125,10 @@
      * String returnValueOfFunction = LargeMessages.processLargeMessage(sqsMessage, processedMsg -&gt; processOrder(processedMsg, orderId, amount));
      * </pre>
      *
-     * @param message the message to process (SQSMessage or SNSRecord)
+     * @param message  the message to process (SQSMessage or SNSRecord)
      * @param function the function to execute with the processed message
-     * @param <T> the message type
-     * @param <R> the return type of the function
+     * @param <T>      the message type
+     * @param <R>      the return type of the function
      * @return the result of the function execution
      */
     public static <T, R> R processLargeMessage(T message, Function<T, R> function) {
@@ -127,11 +144,11 @@
      * String returnValueOfFunction = LargeMessages.processLargeMessage(sqsMessage, processedMsg -&gt; processOrder(processedMsg, orderId, amount), false);
      * </pre>
      *
-     * @param message the message to process (SQSMessage or SNSRecord)
-     * @param function the function to execute with the processed message
+     * @param message        the message to process (SQSMessage or SNSRecord)
+     * @param function       the function to execute with the processed message
      * @param deleteS3Object whether to delete the S3 object after processing
-     * @param <T> the message type
-     * @param <R> the return type of the function
+     * @param <T>            the message type
+     * @param <R>            the return type of the function
      * @return the result of the function execution
      */
     public static <T, R> R processLargeMessage(T message, Function<T, R> function, boolean deleteS3Object) {
@@ -153,4 +170,15 @@
             throw new LargeMessageProcessingException("Failed to process large message", t);
         }
     }
+
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        init();
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }