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

Priming.md

@@ -0,0 +1,59 @@
+# 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 Powertools for AWS Lambda 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](https://openjdk.org/projects/crac/).
+Implementations across the project use the `beforeCheckpoint()` hook, to prime Snapstart-enabled Java functions via Class Priming.
+In order to generate the `classloaded.txt` file for a Java module in this project, 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 `generate-classesloaded-file` profile 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. **Add file to resources**
+    -  Move the cleaned-up file to the corresponding `src/main/resources` directory of the module. See [example](powertools-metrics/src/main/resources/classesloaded.txt).
+
+5. **Register and checkpoint**
+    - A class, usually the entry point of the module, should register the CRaC resource in the constructor. [Example](powertools-metrics/src/main/java/software/amazon/lambda/powertools/metrics/MetricsFactory.java)
+      - Note that AspectJ aspect is not suitable for this purpose, as it does not work with CRaC.
+    - Add the `beforeCheckpoint()` hook in the same class to invoke `ClassPreLoader.preloadClasses()`. The `ClassPreLoader` class is implemented in `powertools-common` module.
+    - This will ensure that the classes are already pre-loaded by the Snapstart RESTORE operation leading to a shorter INIT duration.
+
+
+## Known Issues
+- This is a manual process at the moment, but it can be automated in the future.
+- `classesloaded.txt` file includes test classes as well because the file is generated while running tests. This is not a problem because all the classes that are not found are ignored by `ClassPreLoader.preloadClasses()`. Also `beforeCheckpoint()` hook is not time-sensitive, it only runs once when a new Lambda version gets published. 
+
+## 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,91 @@
+/*
+ * 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 final 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()) {
+                    loadClassIfFound(className);
+                }
+            }
+        } catch (Exception ignored) {
+            // No action is required if preloading fails for any reason
+        }
+    }
+
+    /**
+     * Initializes the class with given name if found, ignores otherwise
+     *
+     * @param className
+     */
+    private static void loadClassIfFound(String className) {
+        try {
+            Class.forName(className, true, ClassPreLoader.class.getClassLoader());
+        } catch (ClassNotFoundException e) {
+            // No action is required if the class with given name cannot be found
+        }
+    }
+}

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

@@ -0,0 +1,34 @@
+package software.amazon.lambda.powertools.common.internal;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+
+class ClassPreLoaderTest {
+
+    // Making this volatile so the Thread Context doesn't need any special handling
+    static volatile boolean dummyClassLoaded = false;
+
+    /**
+     * Dummy class to be loaded by ClassPreLoader in test.
+     * <b>The class name is referenced in <i>powertools-common/src/test/resources/classesloaded.txt</i></b>
+     * This class is used to verify that the ClassPreLoader can load valid classes.
+     * The static block sets a flag to indicate that the class has been loaded.
+     */
+    static class DummyClass {
+        static {
+            dummyClassLoaded = true;
+        }
+    }
+    @Test
+    void preloadClasses_shouldIgnoreInvalidClassesAndLoadValidClasses() {
+
+        dummyClassLoaded = false;
+        // powertools-common/src/test/resources/classesloaded.txt has a class that does not exist
+        // Verify that the missing class did not throw any exception
+        assertDoesNotThrow(ClassPreLoader::preloadClasses);
+
+        // When the classloaded.txt is a mixed bag of valid and invalid classes, Valid class must load
+        assertTrue(dummyClassLoaded, "DummyClass should be loaded");
+    }
+}

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

@@ -0,0 +1,2 @@
+software.amazon.lambda.powertools.common.internal.NonExistingClass
+software.amazon.lambda.powertools.common.internal.ClassPreLoaderTest$DummyClass

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 final 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-logging/powertools-logging-logback/src/test/java/software/amazon/lambda/powertools/logging/internal/LambdaJsonEncoderTest.java

@@ -82,9 +82,9 @@
         // Make sure file is cleaned up before running tests
         try {
             FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
         } catch (NoSuchFileException e) {
             // file may not exist on the first launch
         }
     }
 
     @AfterEach
@@ -368,7 +368,7 @@
         String result = new String(encoded, StandardCharsets.UTF_8);
 
         // THEN
-        assertThat(result).contains("\"thread\":\"main\",\"thread_id\":1,\"thread_priority\":5");
+        assertThat(result).contains("\"thread\":\"main\",\"thread_id\":"+ Thread.currentThread().getId() +",\"thread_priority\":5");
     }
 
     @Test

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,16 @@
 /**
  * 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() {
+    // Dummy instance to register MetricsFactory with CRaC
+    private static final MetricsFactory INSTANCE = new MetricsFactory();
+
+    // Static block to ensure CRaC registration happens at class loading time
+    static {
+        Core.getGlobalContext().register(INSTANCE);
     }
 
     /**
@@ -68,4 +76,15 @@ 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 {
+        MetricsFactory.getMetricsInstance();
+        ClassPreLoader.preloadClasses();
+    }
+
+    @Override
+    public void afterRestore(org.crac.Context<? extends Resource> context) throws Exception {
+        // No action needed after restore
+    }
 }