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). 

pom.xml

@@ -117,6 +117,7 @@
         <mockito.version>5.18.0</mockito.version>
         <mockito-junit-jupiter.version>5.18.0</mockito-junit-jupiter.version>
         <junit-pioneer.version>2.3.0</junit-pioneer.version>
+        <crac.version>1.4.0</crac.version>
 
         <!-- As we have a .mvn directory at the root of the project, this will evaluate to the root directory
             regardless of where maven is run - sub-module, or root. -->
@@ -264,6 +265,11 @@
                 <artifactId>logback-ecs-encoder</artifactId>
                 <version>${elastic.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.crac</groupId>
+                <artifactId>crac</artifactId>
+                <version>${crac.version}</version>
+            </dependency>
             <dependency>
                 <groupId>org.slf4j</groupId>
                 <artifactId>slf4j-api</artifactId>

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

@@ -0,0 +1,77 @@
+/*
+ * 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.common.internal;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.net.URL;
+import java.net.URLConnection;
+import java.util.Enumeration;
+
+/**
+ * Used to preload classes to support automatic priming for SnapStart
+ */
+public class ClassPreLoader {
+    public static final String CLASSES_FILE = "classesloaded.txt";
+
+    private ClassPreLoader() {
+        // Hide default constructor
+    }
+    /**
+     * Initializes the classes listed in the classesloaded resource
+     */
+    public static void preloadClasses() {
+        try {
+            Enumeration<URL> files = ClassPreLoader.class.getClassLoader().getResources(CLASSES_FILE);
+            // If there are multiple files, preload classes from all of them
+            while (files.hasMoreElements()) {
+                URL url = files.nextElement();
+                URLConnection conn = url.openConnection();
+                conn.setUseCaches(false);
+                InputStream is = conn.getInputStream();
+                preloadClassesFromStream(is);
+            }
+        } catch (IOException ignored) {
+            // No action is required if preloading fails for any reason
+        }
+    }
+
+    /**
+     * Loads the list of classes passed as a stream
+     *
+     * @param is
+     */
+    private static void preloadClassesFromStream(InputStream is) {
+        try (is;
+             InputStreamReader isr = new InputStreamReader(is, StandardCharsets.UTF_8);
+             BufferedReader reader = new BufferedReader(isr)) {
+            String line;
+            while ((line = reader.readLine()) != null) {
+                int idx = line.indexOf('#');
+                if (idx != -1) {
+                    line = line.substring(0, idx);
+                }
+                final String className = line.stripTrailing();
+                if (!className.isBlank()) {
+                    Class.forName(className, true, ClassPreLoader.class.getClassLoader());
+                }
+            }
+        } catch (Exception ignored) {
+            // No action is required if preloading fails for any reason
+        }
+    }
+}

powertools-common/src/test/java/software/amazon/lambda/powertools/common/internal/ClassPreLoaderTest.java

@@ -0,0 +1,14 @@
+package software.amazon.lambda.powertools.common.internal;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+
+class ClassPreLoaderTest {
+
+    @Test
+    void preloadClasses_shouldIgnoreInvalidClassesAndLoadValidClasses() {
+        // Verify that the missing class did not throw any exception
+        assertDoesNotThrow(ClassPreLoader::preloadClasses);
+    }
+}

powertools-common/src/test/resources/classesloaded.txt

@@ -0,0 +1 @@
+software.amazon.lambda.powertools.common.internal.NonExistingClass

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

@@ -39,6 +39,10 @@
             <artifactId>aspectjrt</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+        </dependency>
         <dependency>
             <groupId>software.amazon.lambda</groupId>
             <artifactId>powertools-common</artifactId>
@@ -114,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
+    }
 }