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

pom.xml

@@ -102,6 +102,7 @@
         <versions-maven-plugin.version>2.18.0</versions-maven-plugin.version>
         <elastic.version>1.6.0</elastic.version>
         <mockito.version>5.12.0</mockito.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. -->
@@ -249,6 +250,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,74 @@
+/*
+ * 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";
+
+    /**
+     * 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,53 @@
+package software.amazon.lambda.powertools.common.internal;
+
+import static org.mockito.Mockito.*;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.net.URL;
+import java.net.URLConnection;
+import java.util.Collections;
+
+import org.junit.jupiter.api.Test;
+
+class ClassPreLoaderTest {
+
+    @Test
+    void preloadClasses_shouldIgnoreInvalidClassesAndLoadValidClasses() throws Exception {
+        // Mock the class loader with no resources
+        ClassLoader classLoader = mock(ClassLoader.class);
+        URL mockUrl = mock(URL.class);
+        URLConnection mockConnection = mock(URLConnection.class);
+        InputStream mockInputStream = new ByteArrayInputStream("java.lang.String\nInvalid.Class".getBytes());
+
+        when(mockUrl.openConnection()).thenReturn(mockConnection);
+        when(mockConnection.getInputStream()).thenReturn(mockInputStream);
+        when(classLoader.getResources(ClassPreLoader.CLASSES_FILE))
+                .thenReturn(Collections.enumeration(Collections.singletonList(mockUrl)));
+
+        // Inject the mocked class loader
+        Thread.currentThread().setContextClassLoader(classLoader);
+        // Call the method under test
+        ClassPreLoader.preloadClasses();
+
+        // Verify that only the valid class was loaded
+        Class.forName("java.lang.String", true, ClassPreLoader.class.getClassLoader());
+    }
+
+    @Test
+    void preloadClasses_shouldHandleEmptyResources() throws Exception {
+        // Mock the class loader with no resources
+        ClassLoader classLoader = mock(ClassLoader.class);
+        when(classLoader.getResources(ClassPreLoader.CLASSES_FILE))
+                .thenReturn(Collections.emptyEnumeration());
+
+        // Inject the mocked class loader
+        Thread.currentThread().setContextClassLoader(classLoader);
+
+        // Call the method under test
+        ClassPreLoader.preloadClasses();
+
+        // Verify no interactions with the class loader
+        verifyNoInteractions(classLoader);
+    }
+}

powertools-idempotency/pom.xml

@@ -48,6 +48,10 @@
             <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,6 +23,7 @@
 
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.reflect.MethodSignature;
+import org.crac.Resource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,7 +47,7 @@
  * {@link software.amazon.lambda.powertools.idempotency.persistence.PersistenceStore}
  * to store the result of previous calls.
  */
-public class IdempotencyHandler {
+public class IdempotencyHandler implements Resource{
     private static final Logger LOG = LoggerFactory.getLogger(IdempotencyHandler.class);
     private static final int MAX_RETRIES = 2;
 
@@ -63,6 +64,27 @@ 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 {
+    }
+
     /**
      * Main entry point for handling idempotent execution of a function.
      *

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>

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

@@ -30,21 +30,28 @@
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
+import org.crac.Core;
+import org.crac.Resource;
 import software.amazon.cloudwatchlogs.emf.logger.MetricsLogger;
 import software.amazon.cloudwatchlogs.emf.model.DimensionSet;
 import software.amazon.cloudwatchlogs.emf.model.MetricsContext;
 import software.amazon.cloudwatchlogs.emf.model.Unit;
+import software.amazon.lambda.powertools.common.internal.ClassPreLoader;
 import software.amazon.lambda.powertools.common.internal.LambdaHandlerProcessor;
 import software.amazon.lambda.powertools.metrics.Metrics;
 import software.amazon.lambda.powertools.metrics.MetricsUtils;
 import software.amazon.lambda.powertools.metrics.ValidationException;
 
 @Aspect
-public class LambdaMetricsAspect {
+public class LambdaMetricsAspect implements Resource {
     public static final String TRACE_ID_PROPERTY = "xray_trace_id";
     public static final String REQUEST_ID_PROPERTY = "function_request_id";
     private static final String NAMESPACE = System.getenv("POWERTOOLS_METRICS_NAMESPACE");
 
+    public LambdaMetricsAspect() {
+        Core.getGlobalContext().register(this);
+    }
+
     private static String service(Metrics metrics) {
         return !"".equals(metrics.service()) ? metrics.service() : serviceName();
     }
@@ -148,4 +155,13 @@ private void validateMetricsAndRefreshOnFailure(Metrics metrics) {
             throw e;
         }
     }
+
+    @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 {
+    }
 }