feat(idempotency): handle lambda timeout scenarios for INPROGRESS records

aws_lambda_powertools/utilities/idempotency/base.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 from copy import deepcopy
 from typing import Any, Callable, Dict, Optional, Tuple
@@ -73,6 +74,7 @@ def __init__(
         self.data = deepcopy(_prepare_data(function_payload))
         self.fn_args = function_args
         self.fn_kwargs = function_kwargs
+        self.config = config
 
         persistence_store.configure(config, self.function.__name__)
         self.persistence_store = persistence_store
@@ -101,7 +103,9 @@ def _process_idempotency(self):
         try:
             # We call save_inprogress first as an optimization for the most common case where no idempotent record
             # already exists. If it succeeds, there's no need to call get_record.
-            self.persistence_store.save_inprogress(data=self.data)
+            self.persistence_store.save_inprogress(
+                data=self.data, remaining_time_in_millis=self._get_remaining_time_in_millis()
+            )
         except IdempotencyKeyError:
             raise
         except IdempotencyItemAlreadyExistsError:
@@ -113,6 +117,24 @@ def _process_idempotency(self):
 
         return self._get_function_response()
 
+    def _get_remaining_time_in_millis(self) -> Optional[int]:
+        """
+        Tries to determine the remaining time available for the current lambda invocation.
+
+        Currently, it only works if the idempotent handler decorator is used, since we need to acess the lambda context.
+        However, this could be improved if we start storing the lambda context globally during the invocation.
+
+        Returns
+        -------
+        Optional[int]
+            Remaining time in millis, or None if the remaining time cannot be determined.
+        """
+
+        if self.config.lambda_context is not None:
+            return self.config.lambda_context.get_remaining_time_in_millis()
+
+        return None
+
     def _get_idempotency_record(self) -> DataRecord:
         """
         Retrieve the idempotency record from the persistence layer.
@@ -167,6 +189,13 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Dict[Any, Any]
             raise IdempotencyInconsistentStateError("save_inprogress and get_record return inconsistent results.")
 
         if data_record.status == STATUS_CONSTANTS["INPROGRESS"]:
+            if data_record.in_progress_expiry_timestamp is not None and data_record.in_progress_expiry_timestamp < int(
+                datetime.datetime.now().timestamp()
+            ):
+                raise IdempotencyInconsistentStateError(
+                    "item should have been expired in-progress because it already time-outed."
+                )
+
             raise IdempotencyAlreadyInProgressError(
                 f"Execution already in progress with idempotency key: "
                 f"{self.persistence_store.event_key_jmespath}={data_record.idempotency_key}"

aws_lambda_powertools/utilities/idempotency/config.py

@@ -1,5 +1,7 @@
 from typing import Dict, Optional
 
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
 
 class IdempotencyConfig:
     def __init__(
@@ -41,3 +43,7 @@ def __init__(
         self.use_local_cache = use_local_cache
         self.local_cache_max_items = local_cache_max_items
         self.hash_function = hash_function
+        self.lambda_context: Optional[LambdaContext] = None
+
+    def register_lambda_context(self, lambda_context: LambdaContext):
+        self.lambda_context = lambda_context

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -62,6 +62,8 @@ def idempotent(
         return handler(event, context)
 
     config = config or IdempotencyConfig()
+    config.register_lambda_context(context)
+
     args = event, context
     idempotency_handler = IdempotencyHandler(
         function=handler,

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -8,6 +8,7 @@
 import os
 import warnings
 from abc import ABC, abstractmethod
+from math import ceil
 from types import MappingProxyType
 from typing import Any, Dict, Optional
 
@@ -40,6 +41,7 @@ def __init__(
         idempotency_key,
         status: str = "",
         expiry_timestamp: Optional[int] = None,
+        in_progress_expiry_timestamp: Optional[int] = None,
         response_data: Optional[str] = "",
         payload_hash: Optional[str] = None,
     ) -> None:
@@ -53,6 +55,8 @@ def __init__(
             status of the idempotent record
         expiry_timestamp: int, optional
             time before the record should expire, in seconds
+        in_progress_expiry_timestamp: int, optional
+            time before the record should expire while in the INPROGRESS state, in seconds
         payload_hash: str, optional
             hashed representation of payload
         response_data: str, optional
@@ -61,6 +65,7 @@ def __init__(
         self.idempotency_key = idempotency_key
         self.payload_hash = payload_hash
         self.expiry_timestamp = expiry_timestamp
+        self.in_progress_expiry_timestamp = in_progress_expiry_timestamp
         self._status = status
         self.response_data = response_data
 
@@ -328,14 +333,16 @@ def save_success(self, data: Dict[str, Any], result: dict) -> None:
 
         self._save_to_cache(data_record=data_record)
 
-    def save_inprogress(self, data: Dict[str, Any]) -> None:
+    def save_inprogress(self, data: Dict[str, Any], remaining_time_in_millis: Optional[int] = None) -> None:
         """
         Save record of function's execution being in progress
 
         Parameters
         ----------
         data: Dict[str, Any]
             Payload
+        remaining_time_in_millis: Optional[int]
+            If expiry of in-progress invocations is enabled, this will contain the remaining time available in millis
         """
         data_record = DataRecord(
             idempotency_key=self._get_hashed_idempotency_key(data=data),
@@ -344,6 +351,19 @@ def save_inprogress(self, data: Dict[str, Any]) -> None:
             payload_hash=self._get_hashed_payload(data=data),
         )
 
+        if remaining_time_in_millis:
+            now = datetime.datetime.now()
+            period = datetime.timedelta(milliseconds=remaining_time_in_millis)
+
+            # It's very important to use math.ceil here. Otherwise, we might return an integer that will be smaller
+            # than the current time in milliseconds, due to rounding. This will create a scenario where the record
+            # looks already expired in the store, but the invocation is still running.
+            timestamp = ceil((now + period).timestamp())
+
+            data_record.in_progress_expiry_timestamp = timestamp
+        else:
+            warnings.warn("Expires in progress is enabled but we couldn't determine the remaining time left")
+
         logger.debug(f"Saving in progress record for idempotency key: {data_record.idempotency_key}")
 
         if self._retrieve_from_cache(idempotency_key=data_record.idempotency_key):

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -12,7 +12,7 @@
     IdempotencyItemAlreadyExistsError,
     IdempotencyItemNotFoundError,
 )
-from aws_lambda_powertools.utilities.idempotency.persistence.base import DataRecord
+from aws_lambda_powertools.utilities.idempotency.persistence.base import STATUS_CONSTANTS, DataRecord
 
 logger = logging.getLogger(__name__)
 
@@ -25,6 +25,7 @@ def __init__(
         static_pk_value: Optional[str] = None,
         sort_key_attr: Optional[str] = None,
         expiry_attr: str = "expiration",
+        in_progress_expiry_attr: str = "in_progress_expiration",
         status_attr: str = "status",
         data_attr: str = "data",
         validation_key_attr: str = "validation",
@@ -47,6 +48,8 @@ def __init__(
             DynamoDB attribute name for the sort key
         expiry_attr: str, optional
             DynamoDB attribute name for expiry timestamp, by default "expiration"
+        in_progress_expiry_attr: str, optional
+            DynamoDB attribute name for in-progress expiry timestamp, by default "in_progress_expiration"
         status_attr: str, optional
             DynamoDB attribute name for status, by default "status"
         data_attr: str, optional
@@ -85,6 +88,7 @@ def __init__(
         self.static_pk_value = static_pk_value
         self.sort_key_attr = sort_key_attr
         self.expiry_attr = expiry_attr
+        self.in_progress_expiry_attr = in_progress_expiry_attr
         self.status_attr = status_attr
         self.data_attr = data_attr
         self.validation_key_attr = validation_key_attr
@@ -133,6 +137,7 @@ def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
             idempotency_key=item[self.key_attr],
             status=item[self.status_attr],
             expiry_timestamp=item[self.expiry_attr],
+            in_progress_expiry_timestamp=item.get(self.in_progress_expiry_attr),
             response_data=item.get(self.data_attr),
             payload_hash=item.get(self.validation_key_attr),
         )
@@ -153,34 +158,51 @@ def _put_record(self, data_record: DataRecord) -> None:
             self.status_attr: data_record.status,
         }
 
+        if data_record.in_progress_expiry_timestamp is not None:
+            item[self.in_progress_expiry_attr] = data_record.in_progress_expiry_timestamp
+
         if self.payload_validation_enabled:
             item[self.validation_key_attr] = data_record.payload_hash
 
         now = datetime.datetime.now()
         try:
             logger.debug(f"Putting record for idempotency key: {data_record.idempotency_key}")
+
             self.table.put_item(
                 Item=item,
-                ConditionExpression="attribute_not_exists(#id) OR #now < :now",
-                ExpressionAttributeNames={"#id": self.key_attr, "#now": self.expiry_attr},
-                ExpressionAttributeValues={":now": int(now.timestamp())},
+                ConditionExpression=(
+                    "attribute_not_exists(#id) OR #now < :now OR "
+                    "(attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)"
+                ),
+                ExpressionAttributeNames={
+                    "#id": self.key_attr,
+                    "#now": self.expiry_attr,
+                    "#in_progress_expiry": self.in_progress_expiry_attr,
+                    "#status": self.status_attr,
+                },
+                ExpressionAttributeValues={":now": int(now.timestamp()), ":inprogress": STATUS_CONSTANTS["INPROGRESS"]},
             )
         except self.table.meta.client.exceptions.ConditionalCheckFailedException:
             logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")
             raise IdempotencyItemAlreadyExistsError
 
     def _update_record(self, data_record: DataRecord):
         logger.debug(f"Updating record for idempotency key: {data_record.idempotency_key}")
-        update_expression = "SET #response_data = :response_data, #expiry = :expiry, #status = :status"
+        update_expression = (
+            "SET #response_data = :response_data, #expiry = :expiry, "
+            "#status = :status, #in_progress_expiry = :in_progress_expiry"
+        )
         expression_attr_values = {
             ":expiry": data_record.expiry_timestamp,
             ":response_data": data_record.response_data,
             ":status": data_record.status,
+            ":in_progress_expiry": data_record.in_progress_expiry_timestamp,
         }
         expression_attr_names = {
             "#response_data": self.data_attr,
             "#expiry": self.expiry_attr,
             "#status": self.status_attr,
+            "#in_progress_expiry": self.in_progress_expiry_attr,
         }
 
         if self.payload_validation_enabled: