Invoke prime Dynamo persistent store

Author: subhash686

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