Add thread safety to execution operations and unit tests

- Add thread-safe execution operations with proper locking
- Add comprehensive unit tests for execution.py
- Add concurrent execution tests
- Add memory store concurrent tests

Author: Rares Polenciuc

src/aws_durable_execution_sdk_python_testing/execution.py

@@ -3,6 +3,7 @@
 import json
 from dataclasses import replace
 from datetime import UTC, datetime
+from aws_durable_execution_sdk_python.threading import OrderedCounter, OrderedLock
 from typing import Any
 from uuid import uuid4
 
@@ -46,11 +47,24 @@ def __init__(
         self.updates: list[OperationUpdate] = []
         self.used_tokens: set[str] = set()
         # TODO: this will need to persist/rehydrate depending on inmemory vs sqllite store
-        self.token_sequence: int = 0
+
+        self._token_sequence: int = 0
+        self._state_lock: OrderedLock = OrderedLock()
         self.is_complete: bool = False
         self.result: DurableExecutionInvocationOutput | None = None
         self.consecutive_failed_invocation_attempts: int = 0
 
+    @property
+    def token_sequence(self) -> int:
+        """Get current token sequence value."""
+        return self._token_sequence
+    
+    @token_sequence.setter
+    def token_sequence(self, value: int) -> None:
+        """Set token sequence value."""
+        with self._state_lock:
+            self._token_sequence = value
+
     @staticmethod
     def new(input: StartDurableExecutionInput) -> Execution:  # noqa: A002
         # make a nicer arn
@@ -68,7 +82,7 @@ def to_dict(self) -> dict[str, Any]:
             "Operations": [op.to_dict() for op in self.operations],
             "Updates": [update.to_dict() for update in self.updates],
             "UsedTokens": list(self.used_tokens),
-            "TokenSequence": self.token_sequence,
+            "TokenSequence": self._token_sequence,
             "IsComplete": self.is_complete,
             "Result": self.result.to_dict() if self.result else None,
             "ConsecutiveFailedInvocationAttempts": self.consecutive_failed_invocation_attempts,
@@ -95,7 +109,7 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
             OperationUpdate.from_dict(update_data) for update_data in data["Updates"]
         ]
         execution.used_tokens = set(data["UsedTokens"])
-        execution.token_sequence = data["TokenSequence"]
+        execution._token_sequence = data["TokenSequence"]
         execution.is_complete = data["IsComplete"]
         execution.result = (
             DurableExecutionInvocationOutput.from_dict(data["Result"])
@@ -109,23 +123,23 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
         return execution
 
     def start(self) -> None:
-        # not thread safe, prob should be
         if self.start_input.invocation_id is None:
             msg: str = "invocation_id is required"
             raise InvalidParameterValueException(msg)
-        self.operations.append(
-            Operation(
-                operation_id=self.start_input.invocation_id,
-                parent_id=None,
-                name=self.start_input.execution_name,
-                start_timestamp=datetime.now(UTC),
-                operation_type=OperationType.EXECUTION,
-                status=OperationStatus.STARTED,
-                execution_details=ExecutionDetails(
-                    input_payload=json.dumps(self.start_input.input)
-                ),
+        with self._state_lock:
+            self.operations.append(
+                Operation(
+                    operation_id=self.start_input.invocation_id,
+                    parent_id=None,
+                    name=self.start_input.execution_name,
+                    start_timestamp=datetime.now(UTC),
+                    operation_type=OperationType.EXECUTION,
+                    status=OperationStatus.STARTED,
+                    execution_details=ExecutionDetails(
+                        input_payload=json.dumps(self.start_input.input)
+                    ),
+                )
             )
-        )
 
     def get_operation_execution_started(self) -> Operation:
         if not self.operations:
@@ -137,15 +151,15 @@ def get_operation_execution_started(self) -> Operation:
 
     def get_new_checkpoint_token(self) -> str:
         """Generate a new checkpoint token with incremented sequence"""
-        # TODO: not thread safe and it should be
-        self.token_sequence += 1
-        new_token_sequence = self.token_sequence
-        token = CheckpointToken(
-            execution_arn=self.durable_execution_arn, token_sequence=new_token_sequence
-        )
-        token_str = token.to_str()
-        self.used_tokens.add(token_str)
-        return token_str
+        with self._state_lock:
+            self._token_sequence += 1
+            new_token_sequence = self._token_sequence
+            token = CheckpointToken(
+                execution_arn=self.durable_execution_arn, token_sequence=new_token_sequence
+            )
+            token_str = token.to_str()
+            self.used_tokens.add(token_str)
+            return token_str
 
     def get_navigable_operations(self) -> list[Operation]:
         """Get list of operations, but exclude child operations where the parent has already completed."""
@@ -205,17 +219,16 @@ def complete_wait(self, operation_id: str) -> Operation:
             )
             raise IllegalStateException(msg_not_wait)
 
-        # TODO: make thread-safe. Increment sequence
-        self.token_sequence += 1
-
-        # Build and assign updated operation
-        self.operations[index] = replace(
-            operation,
-            status=OperationStatus.SUCCEEDED,
-            end_timestamp=datetime.now(UTC),
-        )
-
-        return self.operations[index]
+        # Thread-safe increment sequence and operation update
+        with self._state_lock:
+            self._token_sequence += 1
+            # Build and assign updated operation
+            self.operations[index] = replace(
+                operation,
+                status=OperationStatus.SUCCEEDED,
+                end_timestamp=datetime.now(UTC),
+            )
+            return self.operations[index]
 
     def complete_retry(self, operation_id: str) -> Operation:
         """Complete STEP retry when timer fires."""
@@ -231,21 +244,21 @@ def complete_retry(self, operation_id: str) -> Operation:
             )
             raise IllegalStateException(msg_not_step)
 
-        # TODO: make thread-safe. Increment sequence
-        self.token_sequence += 1
-
-        # Build updated step_details with cleared next_attempt_timestamp
-        new_step_details = None
-        if operation.step_details:
-            new_step_details = replace(
-                operation.step_details, next_attempt_timestamp=None
+        # Thread-safe increment sequence and operation update
+        with self._state_lock:
+            self._token_sequence += 1
+            # Build updated step_details with cleared next_attempt_timestamp
+            new_step_details = None
+            if operation.step_details:
+                new_step_details = replace(
+                    operation.step_details, next_attempt_timestamp=None
+                )
+
+            # Build updated operation
+            updated_operation = replace(
+                operation, status=OperationStatus.READY, step_details=new_step_details
             )
 
-        # Build updated operation
-        updated_operation = replace(
-            operation, status=OperationStatus.READY, step_details=new_step_details
-        )
-
-        # Assign
-        self.operations[index] = updated_operation
-        return updated_operation
+            # Assign
+            self.operations[index] = updated_operation
+            return updated_operation

src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-import time
 from typing import TYPE_CHECKING, Any, Protocol
 
 import boto3  # type: ignore
@@ -11,11 +10,11 @@
     DurableExecutionInvocationOutput,
     InitialExecutionState,
 )
-from aws_durable_execution_sdk_python.lambda_context import LambdaContext
 
 from aws_durable_execution_sdk_python_testing.exceptions import (
     DurableFunctionsTestError,
 )
+from aws_durable_execution_sdk_python_testing.model import LambdaContext
 
 
 if TYPE_CHECKING:
@@ -46,12 +45,9 @@ def create_test_lambda_context() -> LambdaContext:
     }
 
     return LambdaContext(
-        invoke_id="test-invoke-12345",
+        aws_request_id="test-invoke-12345",
         client_context=client_context_dict,
-        cognito_identity=cognito_identity_dict,
-        epoch_deadline_time_in_ms=int(
-            (time.time() + 900) * 1000
-        ),  # 15 minutes from now
+        identity=cognito_identity_dict,
         invoked_function_arn="arn:aws:lambda:us-west-2:123456789012:function:test-function",
         tenant_id="test-tenant-789",
     )

src/aws_durable_execution_sdk_python_testing/model.py

@@ -25,6 +25,48 @@
 )
 
 
+@dataclass(frozen=True)
+class LambdaContext:
+    """Lambda context for testing."""
+
+    aws_request_id: str
+    log_group_name: str | None = None
+    log_stream_name: str | None = None
+    function_name: str | None = None
+    memory_limit_in_mb: str | None = None
+    function_version: str | None = None
+    invoked_function_arn: str | None = None
+    tenant_id: str | None = None
+    client_context: dict | None = None
+    identity: dict | None = None
+
+    def get_remaining_time_in_millis(self) -> int:
+        return 900000  # 15 minutes default
+
+    def log(self, msg) -> None:
+        pass  # No-op for testing
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        required_fields = ["aws_request_id"]
+        for field in required_fields:
+            if field not in data:
+                msg: str = f"Missing required field: {field}"
+                raise InvalidParameterValueException(msg)
+        return cls(
+            aws_request_id=data["aws_request_id"],
+            log_group_name=data.get("log_group_name"),
+            log_stream_name=data.get("log_stream_name"),
+            function_name=data.get("function_name"),
+            memory_limit_in_mb=data.get("memory_limit_in_mb"),
+            function_version=data.get("function_version"),
+            invoked_function_arn=data.get("invoked_function_arn"),
+            tenant_id=data.get("tenant_id"),
+            client_context=data.get("client_context"),
+            identity=data.get("identity"),
+        )
+
+
 # Web API specific models (not in Smithy but needed for web interface)
 @dataclass(frozen=True)
 class StartDurableExecutionInput:

src/aws_durable_execution_sdk_python_testing/stores/memory.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+from aws_durable_execution_sdk_python.threading import OrderedLock
 
 if TYPE_CHECKING:
     from aws_durable_execution_sdk_python_testing.execution import Execution
@@ -14,15 +15,20 @@ class InMemoryExecutionStore:
 
     def __init__(self) -> None:
         self._store: dict[str, Execution] = {}
+        self._lock: OrderedLock = OrderedLock()
 
     def save(self, execution: Execution) -> None:
-        self._store[execution.durable_execution_arn] = execution
+        with self._lock:
+            self._store[execution.durable_execution_arn] = execution
 
     def load(self, execution_arn: str) -> Execution:
-        return self._store[execution_arn]
+        with self._lock:
+            return self._store[execution_arn]
 
     def update(self, execution: Execution) -> None:
-        self._store[execution.durable_execution_arn] = execution
+        with self._lock:
+            self._store[execution.durable_execution_arn] = execution
 
     def list_all(self) -> list[Execution]:
-        return list(self._store.values())
+        with self._lock:
+            return list(self._store.values())