From a751eaa6ead5e817ac4e3c25a5de3611223e8059 Mon Sep 17 00:00:00 2001
From: Jerome Van Der Linden <jeromevdl@gmail.com>
Date: Thu, 27 Oct 2022 17:15:31 +0200
Subject: [PATCH 1/7] fix idempotency when function timeouts

---
 powertools-idempotency/pom.xml                |  1 +
 .../powertools/idempotency/Idempotency.java   | 10 ++
 .../idempotency/IdempotencyConfig.java        | 15 ++-
 ...IdempotencyInconsistentStateException.java |  2 +-
 .../internal/IdempotencyHandler.java          | 28 +++++-
 .../internal/IdempotentAspect.java            | 19 +++-
 .../persistence/BasePersistenceStore.java     | 17 ++--
 .../idempotency/persistence/DataRecord.java   | 22 ++++-
 .../persistence/DynamoDBPersistenceStore.java | 48 +++++++--
 .../handlers/IdempotencyInternalFunction.java | 10 ++
 .../internal/IdempotencyAspectTest.java       | 87 ++++++++++++++--
 .../persistence/BasePersistenceStoreTest.java | 47 +++++++--
 .../DynamoDBPersistenceStoreTest.java         | 99 ++++++++++++++++++-
 spotbugs-exclude.xml                          |  8 ++
 14 files changed, 371 insertions(+), 42 deletions(-)

diff --git a/powertools-idempotency/pom.xml b/powertools-idempotency/pom.xml
index 7ddc1e793..536ef0e9d 100644
--- a/powertools-idempotency/pom.xml
+++ b/powertools-idempotency/pom.xml
@@ -122,6 +122,7 @@
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-lambda-java-tests</artifactId>
+            <scope>test</scope>
         </dependency>
         <dependency>
             <groupId>com.amazonaws</groupId>
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/Idempotency.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/Idempotency.java
index 1ff2ed47f..ce652791b 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/Idempotency.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/Idempotency.java
@@ -65,6 +65,16 @@ public static Idempotency getInstance() {
         return Holder.instance;
     }
 
+    /**
+     * Can be used in a method which is not the handler to capture the Lambda context,
+     * to calculate the remaining time before the invocation times out.
+     *
+     * @param lambdaContext
+     */
+    public static void registerLambdaContext(Context lambdaContext) {
+        getInstance().getConfig().setLambdaContext(lambdaContext);
+    }
+
     /**
      * Acts like a builder that can be used to configure {@link Idempotency}
      *
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
index 4089d3ed8..aeadfe033 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
@@ -13,8 +13,10 @@
  */
 package software.amazon.lambda.powertools.idempotency;
 
+import com.amazonaws.services.lambda.runtime.Context;
 import software.amazon.lambda.powertools.idempotency.internal.cache.LRUCache;
 
+import java.security.MessageDigest;
 import java.time.Duration;
 
 /**
@@ -28,6 +30,7 @@ public class IdempotencyConfig {
     private final String payloadValidationJMESPath;
     private final boolean throwOnNoIdempotencyKey;
     private final String hashFunction;
+    private Context lambdaContext;
 
     private IdempotencyConfig(String eventKeyJMESPath, String payloadValidationJMESPath, boolean throwOnNoIdempotencyKey, boolean useLocalCache, int localCacheMaxItems, long expirationInSeconds, String hashFunction) {
         this.localCacheMaxItems = localCacheMaxItems;
@@ -71,12 +74,20 @@ public String getHashFunction() {
     /**
      * Create a builder that can be used to configure and create a {@link IdempotencyConfig}.
      *
-     * @return a new instance of {@link IdempotencyConfig.Builder}
+     * @return a new instance of {@link Builder}
      */
     public static Builder builder() {
         return new Builder();
     }
 
+    public void setLambdaContext(Context lambdaContext) {
+        this.lambdaContext = lambdaContext;
+    }
+
+    public Context getLambdaContext() {
+        return lambdaContext;
+    }
+
     public static class Builder {
 
         private int localCacheMaxItems = 256;
@@ -203,7 +214,7 @@ public Builder withThrowOnNoIdempotencyKey() {
         /**
          * Function to use for calculating hashes, by default MD5.
          *
-         * @param hashFunction Can be any algorithm supported by {@link java.security.MessageDigest}, most commons are<ul>
+         * @param hashFunction Can be any algorithm supported by {@link MessageDigest}, most commons are<ul>
          *                     <li>MD5</li>
          *                     <li>SHA-1</li>
          *                     <li>SHA-256</li></ul>
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/exceptions/IdempotencyInconsistentStateException.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/exceptions/IdempotencyInconsistentStateException.java
index c6fe38d23..40c90dcab 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/exceptions/IdempotencyInconsistentStateException.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/exceptions/IdempotencyInconsistentStateException.java
@@ -15,7 +15,7 @@
 
 /**
  * IdempotencyInconsistentStateException can happen under rare but expected cases
- * when persistent state changes in the small-time between put & get requests.
+ * when persistent state changes in the small-time between put &amp; get requests.
  */
 public class IdempotencyInconsistentStateException extends RuntimeException {
     private static final long serialVersionUID = -4293951999802300672L;
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
index 1f3724919..c1eab3ab5 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
@@ -13,6 +13,7 @@
  */
 package software.amazon.lambda.powertools.idempotency.internal;
 
+import com.amazonaws.services.lambda.runtime.Context;
 import com.fasterxml.jackson.databind.JsonNode;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.reflect.MethodSignature;
@@ -25,6 +26,8 @@
 import software.amazon.lambda.powertools.utilities.JsonConfig;
 
 import java.time.Instant;
+import java.util.Optional;
+import java.util.OptionalInt;
 
 import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.EXPIRED;
 import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.INPROGRESS;
@@ -40,10 +43,12 @@ public class IdempotencyHandler {
     private final ProceedingJoinPoint pjp;
     private final JsonNode data;
     private final BasePersistenceStore persistenceStore;
+    private final Context lambdaContext;
 
-    public IdempotencyHandler(ProceedingJoinPoint pjp, String functionName, JsonNode payload) {
+    public IdempotencyHandler(ProceedingJoinPoint pjp, String functionName, JsonNode payload, Context lambdaContext) {
         this.pjp = pjp;
         this.data = payload;
+        this.lambdaContext = lambdaContext;
         persistenceStore = Idempotency.getInstance().getPersistenceStore();
         persistenceStore.configure(Idempotency.getInstance().getConfig(), functionName);
     }
@@ -77,7 +82,7 @@ private Object processIdempotency() throws Throwable {
         try {
             // We call saveInProgress first as an optimization for the most common case where no idempotent record
             // already exists. If it succeeds, there's no need to call getRecord.
-            persistenceStore.saveInProgress(data, Instant.now());
+            persistenceStore.saveInProgress(data, Instant.now(), getRemainingTimeInMillis());
         } catch (IdempotencyItemAlreadyExistsException iaee) {
             DataRecord record = getIdempotencyRecord();
             return handleForStatus(record);
@@ -89,6 +94,21 @@ private Object processIdempotency() throws Throwable {
         return getFunctionResponse();
     }
 
+    /**
+     * Tries to determine the remaining time available for the current lambda invocation.
+     * Currently, it only works if the idempotent handler decorator is used or using {@link Idempotency#registerLambdaContext(Context)}
+     *
+     * @return the remaining time in milliseconds or empty if the context was not provided/found
+     */
+    private OptionalInt getRemainingTimeInMillis() {
+        if (lambdaContext != null) {
+            return OptionalInt.of(lambdaContext.getRemainingTimeInMillis());
+        } else {
+            LOG.warn("Couldn't determine the remaining time left. Did you call registerLambdaContext on Idempotency?");
+        }
+        return OptionalInt.empty();
+    }
+
     /**
      * Retrieve the idempotency record from the persistence layer.
      *
@@ -121,6 +141,10 @@ private Object handleForStatus(DataRecord record) {
         }
 
         if (INPROGRESS.equals(record.getStatus())) {
+            if (record.getInProgressExpiryTimestamp().isPresent()
+                    && record.getInProgressExpiryTimestamp().getAsLong() < Instant.now().toEpochMilli()) {
+                throw new IdempotencyInconsistentStateException("Item should have been expired in-progress because it already time-outed.");
+            }
             throw new IdempotencyAlreadyInProgressException("Execution already in progress with idempotency key: " + record.getIdempotencyKey());
         }
 
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java
index 088e81249..8bc3c8c8a 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotentAspect.java
@@ -19,7 +19,9 @@
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
 import org.aspectj.lang.reflect.MethodSignature;
+import com.amazonaws.services.lambda.runtime.Context;
 import software.amazon.lambda.powertools.idempotency.Constants;
+import software.amazon.lambda.powertools.idempotency.Idempotency;
 import software.amazon.lambda.powertools.idempotency.IdempotencyKey;
 import software.amazon.lambda.powertools.idempotency.Idempotent;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyConfigurationException;
@@ -56,12 +58,20 @@ public Object around(ProceedingJoinPoint pjp,
             throw new IdempotencyConfigurationException("The annotated method doesn't return anything. Unable to perform idempotency on void return type");
         }
 
-        JsonNode payload = getPayload(pjp, method);
+        boolean isHandler = (isHandlerMethod(pjp) && placedOnRequestHandler(pjp));
+        JsonNode payload = getPayload(pjp, method, isHandler);
         if (payload == null) {
             throw new IdempotencyConfigurationException("Unable to get payload from the method. Ensure there is at least one parameter or that you use @IdempotencyKey");
         }
 
-        IdempotencyHandler idempotencyHandler = new IdempotencyHandler(pjp, method.getName(), payload);
+        Context lambdaContext;
+        if (isHandler) {
+            lambdaContext = (Context) pjp.getArgs()[1];
+        } else {
+            lambdaContext = Idempotency.getInstance().getConfig().getLambdaContext();
+        }
+
+        IdempotencyHandler idempotencyHandler = new IdempotencyHandler(pjp, method.getName(), payload, lambdaContext);
         return idempotencyHandler.handle();
     }
 
@@ -71,11 +81,10 @@ public Object around(ProceedingJoinPoint pjp,
      * @param method the annotated method
      * @return the payload used for idempotency
      */
-    private JsonNode getPayload(ProceedingJoinPoint pjp, Method method) {
+    private JsonNode getPayload(ProceedingJoinPoint pjp, Method method, boolean isHandler) {
         JsonNode payload = null;
         // handleRequest or method with one parameter: get the first one
-        if ((isHandlerMethod(pjp) && placedOnRequestHandler(pjp))
-                || pjp.getArgs().length == 1) {
+        if (isHandler || pjp.getArgs().length == 1) {
             payload = JsonConfig.get().getObjectMapper().valueToTree(pjp.getArgs()[0]);
         } else {
             // Look for a parameter annotated with @IdempotencyKey
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
index db0aaa688..fe392262b 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
@@ -35,9 +35,7 @@
 import java.security.NoSuchAlgorithmException;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.Map;
-import java.util.Spliterator;
-import java.util.Spliterators;
+import java.util.*;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -131,19 +129,25 @@ public void saveSuccess(JsonNode data, Object result, Instant now) {
      * @param data Payload
      * @param now
      */
-    public void saveInProgress(JsonNode data, Instant now) throws IdempotencyItemAlreadyExistsException {
+    public void saveInProgress(JsonNode data, Instant now, OptionalInt remainingTimeInMs) throws IdempotencyItemAlreadyExistsException {
         String idempotencyKey = getHashedIdempotencyKey(data);
 
         if (retrieveFromCache(idempotencyKey, now) != null) {
             throw new IdempotencyItemAlreadyExistsException();
         }
 
+        OptionalLong inProgressExpirationMsTimestamp = OptionalLong.empty();
+        if (remainingTimeInMs.isPresent()) {
+            inProgressExpirationMsTimestamp = OptionalLong.of(now.plus(remainingTimeInMs.getAsInt(), ChronoUnit.MILLIS).toEpochMilli());
+        }
+
         DataRecord record = new DataRecord(
                 idempotencyKey,
                 DataRecord.Status.INPROGRESS,
                 getExpiryEpochSecond(now),
                 null,
-                getHashedPayload(data)
+                getHashedPayload(data),
+                inProgressExpirationMsTimestamp
         );
         LOG.debug("saving in progress record for idempotency key: {}", record.getIdempotencyKey());
         putRecord(record, now);
@@ -212,7 +216,8 @@ private String getHashedIdempotencyKey(JsonNode data) {
         }
 
         String hash = generateHash(node);
-        return functionName + "#" + hash;
+        hash = functionName + "#" + hash;
+        return hash;
     }
 
     private boolean isMissingIdemPotencyKey(JsonNode data) {
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DataRecord.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DataRecord.java
index b4f58a73d..934ec3d09 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DataRecord.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DataRecord.java
@@ -13,8 +13,12 @@
  */
 package software.amazon.lambda.powertools.idempotency.persistence;
 
+import software.amazon.lambda.powertools.idempotency.IdempotencyConfig;
+
 import java.time.Instant;
 import java.util.Objects;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
 
 /**
  * Data Class for idempotency records. This is actually the item that will be stored in the persistence layer.
@@ -25,6 +29,7 @@ public class DataRecord {
     private final long expiryTimestamp;
     private final String responseData;
     private final String payloadHash;
+    private final OptionalLong inProgressExpiryTimestamp;
 
     public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, String responseData, String payloadHash) {
         this.idempotencyKey = idempotencyKey;
@@ -32,6 +37,16 @@ public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, St
         this.expiryTimestamp = expiryTimestamp;
         this.responseData = responseData;
         this.payloadHash = payloadHash;
+        this.inProgressExpiryTimestamp = OptionalLong.empty();
+    }
+
+    public DataRecord(String idempotencyKey, Status status, long expiryTimestamp, String responseData, String payloadHash, OptionalLong inProgressExpiryTimestamp) {
+        this.idempotencyKey = idempotencyKey;
+        this.status = status.toString();
+        this.expiryTimestamp = expiryTimestamp;
+        this.responseData = responseData;
+        this.payloadHash = payloadHash;
+        this.inProgressExpiryTimestamp = inProgressExpiryTimestamp;
     }
 
     public String getIdempotencyKey() {
@@ -39,7 +54,7 @@ public String getIdempotencyKey() {
     }
 
     /**
-     * Check if data record is expired (based on expiration configured in the {@link software.amazon.lambda.powertools.idempotency.IdempotencyConfig})
+     * Check if data record is expired (based on expiration configured in the {@link IdempotencyConfig})
      *
      * @return Whether the record is currently expired or not
      */
@@ -60,6 +75,10 @@ public long getExpiryTimestamp() {
         return expiryTimestamp;
     }
 
+    public OptionalLong getInProgressExpiryTimestamp() {
+        return inProgressExpiryTimestamp;
+    }
+
     public String getResponseData() {
         return responseData;
     }
@@ -85,6 +104,7 @@ public int hashCode() {
         return Objects.hash(idempotencyKey, status, expiryTimestamp, responseData, payloadHash);
     }
 
+
     /**
      * Status of the record:
      * <ul>
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
index 4985f845d..244a1420a 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
@@ -27,14 +27,12 @@
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemNotFoundException;
 
 import java.time.Instant;
-import java.util.AbstractMap;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static software.amazon.lambda.powertools.idempotency.Constants.AWS_REGION_ENV;
+import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.INPROGRESS;
 
 /**
  * DynamoDB version of the {@link PersistenceStore}. Will store idempotency data in DynamoDB.<br>
@@ -49,6 +47,8 @@ public class DynamoDBPersistenceStore extends BasePersistenceStore implements Pe
     private final String staticPkValue;
     private final String sortKeyAttr;
     private final String expiryAttr;
+
+    private final String inProgressExpiryAttr;
     private final String statusAttr;
     private final String dataAttr;
     private final String validationAttr;
@@ -62,6 +62,7 @@ private DynamoDBPersistenceStore(String tableName,
                                      String staticPkValue,
                                      String sortKeyAttr,
                                      String expiryAttr,
+                                     String inProgressExpiryAttr,
                                      String statusAttr,
                                      String dataAttr,
                                      String validationAttr,
@@ -71,6 +72,7 @@ private DynamoDBPersistenceStore(String tableName,
         this.staticPkValue = staticPkValue;
         this.sortKeyAttr = sortKeyAttr;
         this.expiryAttr = expiryAttr;
+        this.inProgressExpiryAttr = inProgressExpiryAttr;
         this.statusAttr = statusAttr;
         this.dataAttr = dataAttr;
         this.validationAttr = validationAttr;
@@ -115,6 +117,11 @@ public void putRecord(DataRecord record, Instant now) throws IdempotencyItemAlre
         Map<String, AttributeValue> item = new HashMap<>(getKey(record.getIdempotencyKey()));
         item.put(this.expiryAttr, AttributeValue.builder().n(String.valueOf(record.getExpiryTimestamp())).build());
         item.put(this.statusAttr, AttributeValue.builder().s(record.getStatus().toString()).build());
+
+        if (record.getInProgressExpiryTimestamp().isPresent()) {
+            item.put(this.inProgressExpiryAttr, AttributeValue.builder().n(String.valueOf(record.getInProgressExpiryTimestamp().getAsLong())).build());
+        }
+
         if (this.payloadValidationEnabled) {
             item.put(this.validationAttr, AttributeValue.builder().s(record.getPayloadHash()).build());
         }
@@ -124,16 +131,24 @@ public void putRecord(DataRecord record, Instant now) throws IdempotencyItemAlre
 
             Map<String, String> expressionAttributeNames = Stream.of(
                             new AbstractMap.SimpleEntry<>("#id", this.keyAttr),
-                            new AbstractMap.SimpleEntry<>("#expiry", this.expiryAttr))
+                            new AbstractMap.SimpleEntry<>("#expiry", this.expiryAttr),
+                            new AbstractMap.SimpleEntry<>("#in_progress_expiry", this.inProgressExpiryAttr),
+                            new AbstractMap.SimpleEntry<>("#status", this.statusAttr))
                     .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
+            Map<String, AttributeValue> expressionAttributeValues = Stream.of(
+                    new AbstractMap.SimpleEntry<>(":now", AttributeValue.builder().n(String.valueOf(now.getEpochSecond())).build()),
+                    new AbstractMap.SimpleEntry<>(":inprogress", AttributeValue.builder().s(INPROGRESS.toString()).build())
+            ).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+
             dynamoDbClient.putItem(
                     PutItemRequest.builder()
                             .tableName(tableName)
                             .item(item)
-                            .conditionExpression("attribute_not_exists(#id) OR #expiry < :now")
+                            .conditionExpression("attribute_not_exists(#id) OR #expiry < :now OR (attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)")
                             .expressionAttributeNames(expressionAttributeNames)
-                            .expressionAttributeValues(Collections.singletonMap(":now", AttributeValue.builder().n(String.valueOf(now.getEpochSecond())).build()))
+                            .expressionAttributeValues(expressionAttributeValues)
                             .build()
             );
         } catch (ConditionalCheckFailedException e) {
@@ -212,12 +227,12 @@ private DataRecord itemToRecord(Map<String, AttributeValue> item) {
         // data and validation payload may be null
         AttributeValue data = item.get(this.dataAttr);
         AttributeValue validation = item.get(this.validationAttr);
-
         return new DataRecord(item.get(sortKeyAttr != null ? sortKeyAttr: keyAttr).s(),
                 DataRecord.Status.valueOf(item.get(this.statusAttr).s()),
                 Long.parseLong(item.get(this.expiryAttr).n()),
                 data != null ? data.s() : null,
-                validation != null ? validation.s() : null);
+                validation != null ? validation.s() : null,
+                item.get(this.inProgressExpiryAttr) != null ? OptionalLong.of(Long.parseLong(item.get(this.inProgressExpiryAttr).n())) : OptionalLong.empty());
     }
 
     public static Builder builder() {
@@ -238,6 +253,8 @@ public static class Builder {
         private String staticPkValue = String.format("idempotency#%s", funcEnv != null ? funcEnv : "");
         private String sortKeyAttr;
         private String expiryAttr = "expiration";
+
+        private String inProgressExpiryAttr = "in_progress_expiration";
         private String statusAttr = "status";
         private String dataAttr = "data";
         private String validationAttr = "validation";
@@ -256,7 +273,7 @@ public DynamoDBPersistenceStore build() {
             if (StringUtils.isEmpty(tableName)) {
                 throw new IllegalArgumentException("Table name is not specified");
             }
-            return new DynamoDBPersistenceStore(tableName, keyAttr, staticPkValue, sortKeyAttr, expiryAttr, statusAttr, dataAttr, validationAttr, dynamoDbClient);
+            return new DynamoDBPersistenceStore(tableName, keyAttr, staticPkValue, sortKeyAttr, expiryAttr, inProgressExpiryAttr, statusAttr, dataAttr, validationAttr, dynamoDbClient);
         }
 
         /**
@@ -315,6 +332,17 @@ public Builder withExpiryAttr(String expiryAttr) {
             return this;
         }
 
+        /**
+         * DynamoDB attribute name for in progress expiry timestamp (optional), by default "in_progress_expiration"
+         *
+         * @param inProgressExpiryAttr name of the attribute in the table
+         * @return the builder instance (to chain operations)
+         */
+        public Builder withInProgressExpiryAttr(String inProgressExpiryAttr) {
+            this.inProgressExpiryAttr = inProgressExpiryAttr;
+            return this;
+        }
+
         /**
          * DynamoDB attribute name for status (optional), by default "status"
          *
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyInternalFunction.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyInternalFunction.java
index 549d9e7ed..f3c1bdbc9 100644
--- a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyInternalFunction.java
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyInternalFunction.java
@@ -15,6 +15,7 @@
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.idempotency.Idempotency;
 import software.amazon.lambda.powertools.idempotency.IdempotencyKey;
 import software.amazon.lambda.powertools.idempotency.Idempotent;
 import software.amazon.lambda.powertools.idempotency.model.Basket;
@@ -25,10 +26,19 @@
  */
 public class IdempotencyInternalFunction implements RequestHandler<Product, Basket> {
 
+    private final boolean registerContext;
     private boolean called = false;
 
+    public IdempotencyInternalFunction(boolean registerContext) {
+        this.registerContext = registerContext;
+    }
+
     @Override
     public Basket handleRequest(Product input, Context context) {
+        if (registerContext) {
+            Idempotency.registerLambdaContext(context);
+        }
+        
         return createBasket("fake", input);
     }
 
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
index fc91c6c61..21dc2f698 100644
--- a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
@@ -16,6 +16,7 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junitpioneer.jupiter.SetEnvironmentVariable;
@@ -27,6 +28,7 @@
 import software.amazon.lambda.powertools.idempotency.IdempotencyConfig;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyAlreadyInProgressException;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyConfigurationException;
+import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyInconsistentStateException;
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemAlreadyExistsException;
 import software.amazon.lambda.powertools.idempotency.handlers.*;
 import software.amazon.lambda.powertools.idempotency.model.Basket;
@@ -36,6 +38,8 @@
 import software.amazon.lambda.powertools.utilities.JsonConfig;
 
 import java.time.Instant;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
 
 import static java.time.temporal.ChronoUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -67,17 +71,22 @@ public void firstCall_shouldPutInStore() {
 
         IdempotencyEnabledFunction function = new IdempotencyEnabledFunction();
 
+        when(context.getRemainingTimeInMillis()).thenReturn(30000);
+
         Product p = new Product(42, "fake product", 12);
         Basket basket = function.handleRequest(p, context);
         assertThat(basket.getProducts()).hasSize(1);
         assertThat(function.handlerCalled()).isTrue();
 
         ArgumentCaptor<JsonNode> nodeCaptor = ArgumentCaptor.forClass(JsonNode.class);
-        verify(store).saveInProgress(nodeCaptor.capture(), any());
+        ArgumentCaptor<OptionalInt> expiryCaptor = ArgumentCaptor.forClass(OptionalInt.class);
+        verify(store).saveInProgress(nodeCaptor.capture(), any(), expiryCaptor.capture());
         assertThat(nodeCaptor.getValue().get("id").asLong()).isEqualTo(p.getId());
         assertThat(nodeCaptor.getValue().get("name").asText()).isEqualTo(p.getName());
         assertThat(nodeCaptor.getValue().get("price").asDouble()).isEqualTo(p.getPrice());
 
+        assertThat(expiryCaptor.getValue().orElse(-1)).isEqualTo(30000);
+
         ArgumentCaptor<Basket> resultCaptor = ArgumentCaptor.forClass(Basket.class);
         verify(store).saveSuccess(any(), resultCaptor.capture(), any());
         assertThat(resultCaptor.getValue()).isEqualTo(basket);
@@ -93,7 +102,7 @@ public void secondCall_notExpired_shouldGetFromStore() throws JsonProcessingExce
                         .build()
                 ).configure();
 
-        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any());
+        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
 
         Product p = new Product(42, "fake product", 12);
         Basket b = new Basket(p);
@@ -124,7 +133,7 @@ public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressExcepti
                         .build()
                 ).configure();
 
-        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any());
+        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
 
         Product p = new Product(42, "fake product", 12);
         Basket b = new Basket(p);
@@ -133,7 +142,8 @@ public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressExcepti
                 DataRecord.Status.INPROGRESS,
                 Instant.now().plus(356, SECONDS).getEpochSecond(),
                 JsonConfig.get().getObjectMapper().writer().writeValueAsString(b),
-                null);
+                null,
+                OptionalLong.of(Instant.now().toEpochMilli() + 1000));
         doReturn(record).when(store).getRecord(any(), any());
 
         // THEN
@@ -141,6 +151,34 @@ public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressExcepti
         assertThatThrownBy(() -> function.handleRequest(p, context)).isInstanceOf(IdempotencyAlreadyInProgressException.class);
     }
 
+    @Test
+    public void secondCall_inProgress_lambdaTimeout_timeoutExpired_shouldThrowInconsistentState() throws JsonProcessingException {
+        // GIVEN
+        Idempotency.config()
+                .withPersistenceStore(store)
+                .withConfig(IdempotencyConfig.builder()
+                        .withEventKeyJMESPath("id")
+                        .build()
+                ).configure();
+
+        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
+
+        Product p = new Product(42, "fake product", 12);
+        Basket b = new Basket(p);
+        DataRecord record = new DataRecord(
+                "42",
+                DataRecord.Status.INPROGRESS,
+                Instant.now().plus(356, SECONDS).getEpochSecond(),
+                JsonConfig.get().getObjectMapper().writer().writeValueAsString(b),
+                null,
+                OptionalLong.of(Instant.now().toEpochMilli() - 100));
+        doReturn(record).when(store).getRecord(any(), any());
+
+        // THEN
+        IdempotencyEnabledFunction function = new IdempotencyEnabledFunction();
+        assertThatThrownBy(() -> function.handleRequest(p, context)).isInstanceOf(IdempotencyInconsistentStateException.class);
+    }
+
     @Test
     public void functionThrowException_shouldDeleteRecord_andThrowFunctionException() {
         // GIVEN
@@ -190,7 +228,37 @@ public void idempotencyOnSubMethodAnnotated_firstCall_shouldPutInStore() {
                 .configure();
 
         // WHEN
-        IdempotencyInternalFunction function = new IdempotencyInternalFunction();
+        boolean registerContext = true;
+        when(context.getRemainingTimeInMillis()).thenReturn(30000);
+
+        IdempotencyInternalFunction function = new IdempotencyInternalFunction(registerContext);
+        Product p = new Product(42, "fake product", 12);
+        Basket basket = function.handleRequest(p, context);
+
+        // THEN
+        assertThat(basket.getProducts()).hasSize(2);
+        assertThat(function.subMethodCalled()).isTrue();
+
+        ArgumentCaptor<JsonNode> nodeCaptor = ArgumentCaptor.forClass(JsonNode.class);
+        ArgumentCaptor<OptionalInt> expiryCaptor = ArgumentCaptor.forClass(OptionalInt.class);
+        verify(store).saveInProgress(nodeCaptor.capture(), any(), expiryCaptor.capture());
+        assertThat(nodeCaptor.getValue().asText()).isEqualTo("fake");
+        assertThat(expiryCaptor.getValue().orElse(-1)).isEqualTo(30000);
+
+        ArgumentCaptor<Basket> resultCaptor = ArgumentCaptor.forClass(Basket.class);
+        verify(store).saveSuccess(any(), resultCaptor.capture(), any());
+        assertThat(resultCaptor.getValue().getProducts()).contains(basket.getProducts().get(0), new Product(0, "fake", 0));
+    }
+
+    @Test
+    public void idempotencyOnSubMethodAnnotated_firstCall_contextNotRegistered_shouldPutInStore() {
+        Idempotency.config()
+                .withPersistenceStore(store)
+                .configure();
+
+        // WHEN
+        boolean registerContext = false;
+        IdempotencyInternalFunction function = new IdempotencyInternalFunction(registerContext);
         Product p = new Product(42, "fake product", 12);
         Basket basket = function.handleRequest(p, context);
 
@@ -199,8 +267,10 @@ public void idempotencyOnSubMethodAnnotated_firstCall_shouldPutInStore() {
         assertThat(function.subMethodCalled()).isTrue();
 
         ArgumentCaptor<JsonNode> nodeCaptor = ArgumentCaptor.forClass(JsonNode.class);
-        verify(store).saveInProgress(nodeCaptor.capture(), any());
+        ArgumentCaptor<OptionalInt> expiryCaptor = ArgumentCaptor.forClass(OptionalInt.class);
+        verify(store).saveInProgress(nodeCaptor.capture(), any(), expiryCaptor.capture());
         assertThat(nodeCaptor.getValue().asText()).isEqualTo("fake");
+        assertThat(expiryCaptor.getValue()).isEmpty();
 
         ArgumentCaptor<Basket> resultCaptor = ArgumentCaptor.forClass(Basket.class);
         verify(store).saveSuccess(any(), resultCaptor.capture(), any());
@@ -214,7 +284,7 @@ public void idempotencyOnSubMethodAnnotated_secondCall_notExpired_shouldGetFromS
                 .withPersistenceStore(store)
                 .configure();
 
-        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any());
+        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
 
         Product p = new Product(42, "fake product", 12);
         Basket b = new Basket(p);
@@ -227,7 +297,7 @@ public void idempotencyOnSubMethodAnnotated_secondCall_notExpired_shouldGetFromS
         doReturn(record).when(store).getRecord(any(), any());
 
         // WHEN
-        IdempotencyInternalFunction function = new IdempotencyInternalFunction();
+        IdempotencyInternalFunction function = new IdempotencyInternalFunction(false);
         Basket basket = function.handleRequest(p, context);
 
         // THEN
@@ -285,4 +355,5 @@ public void idempotencyOnSubMethodVoid_shouldThrowException() {
         // THEN
         assertThatThrownBy(() -> function.handleRequest(p, context)).isInstanceOf(IdempotencyConfigurationException.class);
     }
+
 }
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStoreTest.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStoreTest.java
index dac9a9288..e77a995a4 100644
--- a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStoreTest.java
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStoreTest.java
@@ -32,6 +32,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
+import java.util.OptionalInt;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -83,12 +84,30 @@ public void saveInProgress_defaultConfig() {
         persistenceStore.configure(IdempotencyConfig.builder().build(), null);
 
         Instant now = Instant.now();
-        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now);
+        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty());
         assertThat(dr.getStatus()).isEqualTo(DataRecord.Status.INPROGRESS);
         assertThat(dr.getExpiryTimestamp()).isEqualTo(now.plus(3600, ChronoUnit.SECONDS).getEpochSecond());
         assertThat(dr.getResponseData()).isNull();
         assertThat(dr.getIdempotencyKey()).isEqualTo("testFunction#47261bd5b456f400f8d191cfb3a7482f");
         assertThat(dr.getPayloadHash()).isEqualTo("");
+        assertThat(dr.getInProgressExpiryTimestamp()).isEmpty();
+        assertThat(status).isEqualTo(1);
+    }
+
+    @Test
+    public void saveInProgress_withRemainingTime() {
+        APIGatewayProxyRequestEvent event = EventLoader.loadApiGatewayRestEvent("apigw_event.json");
+        persistenceStore.configure(IdempotencyConfig.builder().build(), null);
+
+        int lambdaTimeoutMs = 30000;
+        Instant now = Instant.now();
+        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.of(lambdaTimeoutMs));
+        assertThat(dr.getStatus()).isEqualTo(DataRecord.Status.INPROGRESS);
+        assertThat(dr.getExpiryTimestamp()).isEqualTo(now.plus(3600, ChronoUnit.SECONDS).getEpochSecond());
+        assertThat(dr.getResponseData()).isNull();
+        assertThat(dr.getIdempotencyKey()).isEqualTo("testFunction#47261bd5b456f400f8d191cfb3a7482f");
+        assertThat(dr.getPayloadHash()).isEqualTo("");
+        assertThat(dr.getInProgressExpiryTimestamp().orElse(-1)).isEqualTo(now.plus(lambdaTimeoutMs, ChronoUnit.MILLIS).toEpochMilli());
         assertThat(status).isEqualTo(1);
     }
 
@@ -100,7 +119,7 @@ public void saveInProgress_jmespath() {
                 .build(), "myfunc");
 
         Instant now = Instant.now();
-        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now);
+        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty());
         assertThat(dr.getStatus()).isEqualTo(DataRecord.Status.INPROGRESS);
         assertThat(dr.getExpiryTimestamp()).isEqualTo(now.plus(3600, ChronoUnit.SECONDS).getEpochSecond());
         assertThat(dr.getResponseData()).isNull();
@@ -117,7 +136,7 @@ public void saveInProgress_jmespath_NotFound_shouldThrowException() {
                 .withThrowOnNoIdempotencyKey(true) // should throw
                 .build(), "");
         Instant now = Instant.now();
-        assertThatThrownBy(() -> persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now))
+        assertThatThrownBy(() -> persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty()))
                 .isInstanceOf(IdempotencyKeyException.class)
                 .hasMessageContaining("No data found to create a hashed idempotency key");
         assertThat(status).isEqualTo(-1);
@@ -130,7 +149,7 @@ public void saveInProgress_jmespath_NotFound_shouldNotThrowException() {
                 .withEventKeyJMESPath("unavailable")
                 .build(), "");
         Instant now = Instant.now();
-        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now);
+        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty());
         assertThat(dr.getStatus()).isEqualTo(DataRecord.Status.INPROGRESS);
         assertThat(status).isEqualTo(1);
     }
@@ -151,7 +170,7 @@ public void saveInProgress_withLocalCache_NotExpired_ShouldThrowException() {
                         now.plus(3600, ChronoUnit.SECONDS).getEpochSecond(),
                         null, null)
         );
-        assertThatThrownBy(() -> persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now))
+        assertThatThrownBy(() -> persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty()))
                 .isInstanceOf(IdempotencyItemAlreadyExistsException.class);
         assertThat(status).isEqualTo(-1);
     }
@@ -173,7 +192,7 @@ public void saveInProgress_withLocalCache_Expired_ShouldRemoveFromCache() {
                         now.minus(3, ChronoUnit.SECONDS).getEpochSecond(),
                         null, null)
         );
-        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now);
+        persistenceStore.saveInProgress(JsonConfig.get().getObjectMapper().valueToTree(event), now, OptionalInt.empty());
         assertThat(dr.getStatus()).isEqualTo(DataRecord.Status.INPROGRESS);
         assertThat(cache).isEmpty();
         assertThat(status).isEqualTo(1);
@@ -362,4 +381,20 @@ public void generateHashDouble_shouldGenerateMd5ofDouble() {
         String generatedHash = persistenceStore.generateHash(new DoubleNode(256.42));
         assertThat(generatedHash).isEqualTo(expectedHash);
     }
+
+    @Test
+    public void generateHashString_withSha256Algorithm_shouldGenerateSha256ofString() {
+        persistenceStore.configure(IdempotencyConfig.builder().withHashFunction("SHA-256").build(), null);
+        String expectedHash = "e6139efa88ef3337e901e826e6f327337f414860fb499d9f26eefcff21d719af"; // SHA-256(Lambda rocks)
+        String generatedHash = persistenceStore.generateHash(new TextNode("Lambda rocks"));
+        assertThat(generatedHash).isEqualTo(expectedHash);
+    }
+
+    @Test
+    public void generateHashString_unknownAlgorithm_shouldGenerateMd5ofString() {
+        persistenceStore.configure(IdempotencyConfig.builder().withHashFunction("HASH").build(), null);
+        String expectedHash = "70c24d88041893f7fbab4105b76fd9e1"; // MD5(Lambda rocks)
+        String generatedHash = persistenceStore.generateHash(new TextNode("Lambda rocks"));
+        assertThat(generatedHash).isEqualTo(expectedHash);
+    }
 }
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStoreTest.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStoreTest.java
index 0da516324..86e35cd33 100644
--- a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStoreTest.java
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStoreTest.java
@@ -57,6 +57,68 @@ public void putRecord_shouldCreateRecordInDynamoDB() throws IdempotencyItemAlrea
         assertThat(item.get("expiration").n()).isEqualTo(String.valueOf(expiry));
     }
 
+    @Test
+    public void putRecord_shouldCreateRecordInDynamoDB_IfPreviousExpired() {
+        key = Collections.singletonMap("id", AttributeValue.builder().s("key").build());
+
+        // GIVEN: Insert a fake item with same id and expired
+        Map<String, AttributeValue> item = new HashMap<>(key);
+        Instant now = Instant.now();
+        long expiry = now.minus(30, ChronoUnit.SECONDS).getEpochSecond();
+        item.put("expiration", AttributeValue.builder().n(String.valueOf(expiry)).build());
+        item.put("status", AttributeValue.builder().s(DataRecord.Status.COMPLETED.toString()).build());
+        item.put("data", AttributeValue.builder().s("Fake Data").build());
+        client.putItem(PutItemRequest.builder().tableName(TABLE_NAME).item(item).build());
+
+        // WHEN: call putRecord
+        long expiry2 = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
+        dynamoDBPersistenceStore.putRecord(
+                new DataRecord("key",
+                        DataRecord.Status.INPROGRESS,
+                        expiry2,
+                        null,
+                        null
+                ), now);
+
+        // THEN: an item is inserted
+        Map<String, AttributeValue> itemInDb = client.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
+        assertThat(itemInDb).isNotNull();
+        assertThat(itemInDb.get("status").s()).isEqualTo("INPROGRESS");
+        assertThat(itemInDb.get("expiration").n()).isEqualTo(String.valueOf(expiry2));
+    }
+
+    @Test
+    public void putRecord_shouldCreateRecordInDynamoDB_IfLambdaWasInProgressAndTimedOut() {
+        key = Collections.singletonMap("id", AttributeValue.builder().s("key").build());
+
+        // GIVEN: Insert a fake item with same id and progress expired (Lambda timed out before and we allow a new execution)
+        Map<String, AttributeValue> item = new HashMap<>(key);
+        Instant now = Instant.now();
+        long expiry = now.plus(30, ChronoUnit.SECONDS).getEpochSecond();
+        long progressExpiry = now.minus(30, ChronoUnit.SECONDS).getEpochSecond();
+        item.put("expiration", AttributeValue.builder().n(String.valueOf(expiry)).build());
+        item.put("status", AttributeValue.builder().s(DataRecord.Status.INPROGRESS.toString()).build());
+        item.put("data", AttributeValue.builder().s("Fake Data").build());
+        item.put("in_progress_expiration", AttributeValue.builder().n(String.valueOf(progressExpiry)).build());
+        client.putItem(PutItemRequest.builder().tableName(TABLE_NAME).item(item).build());
+
+        // WHEN: call putRecord
+        long expiry2 = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
+        dynamoDBPersistenceStore.putRecord(
+                new DataRecord("key",
+                        DataRecord.Status.INPROGRESS,
+                        expiry2,
+                        null,
+                        null
+                ), now);
+
+        // THEN: an item is inserted
+        Map<String, AttributeValue> itemInDb = client.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
+        assertThat(itemInDb).isNotNull();
+        assertThat(itemInDb.get("status").s()).isEqualTo("INPROGRESS");
+        assertThat(itemInDb.get("expiration").n()).isEqualTo(String.valueOf(expiry2));
+    }
+
     @Test
     public void putRecord_shouldThrowIdempotencyItemAlreadyExistsException_IfRecordAlreadyExist() {
         key = Collections.singletonMap("id", AttributeValue.builder().s("key").build());
@@ -65,7 +127,7 @@ public void putRecord_shouldThrowIdempotencyItemAlreadyExistsException_IfRecordA
         Map<String, AttributeValue> item = new HashMap<>(key);
         Instant now = Instant.now();
         long expiry = now.plus(30, ChronoUnit.SECONDS).getEpochSecond();
-        item.put("expiration", AttributeValue.builder().n(String.valueOf(expiry)).build());
+        item.put("expiration", AttributeValue.builder().n(String.valueOf(expiry)).build()); // not expired
         item.put("status", AttributeValue.builder().s(DataRecord.Status.COMPLETED.toString()).build());
         item.put("data", AttributeValue.builder().s("Fake Data").build());
         client.putItem(PutItemRequest.builder().tableName(TABLE_NAME).item(item).build());
@@ -89,6 +151,41 @@ public void putRecord_shouldThrowIdempotencyItemAlreadyExistsException_IfRecordA
         assertThat(itemInDb.get("data").s()).isEqualTo("Fake Data");
     }
 
+    @Test
+    public void putRecord_shouldThrowIdempotencyItemAlreadyExistsException_IfRecordAlreadyExistAndProgressNotExpiredAfterLambdaTimedOut() {
+        key = Collections.singletonMap("id", AttributeValue.builder().s("key").build());
+
+        // GIVEN: Insert a fake item with same id
+        Map<String, AttributeValue> item = new HashMap<>(key);
+        Instant now = Instant.now();
+        long expiry = now.plus(30, ChronoUnit.SECONDS).getEpochSecond(); // not expired
+        long progressExpiry = now.plus(30, ChronoUnit.SECONDS).getEpochSecond(); // not expired
+        item.put("expiration", AttributeValue.builder().n(String.valueOf(expiry)).build());
+        item.put("status", AttributeValue.builder().s(DataRecord.Status.INPROGRESS.toString()).build());
+        item.put("data", AttributeValue.builder().s("Fake Data").build());
+        item.put("in_progress_expiration", AttributeValue.builder().n(String.valueOf(progressExpiry)).build());
+        client.putItem(PutItemRequest.builder().tableName(TABLE_NAME).item(item).build());
+
+        // WHEN: call putRecord
+        long expiry2 = now.plus(3600, ChronoUnit.SECONDS).getEpochSecond();
+        assertThatThrownBy(() -> dynamoDBPersistenceStore.putRecord(
+                new DataRecord("key",
+                        DataRecord.Status.INPROGRESS,
+                        expiry2,
+                        null,
+                        null
+                ), now)
+        ).isInstanceOf(IdempotencyItemAlreadyExistsException.class);
+
+        // THEN: item was not updated, retrieve the initial one
+        Map<String, AttributeValue> itemInDb = client.getItem(GetItemRequest.builder().tableName(TABLE_NAME).key(key).build()).item();
+        assertThat(itemInDb).isNotNull();
+        assertThat(itemInDb.get("status").s()).isEqualTo("INPROGRESS");
+        assertThat(itemInDb.get("expiration").n()).isEqualTo(String.valueOf(expiry));
+        assertThat(itemInDb.get("data").s()).isEqualTo("Fake Data");
+    }
+
+
     //</editor-fold>
     // =================================================================
 
diff --git a/spotbugs-exclude.xml b/spotbugs-exclude.xml
index 0c8d1d8f8..fe4194bfd 100644
--- a/spotbugs-exclude.xml
+++ b/spotbugs-exclude.xml
@@ -32,6 +32,10 @@
                 <Class name="software.amazon.lambda.powertools.idempotency.Idempotency"/>
                 <Method name="getPersistenceStore"/>
             </And>
+            <And>
+                <Class name="software.amazon.lambda.powertools.idempotency.Idempotency"/>
+                <Method name="getConfig"/>
+            </And>
         </Or>
     </Match>
     <Match>
@@ -69,6 +73,10 @@
                 <Class name="software.amazon.lambda.powertools.idempotency.Idempotency$Config"/>
                 <Field name="store"/>
             </And>
+            <And>
+                <Class name="software.amazon.lambda.powertools.idempotency.Idempotency$Config"/>
+                <Field name="config"/>
+            </And>
             <And>
                 <Class name="software.amazon.lambda.powertools.idempotency.internal.IdempotencyHandler"/>
                 <Field name="pjp"/>

From 62d86f674e82575c918be486ba793730cb69b3c2 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 2 Nov 2022 11:27:49 +0100
Subject: [PATCH 2/7] build(deps): bump aws.sdk.version from 2.18.4 to 2.18.7
 (#967)

Bumps `aws.sdk.version` from 2.18.4 to 2.18.7.

Updates `software.amazon.awssdk:bom` from 2.18.4 to 2.18.7
- [Release notes](https://github.com/aws/aws-sdk-java-v2/releases)
- [Changelog](https://github.com/aws/aws-sdk-java-v2/blob/master/CHANGELOG.md)
- [Commits](https://github.com/aws/aws-sdk-java-v2/compare/2.18.4...2.18.7)

Updates `http-client-spi` from 2.18.4 to 2.18.7
- [Release notes](https://github.com/aws/aws-sdk-java-v2/releases)
- [Changelog](https://github.com/aws/aws-sdk-java-v2/blob/master/CHANGELOG.md)
- [Commits](https://github.com/aws/aws-sdk-java-v2/compare/2.18.4...2.18.7)

Updates `url-connection-client` from 2.18.4 to 2.18.7

---
updated-dependencies:
- dependency-name: software.amazon.awssdk:bom
  dependency-type: direct:production
  update-type: version-update:semver-patch
- dependency-name: software.amazon.awssdk:http-client-spi
  dependency-type: direct:production
  update-type: version-update:semver-patch
- dependency-name: software.amazon.awssdk:url-connection-client
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 pom.xml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pom.xml b/pom.xml
index 9af93bf6f..bd6968161 100644
--- a/pom.xml
+++ b/pom.xml
@@ -58,7 +58,7 @@
         <log4j.version>2.19.0</log4j.version>
         <jackson.version>2.13.4.2</jackson.version>
         <aspectj.version>1.9.7</aspectj.version>
-        <aws.sdk.version>2.18.4</aws.sdk.version>
+        <aws.sdk.version>2.18.7</aws.sdk.version>
         <aws.xray.recorder.version>2.12.0</aws.xray.recorder.version>
         <payloadoffloading-common.version>2.1.2</payloadoffloading-common.version>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>

From 997049fe57e0e5cd431dd1bb8dfd5ede3fbc6091 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Van=20Der=20Linden?=
 <jeromevdl@gmail.com>
Date: Wed, 2 Nov 2022 11:41:14 +0100
Subject: [PATCH 3/7] fix: envelope is not taken into account with built-in
 types (#960)

* add precedence for the envelope over built-in types

* cannot use same envelope for in and out

* remove extra new line

* handle event field names properly
ex: SQS Records with big R

* more explicit exception message
---
 powertools-validation/pom.xml                 | 10 +++----
 .../validation/ValidationUtils.java           | 12 +++++++-
 .../validation/internal/ValidationAspect.java | 12 ++++++--
 .../SQSWithCustomEnvelopeHandler.java         | 28 ++++++++++++++++++
 .../handlers/SQSWithWrongEnvelopeHandler.java | 29 +++++++++++++++++++
 .../internal/ValidationAspectTest.java        | 18 ++++++++++++
 .../src/test/resources/sqs_message.json       | 22 ++++++++++++++
 7 files changed, 122 insertions(+), 9 deletions(-)
 create mode 100644 powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithCustomEnvelopeHandler.java
 create mode 100644 powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithWrongEnvelopeHandler.java
 create mode 100644 powertools-validation/src/test/resources/sqs_message.json

diff --git a/powertools-validation/pom.xml b/powertools-validation/pom.xml
index 61beb34ba..c5a0a767c 100644
--- a/powertools-validation/pom.xml
+++ b/powertools-validation/pom.xml
@@ -75,6 +75,10 @@
             <artifactId>json-schema-validator</artifactId>
             <version>1.0.73</version>
         </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-lambda-java-serialization</artifactId>
+        </dependency>
 
         <!-- Test dependencies -->
         <dependency>
@@ -87,11 +91,7 @@
             <artifactId>junit-jupiter-engine</artifactId>
             <scope>test</scope>
         </dependency>
-        <dependency>
-            <groupId>com.amazonaws</groupId>
-            <artifactId>aws-lambda-java-serialization</artifactId>
-            <scope>test</scope>
-        </dependency>
+
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-lang3</artifactId>
diff --git a/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/ValidationUtils.java b/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/ValidationUtils.java
index 83f34ebfd..12c51c632 100644
--- a/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/ValidationUtils.java
+++ b/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/ValidationUtils.java
@@ -13,6 +13,7 @@
  */
 package software.amazon.lambda.powertools.validation;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collections;
@@ -21,9 +22,12 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
+import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
+import com.amazonaws.services.lambda.runtime.serialization.events.LambdaEventSerializers;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.JsonNodeType;
+import com.fasterxml.jackson.databind.node.NullNode;
 import com.networknt.schema.JsonSchema;
 import com.networknt.schema.ValidationMessage;
 import io.burt.jmespath.Expression;
@@ -65,9 +69,15 @@ public static void validate(Object obj, JsonSchema jsonSchema, String envelope)
         }
         JsonNode subNode;
         try {
-            JsonNode jsonNode = ValidationConfig.get().getObjectMapper().valueToTree(obj);
+            PojoSerializer pojoSerializer = LambdaEventSerializers.serializerFor(obj.getClass(), ClassLoader.getSystemClassLoader());
+            ByteArrayOutputStream out = new ByteArrayOutputStream();
+            pojoSerializer.toJson(obj, out);
+            JsonNode jsonNode = ValidationConfig.get().getObjectMapper().readTree(out.toString("UTF-8"));
             Expression<JsonNode> expression = ValidationConfig.get().getJmesPath().compile(envelope);
             subNode = expression.search(jsonNode);
+            if (subNode == null || subNode instanceof NullNode) {
+                throw new ValidationException("Envelope not found in the object");
+            }
         } catch (Exception e) {
             throw new ValidationException("Cannot find envelope <"+envelope+"> in the object <"+obj+">", e);
         }
diff --git a/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java b/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java
index 0a1f00599..a9d43271b 100644
--- a/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java
+++ b/powertools-validation/src/main/java/software/amazon/lambda/powertools/validation/internal/ValidationAspect.java
@@ -19,6 +19,8 @@
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Pointcut;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.lambda.powertools.validation.Validation;
 import software.amazon.lambda.powertools.validation.ValidationConfig;
 
@@ -36,6 +38,8 @@
  */
 @Aspect
 public class ValidationAspect {
+    private static final Logger LOG = LoggerFactory.getLogger(ValidationAspect.class);
+
     @SuppressWarnings({"EmptyMethod"})
     @Pointcut("@annotation(validation)")
     public void callAt(Validation validation) {
@@ -59,7 +63,9 @@ && placedOnRequestHandler(pjp)) {
                 JsonSchema inboundJsonSchema = getJsonSchema(validation.inboundSchema(), true);
 
                 Object obj = pjp.getArgs()[0];
-                if (obj instanceof APIGatewayProxyRequestEvent) {
+                if (validation.envelope() != null && !validation.envelope().isEmpty()) {
+                    validate(obj, inboundJsonSchema, validation.envelope());
+                } else if (obj instanceof APIGatewayProxyRequestEvent) {
                     APIGatewayProxyRequestEvent event = (APIGatewayProxyRequestEvent) obj;
                     validate(event.getBody(), inboundJsonSchema);
                 } else if (obj instanceof APIGatewayV2HTTPEvent) {
@@ -105,7 +111,7 @@ && placedOnRequestHandler(pjp)) {
                     KinesisAnalyticsStreamsInputPreprocessingEvent event = (KinesisAnalyticsStreamsInputPreprocessingEvent) obj;
                     event.getRecords().forEach(record -> validate(decode(record.getData()), inboundJsonSchema));
                 } else {
-                    validate(obj, inboundJsonSchema, validation.envelope());
+                    LOG.warn("Unhandled event type {}, please use the 'envelope' parameter to specify what to validate", obj.getClass().getName());
                 }
             }
         }
@@ -131,7 +137,7 @@ && placedOnRequestHandler(pjp)) {
                 KinesisAnalyticsInputPreprocessingResponse response = (KinesisAnalyticsInputPreprocessingResponse) result;
                 response.getRecords().forEach(record -> validate(decode(record.getData()), outboundJsonSchema));
             } else {
-                validate(result, outboundJsonSchema, validation.envelope());
+                LOG.warn("Unhandled response type {}, please use the 'envelope' parameter to specify what to validate", result.getClass().getName());
             }
         }
 
diff --git a/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithCustomEnvelopeHandler.java b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithCustomEnvelopeHandler.java
new file mode 100644
index 000000000..9eb96c0e8
--- /dev/null
+++ b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithCustomEnvelopeHandler.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright 2020 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.validation.handlers;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import software.amazon.lambda.powertools.validation.Validation;
+
+public class SQSWithCustomEnvelopeHandler implements RequestHandler<SQSEvent, String> {
+
+    @Override
+    @Validation(inboundSchema = "classpath:/schema_v7.json", envelope = "Records[*].powertools_json(body).powertools_json(Message)")
+    public String handleRequest(SQSEvent input, Context context) {
+        return "OK";
+    }
+}
diff --git a/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithWrongEnvelopeHandler.java b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithWrongEnvelopeHandler.java
new file mode 100644
index 000000000..c1859f29a
--- /dev/null
+++ b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/handlers/SQSWithWrongEnvelopeHandler.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2020 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.validation.handlers;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import software.amazon.lambda.powertools.validation.Validation;
+
+public class SQSWithWrongEnvelopeHandler implements RequestHandler<SQSEvent, String> {
+
+    @Override
+    // real event contains Records with big R (https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html)
+    @Validation(inboundSchema = "classpath:/schema_v7.json", envelope = "records[*].powertools_json(body).powertools_json(Message)")
+    public String handleRequest(SQSEvent input, Context context) {
+        return "OK";
+    }
+}
diff --git a/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/internal/ValidationAspectTest.java b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/internal/ValidationAspectTest.java
index c8741f1e3..2c66885d6 100644
--- a/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/internal/ValidationAspectTest.java
+++ b/powertools-validation/src/test/java/software/amazon/lambda/powertools/validation/internal/ValidationAspectTest.java
@@ -101,6 +101,24 @@ public void validate_SQS() {
         assertThat(handler.handleRequest(event, context)).isEqualTo("OK");
     }
 
+    @Test
+    public void validate_SQS_CustomEnvelopeTakePrecedence() {
+        PojoSerializer<SQSEvent> pojoSerializer = LambdaEventSerializers.serializerFor(SQSEvent.class, ClassLoader.getSystemClassLoader());
+        SQSEvent event = pojoSerializer.fromJson(this.getClass().getResourceAsStream("/sqs_message.json"));
+
+        SQSWithCustomEnvelopeHandler handler = new SQSWithCustomEnvelopeHandler();
+        assertThat(handler.handleRequest(event, context)).isEqualTo("OK");
+    }
+
+    @Test
+    public void validate_SQS_WrongEnvelope_shouldThrowValidationException() {
+        PojoSerializer<SQSEvent> pojoSerializer = LambdaEventSerializers.serializerFor(SQSEvent.class, ClassLoader.getSystemClassLoader());
+        SQSEvent event = pojoSerializer.fromJson(this.getClass().getResourceAsStream("/sqs_message.json"));
+
+        SQSWithWrongEnvelopeHandler handler = new SQSWithWrongEnvelopeHandler();
+        assertThatExceptionOfType(ValidationException.class).isThrownBy(() -> handler.handleRequest(event, context));
+    }
+
     @Test
     public void validate_Kinesis() {
         PojoSerializer<KinesisEvent> pojoSerializer = LambdaEventSerializers.serializerFor(KinesisEvent.class, ClassLoader.getSystemClassLoader());
diff --git a/powertools-validation/src/test/resources/sqs_message.json b/powertools-validation/src/test/resources/sqs_message.json
new file mode 100644
index 000000000..068279b74
--- /dev/null
+++ b/powertools-validation/src/test/resources/sqs_message.json
@@ -0,0 +1,22 @@
+{
+  "Records": [
+    {
+      "messageId": "d9144555-9a4f-4ec3-99a0-fc4e625a8db2",
+      "receiptHandle": "7kam5bfzbDsjtcjElvhSbxeLJbeey3A==",
+      "body": "{\n \"Message\": \"{\\n  \\\"id\\\": 43242,\\n  \\\"name\\\": \\\"FooBar XY\\\",\\n  \\\"price\\\": 258\\n}\"}",
+      "attributes": {
+        "ApproximateReceiveCount": "1",
+        "SentTimestamp": "1601975709495",
+        "SenderId": "AROAIFU457DVZ5L2J53F2",
+        "ApproximateFirstReceiveTimestamp": "1601975709499"
+      },
+      "messageAttributes": {
+
+      },
+      "md5OfBody": "0f96e88a291edb4429f2f7b9fdc3df96",
+      "eventSource": "aws:sqs",
+      "eventSourceARN": "arn:aws:sqs:eu-central-1:123456789012:TestLambda",
+      "awsRegion": "eu-central-1"
+    }
+  ]
+}
\ No newline at end of file

From b65aec6598daab459ee3750e34eaa32982e4a8ee Mon Sep 17 00:00:00 2001
From: Jerome Van Der Linden <jeromevdl@gmail.com>
Date: Wed, 2 Nov 2022 12:01:21 +0100
Subject: [PATCH 4/7] fix idempotency when only string

---
 .../internal/IdempotencyHandler.java          |  3 +-
 .../handlers/IdempotencyStringFunction.java   | 22 +++++++++++
 .../internal/IdempotencyAspectTest.java       | 37 +++++++++++++++++--
 3 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyStringFunction.java

diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
index c1eab3ab5..c7e910bce 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyHandler.java
@@ -26,7 +26,6 @@
 import software.amazon.lambda.powertools.utilities.JsonConfig;
 
 import java.time.Instant;
-import java.util.Optional;
 import java.util.OptionalInt;
 
 import static software.amazon.lambda.powertools.idempotency.persistence.DataRecord.Status.EXPIRED;
@@ -151,6 +150,8 @@ private Object handleForStatus(DataRecord record) {
         Class<?> returnType = ((MethodSignature) pjp.getSignature()).getReturnType();
         try {
             LOG.debug("Response for key '{}' retrieved from idempotency store, skipping the function", record.getIdempotencyKey());
+            if (returnType.equals(String.class))
+                return record.getResponseData();
             return JsonConfig.get().getObjectMapper().reader().readValue(record.getResponseData(), returnType);
         } catch (Exception e) {
             throw new IdempotencyPersistenceLayerException("Unable to get function response as " + returnType.getSimpleName(), e);
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyStringFunction.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyStringFunction.java
new file mode 100644
index 000000000..e6d719598
--- /dev/null
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/handlers/IdempotencyStringFunction.java
@@ -0,0 +1,22 @@
+package software.amazon.lambda.powertools.idempotency.handlers;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import software.amazon.lambda.powertools.idempotency.Idempotent;
+import software.amazon.lambda.powertools.idempotency.model.Product;
+
+public class IdempotencyStringFunction implements RequestHandler<Product, String> {
+
+    private boolean handlerCalled = false;
+
+    public boolean handlerCalled() {
+        return handlerCalled;
+    }
+
+    @Override
+    @Idempotent
+    public String handleRequest(Product input, Context context) {
+        handlerCalled = true;
+        return input.getName();
+    }
+}
diff --git a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
index 21dc2f698..56687e338 100644
--- a/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
+++ b/powertools-idempotency/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java
@@ -16,7 +16,6 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
-import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junitpioneer.jupiter.SetEnvironmentVariable;
@@ -123,6 +122,36 @@ public void secondCall_notExpired_shouldGetFromStore() throws JsonProcessingExce
         assertThat(function.handlerCalled()).isFalse();
     }
 
+    @Test
+    public void secondCall_notExpired_shouldGetStringFromStore() {
+        // GIVEN
+        Idempotency.config()
+                .withPersistenceStore(store)
+                .withConfig(IdempotencyConfig.builder()
+                        .withEventKeyJMESPath("id")
+                        .build()
+                ).configure();
+
+        doThrow(IdempotencyItemAlreadyExistsException.class).when(store).saveInProgress(any(), any(), any());
+
+        Product p = new Product(42, "fake product", 12);
+        DataRecord record = new DataRecord(
+                "42",
+                DataRecord.Status.COMPLETED,
+                Instant.now().plus(356, SECONDS).getEpochSecond(),
+                p.getName(),
+                null);
+        doReturn(record).when(store).getRecord(any(), any());
+
+        // WHEN
+        IdempotencyStringFunction function = new IdempotencyStringFunction();
+        String name = function.handleRequest(p, context);
+
+        // THEN
+        assertThat(name).isEqualTo(p.getName());
+        assertThat(function.handlerCalled()).isFalse();
+    }
+
     @Test
     public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressException() throws JsonProcessingException {
         // GIVEN
@@ -137,13 +166,14 @@ public void secondCall_inProgress_shouldThrowIdempotencyAlreadyInProgressExcepti
 
         Product p = new Product(42, "fake product", 12);
         Basket b = new Basket(p);
+        OptionalLong timestampInFuture = OptionalLong.of(Instant.now().toEpochMilli() + 1000); // timeout not expired (in 1sec)
         DataRecord record = new DataRecord(
                 "42",
                 DataRecord.Status.INPROGRESS,
                 Instant.now().plus(356, SECONDS).getEpochSecond(),
                 JsonConfig.get().getObjectMapper().writer().writeValueAsString(b),
                 null,
-                OptionalLong.of(Instant.now().toEpochMilli() + 1000));
+                timestampInFuture);
         doReturn(record).when(store).getRecord(any(), any());
 
         // THEN
@@ -165,13 +195,14 @@ public void secondCall_inProgress_lambdaTimeout_timeoutExpired_shouldThrowIncons
 
         Product p = new Product(42, "fake product", 12);
         Basket b = new Basket(p);
+        OptionalLong timestampInThePast = OptionalLong.of(Instant.now().toEpochMilli() - 100); // timeout expired 100ms ago
         DataRecord record = new DataRecord(
                 "42",
                 DataRecord.Status.INPROGRESS,
                 Instant.now().plus(356, SECONDS).getEpochSecond(),
                 JsonConfig.get().getObjectMapper().writer().writeValueAsString(b),
                 null,
-                OptionalLong.of(Instant.now().toEpochMilli() - 100));
+                timestampInThePast);
         doReturn(record).when(store).getRecord(any(), any());
 
         // THEN

From 94b5e17e704386423292337c17aaf14ed8c6b140 Mon Sep 17 00:00:00 2001
From: Jerome Van Der Linden <jeromevdl@gmail.com>
Date: Wed, 2 Nov 2022 12:01:40 +0100
Subject: [PATCH 5/7] fix review comments

---
 .../lambda/powertools/idempotency/IdempotencyConfig.java   | 3 +--
 .../idempotency/persistence/BasePersistenceStore.java      | 7 ++++++-
 .../idempotency/persistence/DynamoDBPersistenceStore.java  | 5 ++++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
index aeadfe033..28f20ffa9 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java
@@ -16,7 +16,6 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import software.amazon.lambda.powertools.idempotency.internal.cache.LRUCache;
 
-import java.security.MessageDigest;
 import java.time.Duration;
 
 /**
@@ -214,7 +213,7 @@ public Builder withThrowOnNoIdempotencyKey() {
         /**
          * Function to use for calculating hashes, by default MD5.
          *
-         * @param hashFunction Can be any algorithm supported by {@link MessageDigest}, most commons are<ul>
+         * @param hashFunction Can be any algorithm supported by {@link java.security.MessageDigest}, most commons are<ul>
          *                     <li>MD5</li>
          *                     <li>SHA-1</li>
          *                     <li>SHA-256</li></ul>
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
index fe392262b..96981ddce 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
@@ -106,7 +106,12 @@ public void configure(IdempotencyConfig config, String functionName) {
     public void saveSuccess(JsonNode data, Object result, Instant now) {
         ObjectWriter writer = JsonConfig.get().getObjectMapper().writer();
         try {
-            String responseJson = writer.writeValueAsString(result);
+            String responseJson;
+            if (result instanceof String) {
+                responseJson = (String) result;
+            } else {
+                responseJson = writer.writeValueAsString(result);
+            }
             DataRecord record = new DataRecord(
                     getHashedIdempotencyKey(data),
                     DataRecord.Status.COMPLETED,
diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
index 244a1420a..e48cb78e2 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/DynamoDBPersistenceStore.java
@@ -27,7 +27,10 @@
 import software.amazon.lambda.powertools.idempotency.exceptions.IdempotencyItemNotFoundException;
 
 import java.time.Instant;
-import java.util.*;
+import java.util.AbstractMap;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.OptionalLong;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 

From c3483abfaf2d6b560857bb9a31458fc08c35538a Mon Sep 17 00:00:00 2001
From: Jerome Van Der Linden <jeromevdl@gmail.com>
Date: Thu, 24 Nov 2022 15:46:23 +0100
Subject: [PATCH 6/7] fix review comments

---
 .../idempotency/persistence/BasePersistenceStore.java       | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
index 96981ddce..79db29f05 100644
--- a/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
+++ b/powertools-idempotency/src/main/java/software/amazon/lambda/powertools/idempotency/persistence/BasePersistenceStore.java
@@ -35,8 +35,12 @@
 import java.security.NoSuchAlgorithmException;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.*;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.OptionalInt;
 import java.util.stream.Stream;
+import java.util.Spliterators;
+import java.util.Spliterator;
 import java.util.stream.StreamSupport;
 
 /**

From 9a067283b1d647013ab0becbe9cb324771c24e0c Mon Sep 17 00:00:00 2001
From: Jerome Van Der Linden <jeromevdl@gmail.com>
Date: Fri, 25 Nov 2022 08:55:20 +0100
Subject: [PATCH 7/7] update documentation

---
 .gitignore                    |   2 +
 docs/utilities/idempotency.md | 126 +++++++++++++++++++++++++++++++++-
 mkdocs.yml                    |   5 ++
 3 files changed, 131 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 12a60ce4d..04e85211d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -82,6 +82,8 @@ Desktop.ini
 ######################
 /bin/
 /deploy/
+/dist/
+/site/
 
 ######################
 # Logs
diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
index 56b9acb78..17391218d 100644
--- a/docs/utilities/idempotency.md
+++ b/docs/utilities/idempotency.md
@@ -355,19 +355,141 @@ Imagine the function executes successfully, but the client never receives the re
 
 This sequence diagram shows an example flow of what happens in the payment scenario:
 
-![Idempotent sequence](../media/idempotent_sequence.png)
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Persistence Layer
+    alt initial request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+        activate Persistence Layer
+        Note right of Persistence Layer: Locked to prevent concurrent<br/>invocations with <br/> the same payload.
+        Lambda-->>Lambda: Call handler (event)
+        Lambda->>Persistence Layer: Update record with result
+        deactivate Persistence Layer
+        Persistence Layer-->>Persistence Layer: Update record with result
+        Lambda-->>Client: Response sent to client
+    else retried request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+        Persistence Layer-->>Lambda: Already exists in persistence layer. Return result
+        Lambda-->>Client: Response sent to client
+    end
+```
+<i>Idempotent sequence</i>
+</center>
 
 The client was successful in receiving the result after the retry. Since the Lambda handler was only executed once, our customer hasn't been charged twice.
 
 !!! note
     Bear in mind that the entire Lambda handler is treated as a single idempotent operation. If your Lambda handler can cause multiple side effects, consider splitting it into separate functions.
 
+#### Lambda timeouts
+
+This is automatically done when you annotate your Lambda handler with [@Idempotent annotation](#idempotent-annotation).
+
+To prevent against extended failed retries when a [Lambda function times out](https://aws.amazon.com/premiumsupport/knowledge-center/lambda-verify-invocation-timeouts/), Powertools calculates and includes the remaining invocation available time as part of the idempotency record.
+
+!!! example
+    If a second invocation happens **after** this timestamp, and the record is marked as `INPROGRESS`, we will execute the invocation again as if it was in the `EXPIRED` state.
+    This means that if an invocation expired during execution, it will be quickly executed again on the next retry.
+
+!!! important
+    If you are using the [@Idempotent annotation on another method](#idempotent-annotation-on-another-method) to guard isolated parts of your code, you must use `registerLambdaContext` method available in the `Idempotency` object to benefit from this protection.
+
+    Here is an example on how you register the Lambda context in your handler:
+    
+    ```java hl_lines="13-19" title="Registering the Lambda context"
+    public class PaymentHandler implements RequestHandler<SQSEvent, List<String>> {
+    
+        public PaymentHandler() {
+            Idempotency.config()
+                    .withPersistenceStore(
+                            DynamoDBPersistenceStore.builder()
+                                    .withTableName(System.getenv("IDEMPOTENCY_TABLE"))
+                                    .build())
+                    .configure();
+        }
+        
+        @Override
+        public List<String> handleRequest(SQSEvent sqsEvent, Context context) {
+            Idempotency.registerLambdaContext(context);
+            return sqsEvent.getRecords().stream().map(record -> process(record.getMessageId(), record.getBody())).collect(Collectors.toList());
+        }
+    
+        @Idempotent
+        private String process(String messageId, @IdempotencyKey String messageBody) {
+            logger.info("Processing messageId: {}", messageId);
+            PaymentRequest request = extractDataFrom(messageBody).as(PaymentRequest.class);
+            return paymentService.process(request);
+        }
+    
+    }
+    ```
+
+#### Lambda timeout sequence diagram
+
+This sequence diagram shows an example flow of what happens if a Lambda function times out:
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Persistence Layer
+    alt initial request
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+        activate Persistence Layer
+        Note right of Persistence Layer: Locked to prevent concurrent<br/>invocations with <br/> the same payload.
+        Note over Lambda: Time out
+        Lambda--xLambda: Call handler (event)
+        Lambda-->>Client: Return error response
+        deactivate Persistence Layer
+    else concurrent request before timeout
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+        Persistence Layer-->>Lambda: Request already INPROGRESS
+        Lambda--xClient: Return IdempotencyAlreadyInProgressError
+    else retry after Lambda timeout
+        Client->>Lambda: Invoke (event)
+        Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+        activate Persistence Layer
+        Note right of Persistence Layer: Locked to prevent concurrent<br/>invocations with <br/> the same payload.
+        Lambda-->>Lambda: Call handler (event)
+        Lambda->>Persistence Layer: Update record with result
+        deactivate Persistence Layer
+        Persistence Layer-->>Persistence Layer: Update record with result
+        Lambda-->>Client: Response sent to client
+    end
+```
+<i>Idempotent sequence for Lambda timeouts</i>
+</center>
+
 ### Handling exceptions
 
 If you are using the `@Idempotent` annotation on your Lambda handler or any other method, any unhandled exceptions that are thrown during the code execution will cause **the record in the persistence layer to be deleted**.
 This means that new invocations will execute your code again despite having the same payload. If you don't want the record to be deleted, you need to catch exceptions within the idempotent function and return a successful response.
 
-![Idempotent sequence exception](../media/idempotent_sequence_exception.png)
+<center>
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Lambda
+    participant Persistence Layer
+    Client->>Lambda: Invoke (event)
+    Lambda->>Persistence Layer: Get or set (id=event.search(payload))
+    activate Persistence Layer
+    Note right of Persistence Layer: Locked during this time. Prevents multiple<br/>Lambda invocations with the same<br/>payload running concurrently.
+    Lambda--xLambda: Call handler (event).<br/>Raises exception
+    Lambda->>Persistence Layer: Delete record (id=event.search(payload))
+    deactivate Persistence Layer
+    Lambda-->>Client: Return error response
+```
+<i>Idempotent sequence exception</i>
+</center>
 
 If an Exception is raised _outside_ the scope of a decorated method and after your method has been called, the persistent record will not be affected. In this case, idempotency will be maintained for your decorated function. Example:
 
diff --git a/mkdocs.yml b/mkdocs.yml
index 574b8b30a..58e7fe757 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -56,6 +56,11 @@ markdown_extensions:
   - pymdownx.snippets:
       base_path: '.'
       check_paths: true
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format
   - meta
   - toc:
       permalink: true