aws-durable-execution-sdk-python-testing 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

aws_durable_execution_sdk_python_testing/worker/checkpoint_tasks.py

@@ -0,0 +1,119 @@
+"""Worker tasks for the checkpoint and get-state operations.
+
+Routing these through a per-execution worker serializes all checkpoint
+and get-state calls for one execution, so concurrent callers can never
+observe or write half-applied state.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    CheckpointOutput,
+    StateOutput,
+)
+
+from aws_durable_execution_sdk_python_testing.worker.status import InvocationState
+from aws_durable_execution_sdk_python_testing.worker.task import (
+    ExecutionTask,
+    TaskOutcome,
+    WorkerContext,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python.lambda_service import OperationUpdate
+
+    from aws_durable_execution_sdk_python_testing.checkpoint.processor import (
+        CheckpointProcessor,
+    )
+
+T = TypeVar("T")
+
+
+def _status_after(worker: WorkerContext) -> InvocationState:
+    """The status to carry forward after a task.
+
+    COMPLETED once the execution is terminal (so the worker tears down);
+    otherwise whatever status the task left on the worker. A store error
+    during the completion check is treated as not-yet-terminal.
+    """
+    try:
+        is_complete: bool = worker.reload().is_complete
+    except Exception:  # noqa: BLE001 — store may be unavailable; keep the worker
+        return worker.status
+    return InvocationState.COMPLETED if is_complete else worker.status
+
+
+class CallableTask(ExecutionTask[T]):
+    """Runs a checkpoint-style callable on the worker lane.
+
+    Lets a caller serialize an existing operation on the execution's
+    lane without restating it as a task; the lane is released once the
+    execution completes.
+    """
+
+    def __init__(self, fn: Callable[[], T]) -> None:
+        self._fn: Callable[[], T] = fn
+
+    def execute(
+        self,
+        status: InvocationState,  # noqa: ARG002
+        worker: WorkerContext,
+    ) -> TaskOutcome[T]:
+        result: T = self._fn()
+        return TaskOutcome(_status_after(worker), result)
+
+
+class CheckpointTask(ExecutionTask[CheckpointOutput]):
+    """Applies a checkpoint for one execution."""
+
+    def __init__(
+        self,
+        processor: CheckpointProcessor,
+        checkpoint_token: str,
+        updates: list[OperationUpdate],
+        client_token: str | None,
+    ) -> None:
+        self._processor: CheckpointProcessor = processor
+        self._checkpoint_token: str = checkpoint_token
+        self._updates: list[OperationUpdate] = updates
+        self._client_token: str | None = client_token
+
+    def execute(
+        self,
+        status: InvocationState,  # noqa: ARG002
+        worker: WorkerContext,
+    ) -> TaskOutcome[CheckpointOutput]:
+        response: CheckpointOutput = self._processor.process_checkpoint(
+            self._checkpoint_token, self._updates, self._client_token
+        )
+        return TaskOutcome(_status_after(worker), response)
+
+
+class GetStateTask(ExecutionTask[StateOutput]):
+    """Reads the current state for one execution."""
+
+    def __init__(
+        self,
+        processor: CheckpointProcessor,
+        checkpoint_token: str,
+        next_marker: str,
+        max_items: int = 1000,
+    ) -> None:
+        self._processor: CheckpointProcessor = processor
+        self._checkpoint_token: str = checkpoint_token
+        self._next_marker: str = next_marker
+        self._max_items: int = max_items
+
+    def execute(
+        self,
+        status: InvocationState,  # noqa: ARG002
+        worker: WorkerContext,
+    ) -> TaskOutcome[StateOutput]:
+        response: StateOutput = self._processor.get_execution_state(
+            self._checkpoint_token, self._next_marker, self._max_items
+        )
+        return TaskOutcome(_status_after(worker), response)

aws_durable_execution_sdk_python_testing/worker/execution_worker.py

@@ -0,0 +1,117 @@
+"""The per-execution serial worker."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from aws_durable_execution_sdk_python_testing.worker.lane import SerialTaskLane
+from aws_durable_execution_sdk_python_testing.worker.status import InvocationState
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+    from aws_durable_execution_sdk_python_testing.scheduler import Scheduler
+    from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
+    from aws_durable_execution_sdk_python_testing.worker.registry import (
+        ExecutionRegistry,
+    )
+    from aws_durable_execution_sdk_python_testing.worker.task import (
+        ExecutionTask,
+        TaskOutcome,
+    )
+
+T = TypeVar("T")
+
+
+class ExecutionWorker:
+    """Runs every operation for one execution, one at a time.
+
+    Operations are submitted as :class:`ExecutionTask` instances and run
+    on a single lane, so they never overlap and the execution's state
+    stays consistent. The worker carries the invocation state returned by
+    each task into the next, and once an execution completes it removes
+    itself from the registry and stops its lane.
+
+    Build instances with :meth:`create`, which starts the worker's lane;
+    the constructor only initializes fields.
+    """
+
+    def __init__(
+        self,
+        execution_arn: str,
+        store: ExecutionStore,
+        scheduler: Scheduler,
+        registry: ExecutionRegistry,
+        lane: SerialTaskLane,
+    ) -> None:
+        self._arn: str = execution_arn
+        self._store: ExecutionStore = store
+        self._scheduler: Scheduler = scheduler
+        self._registry: ExecutionRegistry = registry
+        self._lane: SerialTaskLane = lane
+        self._status: InvocationState = InvocationState.PRE_INVOKE
+
+    @classmethod
+    def create(
+        cls,
+        execution_arn: str,
+        store: ExecutionStore,
+        scheduler: Scheduler,
+        registry: ExecutionRegistry,
+    ) -> ExecutionWorker:
+        """Build a worker and start its lane."""
+        lane: SerialTaskLane = SerialTaskLane.create(
+            name=f"execution-worker-{execution_arn}"
+        )
+        return cls(execution_arn, store, scheduler, registry, lane)
+
+    @property
+    def status(self) -> InvocationState:
+        """The current invocation state.
+
+        Authoritative only when read on the lane (inside a task); reads
+        from other threads may observe a value about to change.
+        """
+        return self._status
+
+    def set_status(self, status: InvocationState) -> None:
+        """Set the invocation state directly.
+
+        Used by the checkpoint gate before the invoke lifecycle runs on
+        the lane; the lane otherwise threads status through task outcomes.
+        """
+        self._status = status
+
+    def submit(self, task: ExecutionTask[T]) -> Future[T]:
+        """Enqueue ``task`` on this worker's lane and return its future."""
+        return self._lane.submit(lambda: self._run(task))
+
+    def reload(self) -> Execution:
+        """Load this execution's current state from the store."""
+        return self._store.load(self._arn)
+
+    def stop(self) -> None:
+        """Stop this worker's lane without touching the registry.
+
+        Used by :meth:`ExecutionRegistry.shutdown`, which has already
+        dropped the worker from its map.
+        """
+        self._lane.stop(wait=False)
+
+    def _run(self, task: ExecutionTask[T]) -> T:
+        outcome: TaskOutcome[T] = task.execute(self._status, self)
+        self._status = outcome.next_status
+        if self._status is InvocationState.COMPLETED:
+            self._teardown()
+        return outcome.value
+
+    def _teardown(self) -> None:
+        # Runs on the lane thread. Stop the lane first (without joining,
+        # which would wait on the current thread) so the stop flag is
+        # set before the worker leaves the registry. A caller that was
+        # handed this worker then fails fast on submit and the registry
+        # hands it a fresh lane, rather than enqueueing onto a lane that
+        # is about to exit.
+        self._lane.stop(wait=False)
+        self._registry.remove(self._arn)

aws_durable_execution_sdk_python_testing/worker/lane.py

@@ -0,0 +1,118 @@
+"""A serial task lane.
+
+One daemon consumer thread drains a FIFO queue, running submitted
+callables one at a time. :meth:`SerialTaskLane.submit` returns a
+:class:`concurrent.futures.Future` the caller may block on (such as a
+handler or HTTP thread) or ignore (such as a timer callback). Routing
+every operation for one execution through a single lane makes them run
+in submit order with no overlap, which keeps that execution's state
+consistent when several callers act concurrently. Independent lanes run
+in parallel.
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from concurrent.futures import Future
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class _Stop:
+    """Sentinel enqueued by :meth:`SerialTaskLane.stop` to end the loop."""
+
+
+class SerialTaskLane:
+    """A single-consumer FIFO lane that runs callables serially.
+
+    Tasks submitted from any number of producer threads run one at a
+    time on a dedicated daemon thread, in the order they were enqueued.
+    An exception raised by one task is captured on its future and does
+    not stop the lane.
+
+    Build instances with :meth:`create`, which starts the consumer
+    thread; the constructor only initializes fields.
+    """
+
+    def __init__(self, name: str | None = None) -> None:
+        self._queue: queue.Queue[tuple[Callable[[], object], Future] | _Stop] = (
+            queue.Queue()
+        )
+        self._stopped: bool = False
+        self._stop_lock: threading.Lock = threading.Lock()
+        self._thread: threading.Thread = threading.Thread(
+            target=self._run,
+            name=name or f"serial-task-lane-{id(self):x}",
+            daemon=True,
+        )
+
+    @classmethod
+    def create(cls, name: str | None = None) -> SerialTaskLane:
+        """Build a lane and start its consumer thread."""
+        lane: SerialTaskLane = cls(name)
+        lane._thread.start()  # noqa: SLF001 — own private member
+        return lane
+
+    def submit(self, fn: Callable[[], T]) -> Future[T]:
+        """Enqueue ``fn`` to run on the lane and return its future.
+
+        Raises:
+            RuntimeError: If the lane has been stopped.
+        """
+        future: Future[T] = Future()
+        with self._stop_lock:
+            if self._stopped:
+                msg: str = "SerialTaskLane is stopped"
+                raise RuntimeError(msg)
+            self._queue.put((fn, future))
+        return future
+
+    def _run(self) -> None:
+        while True:
+            item: tuple[Callable[[], object], Future] | _Stop = self._queue.get()
+            try:
+                if isinstance(item, _Stop):
+                    return
+                fn, future = item
+                if not future.set_running_or_notify_cancel():
+                    continue
+                try:
+                    result: object = fn()
+                except BaseException as exc:  # noqa: BLE001 — relayed to the future
+                    future.set_exception(exc)
+                else:
+                    future.set_result(result)
+            finally:
+                # Drop references before blocking on the next get so a
+                # task's inputs/outputs are not pinned between tasks.
+                del item
+                self._queue.task_done()
+
+    def stop(self, *, wait: bool = True) -> None:
+        """Stop the lane after draining queued tasks.
+
+        Tasks already enqueued run to completion before the lane exits.
+        Submitting after ``stop`` raises ``RuntimeError``.
+
+        Args:
+            wait: Join the consumer thread before returning.
+        """
+        with self._stop_lock:
+            if self._stopped:
+                return
+            self._stopped = True
+            self._queue.put(_Stop())
+        if wait:
+            self._thread.join()
+
+    def is_running(self) -> bool:
+        """Return True while the consumer thread is alive."""
+        return self._thread.is_alive()

aws_durable_execution_sdk_python_testing/worker/registry.py

@@ -0,0 +1,124 @@
+"""The registry of per-execution workers."""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, TypeVar
+
+from aws_durable_execution_sdk_python_testing.worker.execution_worker import (
+    ExecutionWorker,
+)
+
+if TYPE_CHECKING:
+    from concurrent.futures import Future
+
+    from aws_durable_execution_sdk_python_testing.scheduler import Scheduler
+    from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
+    from aws_durable_execution_sdk_python_testing.worker.task import ExecutionTask
+
+T = TypeVar("T")
+
+
+class ExecutionRegistry:
+    """Owns one :class:`ExecutionWorker` per execution ARN.
+
+    Creates a worker the first time an execution is acted on, hands the
+    same worker out for every later operation on that execution, and
+    drops it once the execution completes. Lookups and mutations are
+    guarded so concurrent callers always resolve to the same worker for
+    a given ARN.
+    """
+
+    def __init__(self, store: ExecutionStore, scheduler: Scheduler) -> None:
+        self._store = store
+        self._scheduler = scheduler
+        self._workers: dict[str, ExecutionWorker] = {}
+        self._lock = threading.Lock()
+        self._shut_down: bool = False
+
+    # Submitting can lose a race with a worker tearing its own lane
+    # down (it stops the lane, then leaves the registry), so a
+    # handed-out worker may reject the submit. Evicting the dead worker
+    # and retrying resolves to a fresh lane, which is never stopped, so
+    # a couple of passes always suffice; the bound only guards against a
+    # pathological storm of completions.
+    _MAX_SUBMIT_ATTEMPTS: int = 5
+
+    def get_or_create(self, execution_arn: str) -> ExecutionWorker:
+        """Return the worker for ``execution_arn``, creating it if absent.
+
+        Raises:
+            RuntimeError: After ``shutdown``. A lane created then would
+                outlive the runner that owned it.
+        """
+        with self._lock:
+            if self._shut_down:
+                msg: str = "execution registry is shut down"
+                raise RuntimeError(msg)
+            worker: ExecutionWorker | None = self._workers.get(execution_arn)
+            if worker is None:
+                worker = ExecutionWorker.create(
+                    execution_arn, self._store, self._scheduler, self
+                )
+                self._workers[execution_arn] = worker
+            return worker
+
+    def submit(self, execution_arn: str, task: ExecutionTask[T]) -> Future[T]:
+        """Submit ``task`` to ``execution_arn``'s worker and return its
+        future.
+
+        Resolves the race where the worker tears its own lane down
+        between hand-out and submit: if the handed-out worker rejects
+        the submit, evict it and retry on a fresh worker. The task is
+        enqueued at most once, so a retry never double-runs it.
+
+        Raises:
+            RuntimeError: If a live lane could not be obtained within
+                ``_MAX_SUBMIT_ATTEMPTS`` (not expected in practice).
+        """
+        last_error: RuntimeError | None = None
+        for _ in range(self._MAX_SUBMIT_ATTEMPTS):
+            worker: ExecutionWorker = self.get_or_create(execution_arn)
+            try:
+                return worker.submit(task)
+            except RuntimeError as err:
+                last_error = err
+                self._evict_if_current(execution_arn, worker)
+        msg: str = "could not obtain a live worker lane"
+        raise RuntimeError(msg) from last_error
+
+    def get(self, execution_arn: str) -> ExecutionWorker | None:
+        """Return the worker for ``execution_arn`` if one exists."""
+        with self._lock:
+            return self._workers.get(execution_arn)
+
+    def remove(self, execution_arn: str) -> None:
+        """Drop the worker for ``execution_arn`` if present."""
+        with self._lock:
+            self._workers.pop(execution_arn, None)
+
+    def active_count(self) -> int:
+        """Return the number of live workers."""
+        with self._lock:
+            return len(self._workers)
+
+    def shutdown(self) -> None:
+        """Stop every worker's lane and drop all workers.
+
+        Called when the owning runner shuts down so per-execution lane
+        threads do not outlive it: an execution that never reached a
+        terminal status still has a live lane, which this stops.
+        """
+        with self._lock:
+            self._shut_down = True
+            workers: list[ExecutionWorker] = list(self._workers.values())
+            self._workers.clear()
+        for worker in workers:
+            worker.stop()
+
+    def _evict_if_current(self, execution_arn: str, worker: ExecutionWorker) -> None:
+        """Drop ``worker`` only if it is still the mapped worker, so a
+        replacement created by another thread is not clobbered."""
+        with self._lock:
+            if self._workers.get(execution_arn) is worker:
+                del self._workers[execution_arn]

aws_durable_execution_sdk_python_testing/worker/status.py

@@ -0,0 +1,27 @@
+"""Per-execution handler-invocation state."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class InvocationState(Enum):
+    """Tracks whether a handler is currently running for one execution.
+
+    Held in memory per execution and never persisted: a persisted
+    ``INVOKING`` state could be read back after a crash as "a handler is
+    running" when none is, stranding the execution behind the gate.
+
+    Transitions:
+
+    * ``PRE_INVOKE`` -> ``INVOKING`` when a handler call is dispatched.
+    * ``INVOKING`` -> ``PRE_INVOKE`` when a handler returns ``PENDING``
+      or fails; a retry re-enters ``PRE_INVOKE`` -> ``INVOKING``.
+    * ``INVOKING`` -> ``COMPLETED`` when the execution reaches a terminal
+      status (handler returned ``SUCCEEDED`` / ``FAILED``, a stop was
+      requested, or the invocation timed out).
+    """
+
+    PRE_INVOKE = "PRE_INVOKE"
+    INVOKING = "INVOKING"
+    COMPLETED = "COMPLETED"

aws_durable_execution_sdk_python_testing/worker/task.py

@@ -0,0 +1,57 @@
+"""The execution-task contract and its outcome."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+    from aws_durable_execution_sdk_python_testing.worker.status import InvocationState
+
+T = TypeVar("T")
+
+
+class WorkerContext(Protocol):
+    """The view of its worker that a task is given while it runs.
+
+    Kept minimal (and defined here, below the task it serves) so the
+    task module does not depend on the concrete worker, avoiding a
+    circular import.
+    """
+
+    @property
+    def status(self) -> InvocationState: ...  # pragma: no cover
+
+    def reload(self) -> Execution: ...  # pragma: no cover
+
+
+@dataclass(frozen=True)
+class TaskOutcome(Generic[T]):
+    """What a task produced: the invocation state to carry forward and
+    the value to return to the task's caller.
+    """
+
+    next_status: InvocationState
+    value: T
+
+
+class ExecutionTask(ABC, Generic[T]):
+    """A single operation against one execution, run on its worker.
+
+    A task runs with exclusive access to its execution's state for the
+    duration of :meth:`execute` (the worker runs one at a time), so it
+    may load, read, and mutate the execution without further
+    synchronization. It receives the current invocation state and returns
+    a :class:`TaskOutcome` carrying the next state and the value to
+    deliver to the caller.
+    """
+
+    @abstractmethod
+    def execute(
+        self,
+        status: InvocationState,
+        worker: WorkerContext,
+    ) -> TaskOutcome[T]:
+        """Run the operation and return its outcome."""