feat: Support CRaC and priming of powertools metrics and idempotency-dynamodb

Priming.md

@@ -0,0 +1,58 @@
+# Automatic Priming for AWS Lambda Powertools Java
+
+## Table of Contents
+- [Overview](#overview)
+- [Implementation Steps](#general-implementation-steps)
+- [Known Issues](#known-issues-and-solutions)
+- [Reference Implementation](#reference-implementation)
+
+## Overview
+Priming is the process of preloading dependencies and initializing resources during the INIT phase, rather than during the INVOKE phase to further optimize startup performance with SnapStart. 
+This is required because Java frameworks that use dependency injection load classes into memory when these classes are explicitly invoked, which typically happens during Lambda’s INVOKE phase.
+
+This documentation provides guidance for automatic class priming in AWS Lambda Powertools Java modules.
+
+
+## Implementation Steps
+Classes are proactively loaded using Java runtime hooks which are part of the open source CRaC (Coordinated Restore at Checkpoint) project.
+This implementation uses this hook, called `beforeCheckpoint()`, to prime SnapStart-enabled Java functions via Class Priming.
+In order to generate the `classloaded` files for Powertools java module, follow these general steps.
+
+1. **Add Maven Profile**
+    - Add maven test profile with the following VM argument for generating classes loaded files. 
+   ```shell
+      -Xlog:class+load=info:classesloaded.txt
+      ```
+    - You can find an example of this in profile `generate-classesloaded-file` in this [pom.xml](powertools-metrics/pom.xml).
+
+2. **Generate classes loaded file**
+    - Run tests with `-Pgenerate-classesloaded-file` profile.
+   ```shell
+   mvn -Pgenerate-classesloaded-file clean test
+   ```
+   - This will generate a file named `classesloaded.txt` in the target directory of the module.
+
+3. **Cleanup the file**
+    - The classes loaded file generated in Step 2 has the format
+   `[0.054s][info][class,load] java.lang.Object source: shared objects file` 
+   but we are only interested in `java.lang.Object` the fully qualified class name.
+    - To strip the lines to include only the fully qualified class name, 
+      Use the following regex to replace with empty string.
+      - `^\[[\[\]0-9.a-z,]+ ` (to replace the left part)
+      - `( source: )[0-9a-z :/._$-]+` (to replace the right part)
+
+4. **Place the file**
+    -  Move the cleanedup file containing to the corresponding resources directory of the module. See [example](powertools-metrics/src/resources/classesloaded.txt).
+
+5. **Register and checkpoint**
+    - Register the CRaC Resource in the constructor of the main class.
+    - Add the `beforeCheckpoint()` hook in the same class to invoke `ClassPreLoader.preloadClasses()`.
+    - This will ensure that the classes are loaded during the INIT phase of the Lambda function. See [example](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java)
+  
+
+## Known Issues
+- This is a manual process at the moment, but it can be automated in the future.
+- Because the file is generated while running tests, it includes test classes as well.
+
+## Reference Implementation
+Working example is available in the [powertools-metrics](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java). 

powertools-idempotency/pom.xml

@@ -46,10 +46,6 @@
             <groupId>software.amazon.lambda</groupId>
             <artifactId>powertools-common</artifactId>
         </dependency>
-        <dependency>
-            <groupId>org.crac</groupId>
-            <artifactId>crac</artifactId>
-        </dependency>
         <!-- Test dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>

powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java

@@ -23,7 +23,6 @@
 
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.reflect.MethodSignature;
-import org.crac.Resource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -47,7 +46,7 @@
  * {@link software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore}
  * to store the result of previous calls.
  */
-public class IdempotencyHandler implements Resource{
+public class IdempotencyHandler {
     private static final Logger LOG = LoggerFactory.getLogger(IdempotencyHandler.class);
     private static final int MAX_RETRIES = 2;
 
@@ -64,28 +63,6 @@ public IdempotencyHandler(ProceedingJoinPoint pjp, String functionName, JsonNode
         persistenceStore.configure(Idempotency.getInstance().getConfig(), functionName);
     }
 
-    /**
-     * Primes the persistent store by invoking the get record method with a key that doesn't exist.
-     *
-     * @param context
-     * @throws Exception
-     */
-    @Override
-    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
-        try {
-            persistenceStore.getRecord("__invoke_prime__");
-        } catch (IdempotencyItemNotFoundException keyNotFound) {
-            // This is expected, as we are priming the store
-        } catch (Exception unknown) {
-            // This is unexpected but we must continue without any interruption
-        }
-    }
-
-    @Override
-    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
-        // This is a no-op, as we don't need to do anything after restore
-    }
-
     /**
      * Main entry point for handling idempotent execution of a function.
      *

powertools-idempotency/powertools-idempotency-dynamodb/pom.xml

@@ -57,6 +57,10 @@
                 </exclusion>
             </exclusions>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>

powertools-idempotency/powertools-idempotency-dynamodb/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/dynamodb/DynamoDBPersistenceStore.java

@@ -14,6 +14,8 @@
 
 package software.amazon.lambda.powertools.idempotency.persistence.dynamodb;
 
+import org.crac.Core;
+import org.crac.Resource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
@@ -37,6 +39,7 @@
 import software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore;
 
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.AbstractMap;
 import java.util.HashMap;
 import java.util.Map;
@@ -52,7 +55,7 @@
  * DynamoDB version of the {@link PersistenceStore}. Will store idempotency data in DynamoDB.<br>
  * Use the {@link Builder} to create a new instance.
  */
-public class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore {
+public class DynamoDBPersistenceStore extends BasePersistenceStore implements PersistenceStore, Resource {
 
     public static final String IDEMPOTENCY = "idempotency";
     private static final Logger LOG = LoggerFactory.getLogger(DynamoDBPersistenceStore.class);
@@ -109,6 +112,39 @@
                 this.dynamoDbClient = null;
             }
         }
+        Core.getGlobalContext().register(this);
+    }
+
+    /**
+     * Primes the persistent store by invoking the get record method with a key that doesn't exist.
+     *
+     * @param context
+     * @throws Exception
+     */
+    @Override
+    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
+        try {
+            String primingRecordKey = "__invoke_prime__";
+            Instant now = Instant.now();
+            long expiry = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
+            DataRecord primingDataRecord = new DataRecord(
+                    primingRecordKey,
+                    DataRecord.Status.COMPLETED,
+                    expiry,
+                    null, // no data
+                    null // no validation
+            );
+            putRecord(primingDataRecord, Instant.now());
+            getRecord(primingRecordKey);
+            deleteRecord(primingRecordKey);
+        } catch (Exception unknown) {
+            // This is unexpected but we must continue without any interruption
+        }
+    }
+
+    @Override
+    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
+        // This is a no-op, as we don't need to do anything after restore
     }
 
     public static Builder builder() {

powertools-metrics/pom.xml

@@ -118,6 +118,23 @@
     </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>
             <dependencies>

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java

@@ -14,6 +14,9 @@
 
 package software.amazon.lambda.powertools.metrics;
 
+import org.crac.Core;
+import org.crac.Resource;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.common.internal.LambdaConstants;
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
 import software.amazon.lambda.powertools.metrics.model.DimensionSet;
@@ -23,11 +26,12 @@
 /**
  * Factory for accessing the singleton Metrics instance
  */
-public final class MetricsFactory {
+public final class MetricsFactory implements Resource {
     private static MetricsProvider provider = new EmfMetricsProvider();
     private static Metrics metrics;
 
     private MetricsFactory() {
+        Core.getGlobalContext().register(this);
     }
 
     /**
@@ -68,4 +72,14 @@ public static synchronized void setMetricsProvider(MetricsProvider metricsProvid
         // Reset the metrics instance so it will be recreated with the new provider
         metrics = null;
     }
+
+    @Override
+    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }

powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/internal/LambdaMetricsAspect.java

@@ -23,12 +23,8 @@
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
 
-import org.crac.Core;
-import org.crac.Resource;
-
 import com.amazonaws.services.lambda.runtime.Context;
 
-import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.common.internal.LambdaConstants;
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
 import software.amazon.lambda.powertools.metrics.FlushMetrics;
@@ -37,16 +33,12 @@
 import software.amazon.lambda.powertools.metrics.model.DimensionSet;
 
 @Aspect
-public class LambdaMetricsAspect implements Resource {
+public class LambdaMetricsAspect {
     public static final String TRACE_ID_PROPERTY = "xray_trace_id";
     public static final String REQUEST_ID_PROPERTY = "function_request_id";
     private static final String SERVICE_DIMENSION = "Service";
     private static final String FUNCTION_NAME_ENV_VAR = "POWERTOOLS_METRICS_FUNCTION_NAME";
 
-    public LambdaMetricsAspect() {
-        Core.getGlobalContext().register(this);
-    }
-
     private String functionName(FlushMetrics metrics, Context context) {
         if (!"".equals(metrics.functionName())) {
             return metrics.functionName();
@@ -147,14 +139,4 @@ private void captureColdStartMetricIfEnabled(Context extractedContext, FlushMetr
             metricsInstance.captureColdStartMetric(extractedContext, coldStartDimensions);
         }
     }
-
-    @Override
-    public void beforeCheckpoint(org.crac.Context<? extends Resource> context) throws Exception {
-        ClassPreLoader.preloadClasses();
-    }
-
-    @Override
-    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
-        // No action needed after restore
-    }
 }