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

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/pom.xml

@@ -46,6 +46,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,28 @@
         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.
      *
@@ -204,7 +227,7 @@
         Object response;
         try {
             response = pjp.proceed(pjp.getArgs());
         } catch (Throwable handlerException) {
             // We need these nested blocks to preserve function's exception in case the persistence store operation
             // also raises an exception
             try {

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

@@ -23,8 +23,12 @@
 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;
@@ -33,12 +37,16 @@
 import software.amazon.lambda.powertools.metrics.model.DimensionSet;
 
 @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 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();
@@ -139,4 +147,14 @@ 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
+    }
 }