aws-durable-execution-sdk-python 1.0.0__py3-none-any.whl

aws_durable_execution_sdk_python/suspend.py

@@ -0,0 +1,84 @@
+"""
+The model expects >= 1 seconds when receiving an OperationUpdate via api calls
+    (see StepOptions, WaitOptions, etc)
+We don't force a minimum delay_seconds OR time to timestamp here because:
+
+1. suspension can be reached from multiple handlers,
+2. suspension can be reached from multiple "contexts", e.g. top level wait, or a child retry
+3. we use TimedSuspendExecution as an optimization mechanism within concurrent executions (map/parallel)
+    to know when to retry without being told so.
+
+As such, it is up to the caller to ensure consistency wrt what dataplane sees and what happens within a
+function.
+
+Behaviour:
+- When we suspend without a target delay / timestamp THEN we suspend indefinitely.
+- When `delay_seconds` or timestamp exist, THEN we suspend for `max(delay_seconds, 0)` or suspend for "now".
+- When suspension happens within a child branch and the branch suspends for `0` seconds, THEN execution will
+resume immediately as it will be inserted at the top of the queue
+- When suspension happens within a child branch and the branch suspends for > 0 seconds, THEN execution will
+resume as soon as `delay_seconds` have passed
+- When suspension happens at the top level, then the Lambda will terminate and Dataplane is responsible
+    for resuming
+
+"""
+
+import datetime
+from typing import NoReturn
+
+from aws_durable_execution_sdk_python.exceptions import (
+    SuspendExecution,
+    TimedSuspendExecution,
+)
+
+
+def suspend_with_optional_resume_timestamp(
+    msg: str, datetime_timestamp: datetime.datetime | None = None
+) -> NoReturn:
+    """Suspend execution with optional timestamp.
+
+    Args:
+        msg: Descriptive message for the suspension
+        timestamp: Timestamp to suspend until, or None for indefinite
+
+    Raises:
+        TimedSuspendExecution: When timestamp is in the future or now()
+        SuspendExecution: When timestamp is None or in the past
+    """
+
+    if datetime_timestamp is None:
+        msg = f"No timestamp provided. Suspending without retry timestamp. Original operation: [{msg}]"
+        raise SuspendExecution(msg)
+
+    if datetime_timestamp < datetime.datetime.now(tz=datetime.UTC):
+        msg = f"Invalid timestamp {datetime_timestamp}, suspending with immediate retry, original operation: [{msg}]"
+        raise TimedSuspendExecution.from_datetime(
+            msg, datetime.datetime.now(tz=datetime.UTC)
+        )
+
+    raise TimedSuspendExecution.from_datetime(msg, datetime_timestamp)
+
+
+def suspend_with_optional_resume_delay(
+    msg: str, delay_seconds: int | None = None
+) -> NoReturn:
+    """Suspend execution with optional delay.
+
+    Args:
+        msg: Descriptive message for the suspension
+        delay_seconds: Duration to suspend in seconds, or None for indefinite
+
+    Raises:
+        TimedSuspendExecution: When delay_seconds when delay_seconds is not None
+        SuspendExecution: When delay_seconds is None
+    """
+
+    if delay_seconds is None:
+        msg = f"No delay_seconds provided, suspending without retry timestamp, original operation: [{msg}]"
+        raise SuspendExecution(msg)
+
+    if delay_seconds < 0:
+        msg = f"Invalid delay_seconds {delay_seconds}, suspending with delay 0, original operation: [{msg}]"
+        raise TimedSuspendExecution.from_delay(msg, 0)
+
+    raise TimedSuspendExecution.from_delay(msg, delay_seconds)

aws_durable_execution_sdk_python/threading.py

@@ -0,0 +1,222 @@
+"""Concurrency and locking."""
+
+from __future__ import annotations
+
+from collections import deque
+from threading import Event, Lock
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.exceptions import OrderedLockError
+
+if TYPE_CHECKING:
+    from typing import Self
+
+
+class CompletionEvent:
+    """Threading event that can signal completion or propagate errors.
+
+    This event allows a background thread to wake up a waiting thread either
+    with a successful completion signal or by propagating an exception. When
+    an error is set, the waiting thread will raise that exception.
+
+    This is used for checkpoint operations where the background checkpoint
+    thread needs to signal completion to the user thread, or interrupt it
+    with a critical error.
+
+    Example:
+        >>> event = CompletionEvent()
+        >>> # In background thread:
+        >>> try:
+        ...     process_checkpoint()
+        ...     event.set()  # Success
+        ... except Exception as e:
+        ...     event.set(BackgroundThreadError(..., e))  # Error
+        >>>
+        >>> # In user thread:
+        >>> event.wait()  # Raises BackgroundThreadError if set
+    """
+
+    def __init__(self) -> None:
+        """Initialize completion event."""
+        self._event: Event = Event()
+        self._error: BaseException | None = None
+
+    def set(self, error: BaseException | None = None) -> None:
+        """Signal completion, optionally with an error.
+
+        Args:
+            error: Optional exception to propagate to waiting thread.
+                  If provided, wait() will raise this exception.
+        """
+        # Only set error if none is already set (first error wins)
+        if self._error is None:
+            self._error = error
+        self._event.set()
+
+    def wait(self, timeout: float | None = None) -> bool:
+        """Wait for completion and raise if error occurred.
+
+        Args:
+            timeout: Optional timeout in seconds
+
+        Returns:
+            True if event was set, False if timeout occurred
+
+        Raises:
+            BaseException: If an error was set via set(error=...)
+        """
+        result = self._event.wait(timeout)
+        if self._error is not None:
+            raise self._error
+        return result
+
+    def is_set(self) -> bool:
+        """Return True if the event is set, False otherwise."""
+        return self._event.is_set()
+
+
+class OrderedLock:
+    """Lock that guarantees callers acquire in the invocation order.
+
+    Locks acquire in first-in,first-out (FIFO) order.
+
+    This class is necessary because in a standard Lock the order of pending calls
+    acquiring the lock is not necessarily guaranteed by the thread scheduler.
+
+    For example, assume calls to acquire the lock in order A -> B -> C.
+    A blocks with B and C pending. When A releases, the thread scheduler could favour
+    C rather than B next, which is out of order.
+
+    This OrderedLock instead will guarantee that the order in which callers will
+    acquire the lock is the order of invocation. In the case of example, this means
+    that the order of lock acquire would always be A -> B -> C.
+
+    Once an error occurs in a lock, this instance of the lock is broken and no subsequent lock attempts
+    can succeed, because if any subsequent locks acquire it would violate the order guarantee.
+
+    If a lock fails to acquire, OrderedLock will raise the causing exception to the caller.
+    If there are any other blocked callers waiting in queue, those callers will receive a
+    OrderedLockError, which contains the original causing exception too.
+
+    You can use OrderedLock as a context manager.
+    """
+
+    def __init__(self) -> None:
+        """Initialize ordered lock."""
+        self._lock: Lock = Lock()
+        self._waiters: deque[Event] = deque()
+        self._is_broken: bool = False
+        self._exception: Exception | None = None
+
+    def acquire(self) -> bool:
+        """Acquire lock.
+
+        Returns: True if acquired successfully
+
+        Raises:
+            OrderedLockError: When a preceding caller could not release its lock because it errored.
+        """
+        with self._lock:
+            if self._is_broken:
+                # don't grow queue if already broken
+                msg = "Cannot acquire lock in guaranteed order because a previous lock exited with an exception."
+                raise OrderedLockError(msg, self._exception)
+
+            event = Event()
+            self._waiters.append(event)
+
+            if len(self._waiters) == 1:
+                # first waiter, nothing else in queue so no need to wait
+                event.set()
+
+        # block until it's our turn to proceed
+        event.wait()
+
+        # this is the only thread progressing and holding the lock, so doesn't need to be under lock
+        if self._is_broken:
+            msg = "Cannot acquire lock in guaranteed order because a previous lock exited with an exception."
+            raise OrderedLockError(msg, self._exception)
+
+        return True
+
+    def release(self) -> None:
+        """Release lock. This makes the lock available for the next queued up waiter."""
+        with self._lock:
+            if not self._waiters:
+                msg = "You have to acquire a lock before you can release it."
+                raise OrderedLockError(msg)
+            # remove the current lock from the queue, since it's done
+            self._waiters.popleft()
+            if self._waiters and not self._is_broken:
+                # let the next-in-line waiter proceed
+                self._waiters[0].set()
+
+    def reset(self) -> None:
+        """Reset the lock.
+
+        This assumes all waiters have cleared.
+
+        Raises: OrderedLockError when there still are pending waiters.
+        """
+        with self._lock:
+            if self._waiters:
+                msg = (
+                    "Cannot reset lock because there are callers waiting for the lock."
+                )
+                raise OrderedLockError(msg)
+            self._is_broken = False
+            self._exception = None
+
+    def is_broken(self) -> bool:
+        """Return True if the lock is broken."""
+        with self._lock:
+            return self._is_broken
+
+    # region Context Manager
+    def __enter__(self) -> Self:
+        """Acquire lock."""
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context manager by releasing the current lock."""
+        if exc_type is not None:
+            # can't allow any subsequent locks to succeed, because that would break order guarantee
+            with self._lock:
+                self._is_broken = True
+                self._exception = exc_val
+                # break the queue and let all waiters know
+                for waiter in self._waiters:
+                    waiter.set()
+
+        self.release()
+
+    # endregion Context Manager
+
+
+class OrderedCounter:
+    """Thread-safe counter that guarantees callers get the next increment in the invocation order.
+
+    The counter starts at 0.
+    """
+
+    def __init__(self) -> None:
+        self._lock: OrderedLock = OrderedLock()
+        self._counter: int = 0
+
+    def increment(self) -> int:
+        """Increment the counter by 1."""
+        with self._lock:
+            self._counter += 1
+            return self._counter
+
+    def decrement(self) -> int:
+        """Decrement the counter by 1."""
+        with self._lock:
+            self._counter -= 1
+            return self._counter
+
+    def get_current(self) -> int:
+        """Return the current value of the counter."""
+        with self._lock:
+            return self._counter

aws_durable_execution_sdk_python/types.py

@@ -0,0 +1,180 @@
+"""Types and Protocols. Don't import anything other than config here - the reason it exists is to avoid circular references."""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping, Sequence
+
+    from aws_durable_execution_sdk_python.config import (
+        BatchedInput,
+        CallbackConfig,
+        ChildConfig,
+        Duration,
+        MapConfig,
+        ParallelConfig,
+        StepConfig,
+    )
+
+T = TypeVar("T")
+U = TypeVar("U")
+C_co = TypeVar("C_co", covariant=True)
+C_contra = TypeVar("C_contra", contravariant=True)
+
+
+class LoggerInterface(Protocol):
+    def debug(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None: ...  # pragma: no cover
+
+    def info(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None: ...  # pragma: no cover
+
+    def warning(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None: ...  # pragma: no cover
+
+    def error(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None: ...  # pragma: no cover
+
+    def exception(
+        self, msg: object, *args: object, extra: Mapping[str, object] | None = None
+    ) -> None: ...  # pragma: no cover
+
+
+@dataclass(frozen=True)
+class OperationContext:
+    logger: LoggerInterface
+
+
+@dataclass(frozen=True)
+class StepContext(OperationContext):
+    pass
+
+
+@dataclass(frozen=True)
+class WaitForCallbackContext(OperationContext):
+    """Context provided to waitForCallback submitter functions."""
+
+
+@dataclass(frozen=True)
+class WaitForConditionCheckContext(OperationContext):
+    pass
+
+
+class Callback(Protocol, Generic[C_co]):
+    """Protocol for callback futures."""
+
+    callback_id: str
+
+    @abstractmethod
+    def result(self) -> C_co | None:
+        """Return the result of the future. Will block until result is available."""
+        ...  # pragma: no cover
+
+
+class BatchResult(Protocol, Generic[T]):
+    """Protocol for batch operation results."""
+
+    @abstractmethod
+    def get_results(self) -> list[T]:
+        """Get all successful results."""
+        ...  # pragma: no cover
+
+
+class DurableContext(Protocol):
+    """Protocol defining the interface for durable execution contexts."""
+
+    @abstractmethod
+    def step(
+        self,
+        func: Callable[[StepContext], T],
+        name: str | None = None,
+        config: StepConfig | None = None,
+    ) -> T:
+        """Execute a step durably."""
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def run_in_child_context(
+        self,
+        func: Callable[[DurableContext], T],
+        name: str | None = None,
+        config: ChildConfig | None = None,
+    ) -> T:
+        """Run callable in a child context."""
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def map(
+        self,
+        inputs: Sequence[U],
+        func: Callable[[DurableContext, U | BatchedInput[Any, U], int, Sequence[U]], T],
+        name: str | None = None,
+        config: MapConfig | None = None,
+    ) -> BatchResult[T]:
+        """Apply function durably to each item in inputs."""
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def parallel(
+        self,
+        functions: Sequence[Callable[[DurableContext], T]],
+        name: str | None = None,
+        config: ParallelConfig | None = None,
+    ) -> BatchResult[T]:
+        """Execute callables durably in parallel."""
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def wait(self, duration: Duration, name: str | None = None) -> None:
+        """Wait for a specified amount of time."""
+        ...  # pragma: no cover
+
+    @abstractmethod
+    def create_callback(
+        self, name: str | None = None, config: CallbackConfig | None = None
+    ) -> Callback:
+        """Create a callback."""
+        ...  # pragma: no cover
+
+
+class LambdaContext(Protocol):  # pragma: no cover
+    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: Any | None = None
+    identity: Any | None = None
+
+    def get_remaining_time_in_millis(self) -> int: ...
+    def log(self, msg) -> None: ...
+
+
+# region Summary
+
+"""Summary generators for concurrent operations.
+
+Summary generators create compact JSON representations of large BatchResult objects
+when the serialized result exceeds the 256KB checkpoint size limit. This prevents
+large payloads from being stored in checkpoints while maintaining operation metadata.
+
+When a summary is used, the operation is marked with ReplayChildren=true, causing
+the child context to be re-executed during replay to reconstruct the full result.
+"""
+
+
+class SummaryGenerator(Protocol[C_contra]):
+    def __call__(self, result: C_contra) -> str: ...  # pragma: no cover
+
+
+# endregion Summary

aws_durable_execution_sdk_python/waits.py

@@ -0,0 +1,130 @@
+"""Ready-made wait strategies and wait creators."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generic
+
+from aws_durable_execution_sdk_python.config import Duration, JitterStrategy, T
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python.serdes import SerDes
+
+Numeric = int | float
+
+
+@dataclass
+class WaitDecision:
+    """Decision about whether to wait a step and with what delay."""
+
+    should_wait: bool
+    delay: Duration
+
+    @property
+    def delay_seconds(self) -> int:
+        """Get delay in seconds."""
+        return self.delay.to_seconds()
+
+    @classmethod
+    def wait(cls, delay: Duration) -> WaitDecision:
+        """Create a wait decision."""
+        return cls(should_wait=True, delay=delay)
+
+    @classmethod
+    def no_wait(cls) -> WaitDecision:
+        """Create a no-wait decision."""
+        return cls(should_wait=False, delay=Duration())
+
+
+@dataclass
+class WaitStrategyConfig(Generic[T]):
+    should_continue_polling: Callable[[T], bool]
+    max_attempts: int = 60
+    initial_delay: Duration = field(default_factory=lambda: Duration.from_seconds(5))
+    max_delay: Duration = field(
+        default_factory=lambda: Duration.from_minutes(5)
+    )  # 5 minutes
+    backoff_rate: Numeric = 1.5
+    jitter_strategy: JitterStrategy = field(default=JitterStrategy.FULL)
+    timeout: Duration | None = None  # Not implemented yet
+
+    @property
+    def initial_delay_seconds(self) -> int:
+        """Get initial delay in seconds."""
+        return self.initial_delay.to_seconds()
+
+    @property
+    def max_delay_seconds(self) -> int:
+        """Get max delay in seconds."""
+        return self.max_delay.to_seconds()
+
+    @property
+    def timeout_seconds(self) -> int | None:
+        """Get timeout in seconds."""
+        if self.timeout is None:
+            return None
+        return self.timeout.to_seconds()
+
+
+def create_wait_strategy(
+    config: WaitStrategyConfig[T],
+) -> Callable[[T, int], WaitDecision]:
+    def wait_strategy(result: T, attempts_made: int) -> WaitDecision:
+        # Check if condition is met
+        if not config.should_continue_polling(result):
+            return WaitDecision.no_wait()
+
+        # Check if we've exceeded max attempts
+        if attempts_made >= config.max_attempts:
+            return WaitDecision.no_wait()
+
+        # Calculate delay with exponential backoff
+        base_delay: float = min(
+            config.initial_delay_seconds * (config.backoff_rate ** (attempts_made - 1)),
+            config.max_delay_seconds,
+        )
+
+        # Apply jitter to get final delay
+        delay_with_jitter: float = config.jitter_strategy.apply_jitter(base_delay)
+
+        # Round up and ensure minimum of 1 second
+        final_delay: int = max(1, math.ceil(delay_with_jitter))
+
+        return WaitDecision.wait(Duration(seconds=final_delay))
+
+    return wait_strategy
+
+
+@dataclass(frozen=True)
+class WaitForConditionDecision:
+    """Decision about whether to continue waiting."""
+
+    should_continue: bool
+    delay: Duration
+
+    @property
+    def delay_seconds(self) -> int:
+        """Get delay in seconds."""
+        return self.delay.to_seconds()
+
+    @classmethod
+    def continue_waiting(cls, delay: Duration) -> WaitForConditionDecision:
+        """Create a decision to continue waiting for delay_seconds."""
+        return cls(should_continue=True, delay=delay)
+
+    @classmethod
+    def stop_polling(cls) -> WaitForConditionDecision:
+        """Create a decision to stop polling."""
+        return cls(should_continue=False, delay=Duration())
+
+
+@dataclass(frozen=True)
+class WaitForConditionConfig(Generic[T]):
+    """Configuration for wait_for_condition."""
+
+    wait_strategy: Callable[[T, int], WaitForConditionDecision]
+    initial_state: T
+    serdes: SerDes | None = None