From b52ddaa6539eb5c90391ab5edf2de6b28b16dcb3 Mon Sep 17 00:00:00 2001 From: yaythomas Date: Wed, 17 Jun 2026 21:18:34 +0000 Subject: [PATCH] feat(testing): per-execution worker serialization Replace the per-ARN lock with a per-execution worker that runs every operation for one execution one at a time on its own serial lane. - Worker model: SerialTaskLane, ExecutionWorker, ExecutionRegistry, ExecutionTask, CallableTask. - One checkpoint write transaction in checkpoint/core.py shared by the web and in process paths. - Checkpoint lifecycle actions returned as effects applied after the write. The reentrant notifier is gone. - The invocation gate lives on the worker status. - The blocking invoke runs off the lane so handler callbacks can run. - Store integration test on memory, filesystem, and sqlite. Caught and fixed a non-atomic write in FileSystemStore. - Runner accepts an optional store for store injection testing. - docs/architecture.md: new, describes the worker model. Closes #435 --- .../README.md | 64 +- ...dar-python-test-framework-architecture.svg | 1 - .../dar-python-test-framework-event-flow.svg | 1 - .../docs/architecture.md | 539 +++++++++ .../checkpoint/core.py | 111 +- .../checkpoint/effects.py | 50 + .../checkpoint/processor.py | 70 +- .../checkpoint/transformer.py | 27 +- .../client.py | 30 +- .../execution.py | 7 +- .../executor.py | 1062 +++++++++-------- .../observer.py | 111 +- .../runner.py | 20 +- .../stores/filesystem.py | 10 +- .../web/server.py | 2 +- .../worker/__init__.py | 9 + .../worker/checkpoint_tasks.py | 119 ++ .../worker/execution_worker.py | 106 ++ .../worker/lane.py | 118 ++ .../worker/registry.py | 57 + .../worker/status.py | 27 + .../worker/task.py | 57 + .../tests/checkpoint/processor_test.py | 18 +- .../tests/checkpoint/transformer_test.py | 77 +- .../tests/client_test.py | 16 +- .../tests/e2e/store_backed_lifecycle_test.py | 123 ++ .../tests/executor_checkpoint_test.py | 10 +- .../tests/executor_invariants_test.py | 28 +- .../tests/executor_test.py | 112 +- .../tests/observer_test.py | 284 ++--- .../tests/worker/__init__.py | 1 + .../tests/worker/checkpoint_tasks_test.py | 122 ++ .../tests/worker/concurrency_test.py | 152 +++ .../tests/worker/lane_test.py | 155 +++ .../tests/worker/worker_test.py | 125 ++ 35 files changed, 2774 insertions(+), 1047 deletions(-) delete mode 100644 packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-architecture.svg delete mode 100644 packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-event-flow.svg create mode 100644 packages/aws-durable-execution-sdk-python-testing/docs/architecture.md create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/effects.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/__init__.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/checkpoint_tasks.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/execution_worker.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/lane.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/registry.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/status.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/task.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/store_backed_lifecycle_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/worker/__init__.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/worker/checkpoint_tasks_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/worker/concurrency_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/worker/lane_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/worker/worker_test.py diff --git a/packages/aws-durable-execution-sdk-python-testing/README.md b/packages/aws-durable-execution-sdk-python-testing/README.md index 4731527c..d6e4dc74 100644 --- a/packages/aws-durable-execution-sdk-python-testing/README.md +++ b/packages/aws-durable-execution-sdk-python-testing/README.md @@ -112,66 +112,10 @@ def test_my_durable_functions(): assert three_result.result == '"5 6"' ``` ## Architecture -![Durable Functions Python Test Framework Architecture](assets/dar-python-test-framework-architecture.svg) - -## Event Flow -![Event Flow Sequence Diagram](assets/dar-python-test-framework-event-flow.svg) - -1. **DurableTestRunner** starts execution via **Executor** -2. **Executor** creates **Execution** and schedules initial invocation -3. During execution, checkpoints are processed by **CheckpointProcessor** -4. **Individual Processors** transform operation updates and may trigger events -5. **ExecutionNotifier** broadcasts events to **Executor** (observer) -6. **Executor** updates **Execution** state based on events -7. **Execution** completion triggers final event notifications -8. **DurableTestRunner** run() blocks until it receives completion event, and then returns `DurableFunctionTestResult`. - -## Major Components - -### Core Execution Flow -- **DurableTestRunner** - Main entry point that orchestrates test execution -- **Executor** - Manages execution lifecycle. Mutates Execution. -- **Execution** - Represents the state and operations of a single durable execution - -### Service Client Integration -- **InMemoryServiceClient** - Replaces AWS Lambda service client for local testing. Injected into SDK via `DurableExecutionInvocationInputWithClient` - -### Checkpoint Processing Pipeline -- **CheckpointProcessor** - Orchestrates operation transformations and validation -- **Individual Validators** - Validate operation updates and state transitions -- **Individual Processors** - Transform operation updates into operations (step, wait, callback, context, execution) - -### Execution status changes (Observer Pattern) -- **ExecutionNotifier** - Notifies observers of execution events -- **ExecutionObserver** - Interface for receiving execution lifecycle events -- **Executor** implements `ExecutionObserver` to handle completion events - -## Component Relationships - -### 1. DurableTestRunner → Executor → Execution -- **DurableTestRunner** serves as the main API entry point and sets up all components -- **Executor** manages the execution lifecycle, handling invocations and state transitions -- **Execution** maintains the state of operations and completion status - -### 2. Service Client Injection -- **DurableTestRunner** creates **InMemoryServiceClient** with **CheckpointProcessor** -- **InProcessInvoker** injects the service client into SDK via `DurableExecutionInvocationInputWithClient` -- When durable functions call checkpoint operations, they're intercepted by **InMemoryServiceClient** -- **InMemoryServiceClient** delegates to **CheckpointProcessor** for local processing - -### 3. CheckpointProcessor → Individual Validators → Individual Processors -- **CheckpointProcessor** orchestrates the checkpoint processing pipeline -- **Individual Validators** (CheckpointValidator, TransitionsValidator, and operation-specific validators) ensure operation updates are valid -- **Individual Processors** (StepProcessor, WaitProcessor, etc.) transform `OperationUpdate` into `Operation` - -### 4. Observer Pattern Flow -The observer pattern enables loose coupling between checkpoint processing and execution management: - -1. **CheckpointProcessor** processes operation updates -2. **Individual Processors** detect state changes (completion, failures, timer scheduling) -3. **ExecutionNotifier** broadcasts events to registered observers -4. **Executor** (as ExecutionObserver) receives notifications and updates **Execution** state -5. **Execution** complete_* methods finalize the execution state + +See [docs/architecture.md](docs/architecture.md) for framework +internals. It covers the components, the worker model, the checkpoint +flow, pagination, and a map of the code. ## Documentation diff --git a/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-architecture.svg b/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-architecture.svg deleted file mode 100644 index 0d8fd6d5..00000000 --- a/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-architecture.svg +++ /dev/null @@ -1 +0,0 @@ -Service ClientExecution LifecycleCheckpoint ProcessingOperation Processors (Strategy Pattern)Operation Validators (Strategy Pattern)Observer PatternDurableServiceClientcheckpoint()get_execution_state()stop()checkpoint()get_execution_state()stop()InMemoryServiceClientcheckpoint_processor: CheckpointProcessorcheckpoint_processor: CheckpointProcessorcheckpoint()get_execution_state()stop()checkpoint()get_execution_state()stop()InProcessInvokerhandler: Callableservice_client: InMemoryServiceClienthandler: Callableservice_client: InMemoryServiceClientcreate_invocation_input()invoke()create_invocation_input()invoke()Executorstore: ExecutionStorescheduler: Schedulerinvoker: Invokerstart_execution()complete_execution()fail_execution()on_completed()on_failed()on_wait_timer_scheduled()on_step_retry_scheduled()Executiondurable_execution_arn: stroperations: list[Operation]is_complete: boolstart()complete_success()complete_fail()complete_wait()complete_retry()Schedulercall_later()create_event()CheckpointProcessorstore: ExecutionStorescheduler: Schedulernotifier: ExecutionNotifiertransformer: OperationTransformerprocess_checkpoint()add_execution_observer()Processes operation updatesthrough individual processorsand validators, then notifiesobservers of state changesCheckpointValidatorvalidate_input()TransitionsValidatorvalidate_transitions()OperationProcessor«note: Translates OperationUpdate to Operation»process()StepProcessorWaitProcessorCallbackProcessorContextProcessorExecutionProcessorOperationValidatorvalidate()Strategy Pattern: Each validatorimplements specific validationlogic for different operation typesStepValidatorWaitValidatorCallbackValidatorContextValidatorExecutionValidatorInvokeValidatorExecutionObserveron_completed()on_failed()on_wait_timer_scheduled()on_step_retry_scheduled()ExecutionNotifierobservers: list[ExecutionObserver]add_observer()notify_completed()notify_failed()notify_wait_timer_scheduled()notify_step_retry_scheduled()DurableTestRunnerhandler: Callableservice_client: InMemoryServiceClientexecutor: Executorrun()close()InMemoryServiceClientReplaces AWS Lambda service clientfor local testing. Injected intoSDK via DurableExecutionInvocationInputWithClientto intercept checkpoint callscreatesusesmanagescomplete_success()complete_fail()usesimplementsimplementsdelegates toinjects into SDKusesusesusesusesusesusescall_later/create_eventnotifiesnotifies via ExecutionNotifiernotify_completed()notify_failed()notify_wait_timer_scheduled()notify_step_retry_scheduled() \ No newline at end of file diff --git a/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-event-flow.svg b/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-event-flow.svg deleted file mode 100644 index fbd55ab0..00000000 --- a/packages/aws-durable-execution-sdk-python-testing/assets/dar-python-test-framework-event-flow.svg +++ /dev/null @@ -1 +0,0 @@ -DurableTestRunnerDurableTestRunnerExecutorExecutorExecutionExecutionCheckpointProcessorCheckpointProcessorIndividual ProcessorsIndividual ProcessorsExecutionNotifierExecutionNotifier1. start execution2. create & schedule invocation3. process checkpoints4. transform operation updates4. trigger events5. broadcast events (observer)6. update state based on events7. completion triggers final notifications7. final event notifications8. DurableFunctionTestResult \ No newline at end of file diff --git a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md new file mode 100644 index 00000000..c46f076a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md @@ -0,0 +1,539 @@ +# Architecture + +This is the on ramp for developers working on the testing framework +itself. It describes how the framework is put together and how to find +your way around the code. + +If you are writing tests for your own durable functions, start at the +[README quickstart](../README.md) instead. + +## What this framework is + +The framework runs durable function code locally. It serves the durable +execution service contract that the function's SDK calls back into. That +contract covers checkpoint, get state, and callbacks. A test can drive a +durable function to completion with no deployed service. The observable +behaviour stays the same as production. + +Three runners cover different needs. + +- `DurableFunctionTestRunner` runs the durable function in the test's + own interpreter. It is fast and deterministic. It uses no network. It + is the default for unit and integration tests. +- `WebRunner` serves the same contract over HTTP. A function running in + a separate Lambda process can reach it over the wire. One example is + `sam local start-lambda`. Use it to exercise the real Lambda runtime + and HTTP client path locally. +- `DurableFunctionCloudTestRunner` invokes a function already deployed + to AWS and polls for completion. It is a thin client over the real + service. It is not a local implementation. + +This doc covers the first two. They share nearly all of their internals. +The cloud runner only drives a deployed function. It is described under +[Runners at a glance](#runners-at-a-glance). + +## Core building blocks + +### `Execution` in `execution.py` + +This is the persistent state of one durable execution. It holds the +operation list, the counters, the delivery watermark, and the +idempotency slot. Each store serializes it through `to_json_dict` and +`from_json_dict`. The stores are `InMemoryExecutionStore`, +`SQLiteExecutionStore`, and `FileSystemExecutionStore`. + +Stores are local to the process and load by value. A SQLite or +Filesystem store returns a freshly deserialized `Execution` on every +`load`. So fields on an `Execution` instance cannot coordinate +concurrent callers across a load and save window. Coordination lives on +the worker instead. See below. + +### The worker model + +Several callers act on the same execution at once. The function makes +checkpoint and get state calls. Wait and retry timers fire. Callbacks +arrive. A stop can be requested. To keep one execution's state +consistent, all of its operations run on one worker, one at a time. + +- `ExecutionRegistry` in `worker/registry.py` owns one `ExecutionWorker` + per execution ARN. It creates a worker the first time an execution is + acted on. It hands the same worker out for every later operation. It + drops the worker once the execution completes. Lookups are guarded, so + concurrent callers always resolve to the same worker. +- `ExecutionWorker` in `worker/execution_worker.py` runs every operation + for one execution on its lane. It carries the invocation state from + each task into the next. It tears itself down when the execution + reaches a terminal status. +- `SerialTaskLane` in `worker/lane.py` is a FIFO queue with one consumer + thread. `submit(fn)` returns a `Future`. The caller may block on it or + ignore it. Tasks run in submit order with no overlap. An exception on + one task is captured on its future and does not stop the lane. + Independent lanes run in parallel. +- `ExecutionTask` and `TaskOutcome` in `worker/task.py` are the contract + for work run on a worker. A task runs with exclusive access to its + execution for the duration of `execute`. So it may load and mutate the + execution with no further synchronization. It returns the next + invocation state and the value to deliver to the caller. + +Two task shapes cover everything. They live in +`worker/checkpoint_tasks.py`. + +- `CheckpointTask` and `GetStateTask` apply a checkpoint or read state + through `CheckpointProcessor`. +- `CallableTask` runs an arbitrary callable on the lane. It lets the + `Executor` serialize an existing method on the execution's worker + without restating it as a task. + +After any task, the worker carries `COMPLETED` forward once the +execution is terminal. That triggers teardown. Otherwise it keeps the +status the task left. + +### `InvocationState` in `worker/status.py` + +This is an enum with three values. They are `PRE_INVOKE`, `INVOKING`, +and `COMPLETED`. It lives on the worker. It gates handler invocation and +checkpoint validity. It is never persisted. A persisted `INVOKING` could +be read back after a crash as "a handler is running" when none is. That +would strand the execution behind the gate. An absent worker means +`PRE_INVOKE`. + +### `CheckpointRequestDispatcher` in `checkpoint/transformer.py` + +This applies a batch of `OperationUpdate` values to an `Execution` in +place. For each accepted update it dispatches to a processor for that +operation type. The types are step, wait, callback, context, and +execution. It upserts the resulting `Operation`. It records the payload +size. It touches the operation. It returns the lifecycle effects implied +by the batch. See below. It does not build the response. + +### `CheckpointCore` in `checkpoint/core.py` + +This is the one checkpoint write transaction shared by both entry +points. Its `apply` advances `token_sequence` once. It returns one page +of the operations the handler has not yet seen. It advances the delivery +watermark only for what it returns. It records the idempotency entry so +a retried call can replay the same bytes. Its `match_cached` returns the +cached record for a retried pair of `client_token` and +`checkpoint_token`. The caller owns the gate. The caller owns +persistence. The caller applies the returned effects. + +### Effects in `checkpoint/effects.py` + +Applying a batch can imply actions beyond the state change. The +execution completed, which is `Completed`. The execution failed, which +is `Failed`. A callback was created and needs a timeout, which is +`CallbackCreated`. These are returned as data. The caller applies them +after the write finishes. No follow up action runs in the middle of a +write. + +### `OperationPaginatorState` in `execution.py` + +This is a pinned snapshot of an execution's operation list at one +`token_sequence`. It serves every read the handler sees. It gives a page +bounded by bytes for invocation input and `GetDurableExecutionState`. It +gives the delta of operations past `handler_seen_seq` for a checkpoint +response. It advances the delivery watermark forward. Pinning the +sequence at construction means a page or delta sees one consistent +snapshot even while the execution is mutated. It is frozen and cheap. +Its lifetime is one call. + +## Concurrency model + +Every operation on an execution runs as a task on that execution's +worker lane. + +- Operations for the same execution run one at a time, in submit order. + Nothing overlaps. No caller observes or writes half applied state. +- Operations for different executions run on different lanes, in + parallel. + +The public methods on `Executor` resolve the execution's worker and +submit the work, then block on the returned future. Those methods are +`checkpoint_execution`, `get_execution_state`, and the callback and stop +entry points. Timer and callback callbacks submit the same way but do +not block. + +### Invocation runs off the lane + +When a handler runs in a separate process, it calls back over HTTP to +checkpoint and read state. Those calls must run on the execution's lane. +So the blocking handler invocation must not occupy the lane. Otherwise +the handler would wait on a lane it needs in order to make progress. + +`_invoke_handler` splits the invocation into three steps. + +1. Begin. A task on the lane claims the gate from `PRE_INVOKE` to + `INVOKING` and builds the invocation input. +2. Invoke. The blocking call to the function runs off the lane, bounded + by the invocation timeout. While it blocks, the handler's checkpoint + and get state tasks run on the now free lane. +3. Finish. A task on the lane records the response and releases the + gate. It moves `PENDING` to `PRE_INVOKE`, or a terminal result to + `COMPLETED`. + +```mermaid +sequenceDiagram + autonumber + participant Sched as Scheduler + participant Lane as Execution lane + participant Invk as Invoker + participant HTTP as Handler callbacks + + Sched->>Lane: begin (claim gate, build input) + Sched->>Invk: invoke() BLOCKING, off the lane + Note over Sched,Invk: handler running in its process + HTTP->>Lane: checkpoint / get state (run on the free lane) + Invk-->>Sched: handler returned + Sched->>Lane: finish (record response, release gate) +``` + +### Completion teardown + +When a task leaves the execution terminal, the worker removes itself +from the registry and stops its lane without waiting. It runs on that +lane, so a join would wait on the current thread. The stop takes effect +at once. Any later submit for a completed execution is rejected. It does +not run against torn down state. A late gate write cannot bring a removed +worker back. + +## The two counters + +`Execution` holds two counters that only ever increase. They do +different jobs. Confusing them is the most common source of mistakes. + +### `seq_counter` is the internal event counter + +It increases every time operation state changes. That covers every +update applied in a checkpoint, every wait timer fire, every step retry +made ready, every callback resolution, and every terminal transition. It +is applied through `Execution.touch_operation(op_id)`. That call also +records `operation_last_touched_seq[op_id] = seq_counter`. + +### `token_sequence` is the checkpoint response version + +It increases once per accepted non idempotent checkpoint call. The SDK +validates against it when it presents a checkpoint token. The question +is whether this token's sequence is the current one. It is advanced only +by `Execution.advance_token_sequence()`. That runs exactly once inside +the checkpoint write transaction. + +### Rules of thumb + +- Applying a state change wants `touch_operation`. +- Accepting a checkpoint call wants `advance_token_sequence`. Call it + once, whatever the number of updates the call carried. +- Idempotent replays advance neither. +- Pure reads advance neither. + +## The watermark and delta + +`Execution.handler_seen_seq` records the highest `seq_counter` the +handler has been told about. The handler learns of operations through +`InitialExecutionState` on an invocation, or through the `operations` +list in a checkpoint response. The delta for the next response is the +operations whose `operation_last_touched_seq` is greater than +`handler_seen_seq`. + +The watermark advances only when a checkpoint response is built. It never +advances when invocation input is built. That is what makes retries +correct. A failed invocation never moves the watermark. So the retry +delivers the same state again. + +``` + seq_counter handler_seen_seq delta next checkpoint + . . . + 0 0 (none, just started) + touch A = 1 0 {A} + touch B = 2 0 {A, B} + checkpoint 2 (empty, all delivered) + touch C = 3 2 {C} + retry op C = 4 2 {C} (touched again) + checkpoint 4 (empty) +``` + +## Component diagram + +This is the in process shape used by `DurableFunctionTestRunner`. The +`WebRunner` shape is the same except `InProcessInvoker` is replaced by +`LambdaInvoker` plus a Lambda process reaching back over HTTP. + +```mermaid +flowchart TB + subgraph TestProcess[test Python process] + Handler[durable handler] + TestRunner[DurableFunctionTestRunner] + Invoker[InProcessInvoker] + Executor + Registry[ExecutionRegistry] + Worker[ExecutionWorker and lane] + CP[CheckpointProcessor] + Core[CheckpointCore] + Dispatcher[CheckpointRequestDispatcher] + Paginator[OperationPaginatorState] + Client[InMemoryServiceClient] + Store[(InMemoryExecutionStore)] + Scheduler + end + + TestRunner -->|start_execution| Executor + Executor -->|get_or_create| Registry + Registry --> Worker + Executor -->|create_invocation_input| Invoker + Invoker -->|injects| Client + Invoker -.->|calls| Handler + Handler -->|context.step / wait / checkpoint| Client + Client -->|process_checkpoint| CP + Executor -->|checkpoint_execution| Worker + Worker --> CP + CP --> Core + Core --> Dispatcher + Core --> Paginator + CP --> Store + Executor --> Scheduler +``` + +Two points stand out. + +- There are two checkpoint entry points. `Executor.checkpoint_execution` + is used by the HTTP route. `CheckpointProcessor.process_checkpoint` is + used by `InMemoryServiceClient` in process. Both run on the + execution's worker and share `CheckpointCore`, so behaviour matches. +- The scheduler is async. Its event loop runs on a background thread. + `call_later` returns a `Future`. + +## Invocation state machine + +The gate on each worker governs when a handler can be invoked and when a +checkpoint call is valid. + +```mermaid +stateDiagram-v2 + [*] --> PRE_INVOKE : start_execution + PRE_INVOKE --> INVOKING : begin claims the gate + INVOKING --> PRE_INVOKE : handler returns PENDING + INVOKING --> PRE_INVOKE : handler failed, will retry + INVOKING --> COMPLETED : handler returned SUCCEEDED / FAILED + INVOKING --> COMPLETED : stop / timeout + PRE_INVOKE --> COMPLETED : retry budget exhausted + COMPLETED --> [*] +``` + +The gate enforces a few rules. + +- At most one handler invocation is in flight per execution. A second + invoke attempt that finds `INVOKING` records that a re invoke is needed + and returns. It does not start a new handler. +- `checkpoint_execution` and `get_execution_state` are valid only while + the gate is `INVOKING`. Outside it they reject with + `InvalidParameterValueException`. This rejects stale calls from a + handler process that died and returned after teardown. +- Retries re enter `PRE_INVOKE` then `INVOKING` through the same path. So + every rule still applies. + +## Checkpoint flow + +A checkpoint call runs as a task on the execution's worker. It goes +through a fixed pipeline. The logic is shared by both entry points +through `CheckpointCore`. + +```mermaid +sequenceDiagram + autonumber + participant H as Handler + participant W as Execution worker + participant Core as CheckpointCore + participant S as Store + + H->>W: checkpoint(token, updates, client_token) + W->>S: load(arn) -> Execution + W->>Core: match_cached? + alt cache hit + Core-->>H: cached response + else cache miss + W->>W: gate == INVOKING and token is current + W->>Core: apply(updates) + Note over Core: validate, dispatch, upsert, touch,
advance token_sequence, pin, delta,
one page, advance watermark, record idempotency + Core-->>W: new token, response ops, effects + W->>S: update(execution) + W->>W: apply effects (complete / fail / schedule callback timeout) + W->>W: rearm earliest pending wakeup + W-->>H: CheckpointDurableExecutionResponse + end +``` + +A few points are worth calling out. + +1. Idempotency is checked before the gate and the token. A caller that + retries a call that already succeeded gets the cached response. That + holds even if the execution has moved on. +2. The checkpoint response is one page. If the delta is larger than + `max_invocation_page_bytes`, it is cut to what fits. The watermark + advances only for the returned operations. The rest comes back on the + next checkpoint call. +3. Effects are applied after the write. Completion, failure, and callback + timeout scheduling submit follow up work onto the execution's lane. A + running lane task cannot re enter the lane. So they run once the write + returns. +4. A call with no updates still advances `token_sequence` and can return + operations. An async completion that fired between invocations shows + up as a delta even when `updates` is empty. One example is a wait + timer that succeeded. + +### Operation processors and validators + +The write transaction validates a batch and then applies it. +`CheckpointValidator` runs first. For each update it checks that the +action is allowed for the operation type through +`ValidActionsByOperationTypeValidator`, then runs the validator for that +type. `CheckpointRequestDispatcher` then routes each update to the +processor for its type, which produces the new operation. CHAINED_INVOKE +has a validator but no processor, which is the unimplemented feature. + +```mermaid +flowchart TB + Core[CheckpointCore.apply] + Validator[CheckpointValidator] + Actions[ValidActionsByOperationTypeValidator] + Dispatcher[CheckpointRequestDispatcher] + + subgraph Validators[Per operation validators] + SV[StepOperationValidator] + WV[WaitOperationValidator] + CV[CallbackOperationValidator] + CtxV[ContextOperationValidator] + EV[ExecutionOperationValidator] + IV[ChainedInvokeOperationValidator] + end + + subgraph Processors[Per type processors] + SP[StepProcessor] + WP[WaitProcessor] + CP[CallbackProcessor] + CtxP[ContextProcessor] + EP[ExecutionProcessor] + end + + Core -->|validate_input| Validator + Validator -->|action allowed for type| Actions + Validator -->|run validator per update| Validators + Core -->|apply_updates| Dispatcher + Dispatcher -->|route update by type| Processors +``` + +## Invocation lifecycle + +`_invoke_execution(arn)` is the one scheduling entry point. Every trigger +routes through it. The triggers are the initial start, a re arm after a +checkpoint, a wait timer fire, a step retry fire, and a callback arrival. +Routing through one entry point means the single invocation rule and the +retry budget cannot be bypassed. + +### Retry budget + +A handler exception or an output validation failure increases the +execution's `consecutive_failed_invocation_attempts`. At +`MAX_CONSECUTIVE_FAILED_ATTEMPTS` the execution is failed with the last +observed error. Below that, a retry is scheduled through +`_invoke_execution` after `RETRY_BACKOFF_SECONDS`. Any clean invocation +resets the counter. A clean invocation is one that succeeded or returned +PENDING. So a transient failure followed by a success does not count +against future budget. + +### Invocation timeout + +`execution_timeout` and `invocation_timeout` are set on the runner +constructor and CLI. The invocation timeout bounds one handler call. It +models the function's configured timeout. A call that runs past it is +abandoned. The in flight step stays `STARTED` because no checkpoint was +sent. The execution is invoked again. So on the next invocation the SDK +observes the interrupted step. The execution timeout fails the whole +execution once its deadline passes. + +### Earliest pending scheduler + +When a handler returns `PENDING` it has usually declared async +operations. Each has a future moment to wake on. + +- `wait.scheduled_end_timestamp` for waits in `STARTED`. +- `step.next_attempt_timestamp` for steps in `PENDING`. +- Callbacks keep their own timers. + +`_schedule_earliest_pending(arn)` arms one wake up at the minimum of +those moments. When it fires, it completes every operation whose moment +has passed. Each completion touches its operation. It then routes through +`_invoke_execution` to enter the handler again. It is armed again after a +checkpoint commits, after a `PENDING` invocation, and after it fires. The +cancel then arm sequence runs under `_pending_wakeup_lock`. So concurrent +callers cannot leave an orphan future in the `_pending_wakeup` map. + +## Pagination + +Three reads the handler sees all use `OperationPaginatorState`. The +caller's job differs across them. + +| Read site | Pages? | Marker | +|----------------------------|--------|-----------------------------| +| `InitialExecutionState` | Yes | Real marker if state splits | +| `GetDurableExecutionState` | Yes | Real marker round trip | +| Checkpoint response delta | No | Always `None` | + +### Marker format + +Markers are opaque to the SDK. They are encoded as `"seq:idx"`. `seq` is +the pinned `token_sequence`. `idx` is the next index into the snapshot. A +marker raises `InvalidParameterValueException` in three cases. Its +sequence does not match the pinned snapshot. Its index is out of range. +It does not parse. So a marker is valid only while the execution is at +the same `token_sequence`. Once a checkpoint advances the counter, +outstanding markers become invalid. A marker can never page an old +snapshot. + +### Why the checkpoint response is not paged + +`GetDurableExecutionState` returns one page of a full snapshot. Its +markers index into that snapshot. A checkpoint response returns a delta. +Its markers could not round trip cleanly into `get_execution_state`. So a +delta larger than the byte cap is cut to one page. The watermark advances +only for the returned operations. The rest returns on the next checkpoint +call. + +### `max_invocation_page_bytes` + +This is set on `DurableFunctionTestRunner.__init__` and `WebRunnerConfig`. +Otherwise it is `DEFAULT_MAX_INVOCATION_PAGE_BYTES`. Every read that uses +a paginator honours this cap. A test that wants to exercise a split +across pages passes a small value. + +## Where to look in the code + +| I want to understand | Start here | +|-----------------------------------------------|-----------------------------------------------------------------------------| +| What gets persisted and serialized | `execution.py`, `Execution.to_json_dict` and `from_json_dict` | +| How operations serialize per execution | `worker/registry.py`, `worker/execution_worker.py`, `worker/lane.py` | +| The checkpoint write transaction | `checkpoint/core.py`, `CheckpointCore` | +| The HTTP checkpoint entry point | `executor.py`, `Executor.checkpoint_execution` | +| The in process checkpoint entry point | `checkpoint/processor.py`, `CheckpointProcessor.process_checkpoint` | +| How an `OperationUpdate` becomes an `Operation` | `checkpoint/transformer.py`, `CheckpointRequestDispatcher` | +| State transitions per operation type | `checkpoint/processors/{step,wait,callback,context,execution}.py` | +| Which updates are valid for which state | `checkpoint/validators/valid_actions_by_operation_type.py` | +| The gate, the off lane invoke, the retry budget | `executor.py`, `_invoke_handler`, `_invoke_execution`, `_retry_invocation` | +| Wait and step retry timer scheduling | `executor.py`, `_schedule_earliest_pending`, `_fire_due_and_invoke_handler` | +| Callback timeout and heartbeat timers | `executor.py`, `_schedule_callback_timeouts`, `_on_callback_timeout` | +| The HTTP surface used by `WebRunner` | `web/handlers.py`, `web/routes.py` | +| How tests drive a full run | `tests/executor_checkpoint_test.py`, `tests/executor_invariants_test.py` | +| Worker concurrency invariants | `tests/worker/concurrency_test.py`, `tests/worker/lane_test.py` | + +## Runners at a glance + +| Runner | Handler runs in | Service client | Invoker | Use for | +|----------------------------------|------------------------------|-------------------------|--------------------|--------------------------------------| +| `DurableFunctionTestRunner` | Same Python process as tests | `InMemoryServiceClient` | `InProcessInvoker` | Unit and integration, the default | +| `WebRunner` | Lambda process, SAM CLI etc. | HTTP routes | `LambdaInvoker` | Integration with the Lambda runtime | +| `DurableFunctionCloudTestRunner` | Deployed AWS Lambda | none, calls Lambda directly | real boto3 | Smoke tests against real deployments | + +## See also + +- `docs/error-responses.md` for the error response format the HTTP routes + emit. +- `CONTRIBUTING.md` for the developer workflow. It covers hatch + environments and test commands. diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py index edb7c2bb..60c378eb 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py @@ -1,23 +1,50 @@ -"""The shared checkpoint write-transaction core. +"""The shared checkpoint write-transaction. -Both checkpoint entry points need to detect an idempotent replay before -doing real work. This class holds that shared logic so neither caller -duplicates it. +Both checkpoint entry points apply a batch of updates and compute the +delta of operations the handler has not yet seen. :class:`CheckpointCore` +holds that common logic so the two callers differ only in their gate, +locking, and response type. """ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple + +from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import ( + CheckpointValidator, +) +from aws_durable_execution_sdk_python_testing.execution import ( + CheckpointIdempotencyRecord, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.token import CheckpointToken if TYPE_CHECKING: - from aws_durable_execution_sdk_python_testing.execution import ( - CheckpointIdempotencyRecord, - Execution, + from aws_durable_execution_sdk_python.lambda_service import ( + Operation, + OperationUpdate, + ) + + from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CheckpointEffect, ) + from aws_durable_execution_sdk_python_testing.checkpoint.transformer import ( + CheckpointRequestDispatcher, + ) + from aws_durable_execution_sdk_python_testing.execution import Execution + + +class CheckpointResult(NamedTuple): + """Outcome of applying a checkpoint: the new token, the operations to + return to the handler this round, and the lifecycle effects raised.""" + + checkpoint_token: str + operations: list[Operation] + effects: list[CheckpointEffect] class CheckpointCore: - """Shared checkpoint logic used by both the web and in-process paths.""" + """The checkpoint write-transaction shared by both entry points.""" @staticmethod def match_cached( @@ -26,11 +53,12 @@ def match_cached( client_token: str | None, ) -> CheckpointIdempotencyRecord | None: """Return the cached idempotency record when the incoming - ``(client_token, checkpoint_token)`` matches the last checkpoint. + ``(client_token, checkpoint_token)`` matches the last checkpoint, + else ``None``. - Returns ``None`` when there is no match. A missing or empty - ``client_token`` cannot replay because idempotency requires an - explicit client identifier. + A missing ``client_token`` cannot replay: idempotency requires an + explicit client identifier. The caller builds its own response + type from the returned record. """ if not client_token: return None @@ -43,3 +71,60 @@ def match_cached( ): return None return cached + + @staticmethod + def apply( + execution: Execution, + checkpoint_token: str, + updates: list[OperationUpdate], + client_token: str | None, + dispatcher: CheckpointRequestDispatcher, + max_page_bytes: int, + ) -> CheckpointResult: + """Apply ``updates`` to ``execution`` and compute the response delta. + + Advances ``token_sequence`` exactly once, returns one page of the + operations the handler has not yet seen, advances + ``handler_seen_seq`` only for what is returned, and records the + idempotency entry for a byte-identical replay of a retried call. + The caller is responsible for the invocation gate, locking, + persistence, and applying the returned effects. + """ + effects: list[CheckpointEffect] = [] + if updates: + CheckpointValidator.validate_input(updates, execution) + effects = dispatcher.apply_updates( + execution=execution, + updates=updates, + client_token=client_token, + touch=execution.touch_operation, + ) + + new_token_sequence: int = execution.advance_token_sequence() + + paginator: OperationPaginatorState = OperationPaginatorState.pin(execution) + delta: list[Operation] = paginator.unseen_operations() + response_ops: list[Operation] + response_ops, _ = paginator.page_subset(delta, max_page_bytes) + if response_ops: + highest_delivered_seq: int = max( + execution.operation_last_touched_seq[op.operation_id] + for op in response_ops + ) + paginator.advance_handler_seen(highest_delivered_seq) + + new_token: str = CheckpointToken( + execution_arn=execution.durable_execution_arn, + token_sequence=new_token_sequence, + invocation_id=execution.current_invocation_id, + ).to_str() + + execution.last_checkpoint = CheckpointIdempotencyRecord( + client_token=client_token or "", + inbound_checkpoint_token=checkpoint_token, + outbound_checkpoint_token=new_token, + operations=list(response_ops), + next_marker=None, + ) + + return CheckpointResult(new_token, response_ops, effects) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/effects.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/effects.py new file mode 100644 index 00000000..9472028c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/effects.py @@ -0,0 +1,50 @@ +"""Lifecycle effects produced by applying checkpoint updates. + +Applying a batch of operation updates can imply actions on the execution +beyond the state change itself: the execution completed, it failed, or a +callback was created and now needs a timeout. These are returned as data +so the caller can act on them after the write finishes, rather than +triggering them in the middle of it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.lambda_service import ( + CallbackOptions, + ErrorObject, + ) + + from aws_durable_execution_sdk_python_testing.token import CallbackToken + + +@dataclass(frozen=True) +class Completed: + """The execution completed successfully with ``result``.""" + + execution_arn: str + result: str | None = None + + +@dataclass(frozen=True) +class Failed: + """The execution failed with ``error``.""" + + execution_arn: str + error: ErrorObject + + +@dataclass(frozen=True) +class CallbackCreated: + """A callback was created and its timeout (if any) must be scheduled.""" + + execution_arn: str + operation_id: str + callback_options: CallbackOptions | None + callback_token: CallbackToken + + +CheckpointEffect = Completed | Failed | CallbackCreated diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py index 4b78b6fd..3144e08d 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py @@ -20,17 +20,10 @@ from aws_durable_execution_sdk_python_testing.checkpoint.transformer import ( CheckpointRequestDispatcher, ) -from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import ( - CheckpointValidator, -) from aws_durable_execution_sdk_python_testing.exceptions import ( InvalidParameterValueException, ) -from aws_durable_execution_sdk_python_testing.execution import ( - CheckpointIdempotencyRecord, - OperationPaginatorState, -) -from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier +from aws_durable_execution_sdk_python_testing.observer import apply_effects from aws_durable_execution_sdk_python_testing.token import CheckpointToken @@ -38,6 +31,7 @@ from aws_durable_execution_sdk_python.lambda_service import OperationUpdate from aws_durable_execution_sdk_python_testing.execution import Execution + from aws_durable_execution_sdk_python_testing.observer import ExecutionObserver from aws_durable_execution_sdk_python_testing.scheduler import Scheduler from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore @@ -61,13 +55,13 @@ def __init__( max_page_bytes: int = DEFAULT_MAX_INVOCATION_PAGE_BYTES, ): self._store = store - self._notifier = ExecutionNotifier() + self._observers: list[ExecutionObserver] = [] self._dispatcher = CheckpointRequestDispatcher() self._max_page_bytes = max_page_bytes - def add_execution_observer(self, observer) -> None: + def add_execution_observer(self, observer: ExecutionObserver) -> None: """Add observer for execution events.""" - self._notifier.add_observer(observer) + self._observers.append(observer) def process_checkpoint( self, @@ -99,48 +93,24 @@ def process_checkpoint( msg = "Invalid checkpoint token" raise InvalidParameterValueException(msg) - if updates: - CheckpointValidator.validate_input(updates, execution) - self._dispatcher.apply_updates( - execution=execution, - updates=updates, - client_token=client_token, - notifier=self._notifier, - touch=execution.touch_operation, - ) - - new_token_sequence = execution.advance_token_sequence() - - paginator = OperationPaginatorState.pin(execution) - delta = paginator.unseen_operations() - response_ops, _ = paginator.page_subset(delta, self._max_page_bytes) - if response_ops: - highest_delivered_seq = max( - execution.operation_last_touched_seq[op.operation_id] - for op in response_ops - ) - paginator.advance_handler_seen(highest_delivered_seq) - - new_token = CheckpointToken( - execution_arn=execution.durable_execution_arn, - token_sequence=new_token_sequence, - invocation_id=execution.current_invocation_id, - ).to_str() - - execution.last_checkpoint = CheckpointIdempotencyRecord( - client_token=client_token or "", - inbound_checkpoint_token=checkpoint_token, - outbound_checkpoint_token=new_token, - operations=list(response_ops), - next_marker=None, + result = CheckpointCore.apply( + execution, + checkpoint_token, + updates, + client_token, + self._dispatcher, + self._max_page_bytes, ) self._store.update(execution) + for observer in self._observers: + apply_effects(result.effects, observer) + return CheckpointOutput( - checkpoint_token=new_token, + checkpoint_token=result.checkpoint_token, new_execution_state=CheckpointUpdatedExecutionState( - operations=response_ops, + operations=result.operations, next_marker=None, ), ) @@ -170,8 +140,10 @@ def _maybe_replay_cached( checkpoint_token: str, client_token: str | None, ) -> CheckpointOutput | None: - """Replay cached CheckpointOutput when the incoming call matches - the last checkpoint. Returns ``None`` when there is no match. + """Replay the cached response when the incoming + ``(client_token, checkpoint_token)`` matches the last checkpoint. + + Returns ``None`` when there is no matching cached record. """ cached = CheckpointCore.match_cached(execution, checkpoint_token, client_token) if cached is None: diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/transformer.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/transformer.py index 228fc575..de0909fb 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/transformer.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/transformer.py @@ -1,9 +1,8 @@ """Checkpoint request dispatcher. -Routes each ``OperationUpdate`` in a checkpoint batch to a -type-specific processor (step, wait, callback, context, execution), -which returns the resulting ``Operation`` to upsert into the -execution's operation list. +Routes each ``OperationUpdate`` in a checkpoint to its type-specific +processor (step, wait, callback, context, execution) and applies the +result to the execution. """ from __future__ import annotations @@ -34,6 +33,7 @@ from aws_durable_execution_sdk_python_testing.exceptions import ( InvalidParameterValueException, ) +from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier if TYPE_CHECKING: @@ -44,11 +44,13 @@ OperationUpdate, ) + from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CheckpointEffect, + ) from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import ( OperationProcessor, ) from aws_durable_execution_sdk_python_testing.execution import Execution - from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier class CheckpointRequestDispatcher: @@ -81,9 +83,8 @@ def apply_updates( execution: Execution, updates: list[OperationUpdate], client_token: str | None, # noqa: ARG002 — reserved for future idempotency diagnostics - notifier: ExecutionNotifier, touch: Callable[[str], None], - ) -> None: + ) -> list[CheckpointEffect]: """Apply ``updates`` to ``execution`` in place. Callers are responsible for running @@ -104,10 +105,12 @@ def apply_updates( production caller) bumps :attr:`Execution.seq_counter` and sets ``operation_last_touched_seq[op_id]``. - No response object is returned. Response construction is the - responsibility of the checkpoint orchestrator - (``Executor.checkpoint_execution``). + Returns the lifecycle effects raised by the updates (completion, + failure, callback creation) for the caller to apply once the + write is done. Response construction is the responsibility of the + checkpoint orchestrator. """ + collector = ExecutionNotifier() op_map = {op.operation_id: op for op in execution.operations} for update in updates: @@ -120,7 +123,7 @@ def apply_updates( updated_op = processor.process( update=update, current_op=current_op, - notifier=notifier, + notifier=collector, execution_arn=execution.durable_execution_arn, ) if updated_op is None: @@ -143,6 +146,8 @@ def apply_updates( execution.updates.extend(updates) execution.update_timestamps.extend(datetime.now(UTC) for _ in updates) + return collector.effects + def _estimate_payload_size(update: OperationUpdate) -> int: """Estimate the on-wire size of an ``OperationUpdate``'s payload. diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/client.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/client.py index a68f0ccf..0e2d8c0d 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/client.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/client.py @@ -12,37 +12,49 @@ from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( CheckpointProcessor, ) +from aws_durable_execution_sdk_python_testing.worker.checkpoint_tasks import ( + CheckpointTask, + GetStateTask, +) +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry class InMemoryServiceClient(DurableServiceClient): """An in-memory service client, that can replace the boto lambda service client.""" - def __init__(self, checkpoint_processor: CheckpointProcessor): + def __init__( + self, + checkpoint_processor: CheckpointProcessor, + registry: ExecutionRegistry, + ): self._checkpoint_processor: CheckpointProcessor = checkpoint_processor + self._registry: ExecutionRegistry = registry def checkpoint( self, - durable_execution_arn: str, # noqa: ARG002 + durable_execution_arn: str, checkpoint_token: str, updates: list[OperationUpdate], client_token: str | None, ) -> CheckpointOutput: - # durable_execution_arn is not used in in-memory testing - return self._checkpoint_processor.process_checkpoint( - checkpoint_token, updates, client_token + worker = self._registry.get_or_create(durable_execution_arn) + task = CheckpointTask( + self._checkpoint_processor, checkpoint_token, updates, client_token ) + return worker.submit(task).result() def get_execution_state( self, - durable_execution_arn: str, # noqa: ARG002 + durable_execution_arn: str, checkpoint_token: str, next_marker: str, max_items: int = 1000, ) -> StateOutput: - # durable_execution_arn is not used in in-memory testing - return self._checkpoint_processor.get_execution_state( - checkpoint_token, next_marker, max_items + worker = self._registry.get_or_create(durable_execution_arn) + task = GetStateTask( + self._checkpoint_processor, checkpoint_token, next_marker, max_items ) + return worker.submit(task).result() def stop(self, execution_arn: str, payload: bytes | None) -> datetime.datetime: # noqa: ARG002 # TODO: implement diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index d92f32a2..e8bcaa3b 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -49,10 +49,10 @@ class ExecutionStatus(Enum): class CheckpointIdempotencyRecord: """Single-slot cache of the most recent accepted checkpoint response. - A retried checkpoint call with the same + Single-slot cache of the most recent accepted checkpoint response. ``(client_token, inbound_checkpoint_token)`` pair is entitled to a - byte-identical response. This record is what we compare against and - replay from. + byte-identical response; this record is what we compare + against and replay from. """ client_token: str @@ -611,6 +611,7 @@ class OperationPaginatorState: after a checkpoint response is constructed and its ops are known. + Pinning ``token_sequence`` and ``operations`` at construction guarantees that paginated reads and delta computation see the same snapshot, even if the underlying :class:`Execution` is mutated diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py index 183daf9e..56cdd2cd 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py @@ -7,8 +7,8 @@ import threading import time import uuid +from contextlib import AbstractContextManager, nullcontext from datetime import UTC, datetime -from enum import Enum from typing import TYPE_CHECKING from aws_durable_execution_sdk_python.execution import ( @@ -75,54 +75,40 @@ Execution as ExecutionSummary, ) from aws_durable_execution_sdk_python_testing.observer import ( - ExecutionNotifier, ExecutionObserver, + apply_effects, ) from aws_durable_execution_sdk_python_testing.token import ( CallbackToken, CheckpointToken, ) +from aws_durable_execution_sdk_python_testing.worker.checkpoint_tasks import ( + CallableTask, +) +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry +from aws_durable_execution_sdk_python_testing.worker.status import InvocationState if TYPE_CHECKING: from collections.abc import Awaitable, Callable from concurrent.futures import Future + from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CheckpointEffect, + ) from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( CheckpointProcessor, ) - from aws_durable_execution_sdk_python_testing.invoker import Invoker + from aws_durable_execution_sdk_python_testing.invoker import ( + Invoker, + InvokeResponse, + ) from aws_durable_execution_sdk_python_testing.scheduler import Event, Scheduler from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore logger = logging.getLogger(__name__) -class InvocationState(Enum): - """Per-ARN handler-invocation state machine. - - Lives on :class:`Executor` keyed by execution ARN, never persisted - : a crashed - ``INVOKING`` state would strand the execution forever on restart - because the gate would read "someone is invoking" while nobody is. - - Transitions: - - * ``PRE_INVOKE`` → ``INVOKING`` — when ``_invoke_execution`` - dispatches a handler call. - * ``INVOKING`` → ``PRE_INVOKE`` — when a handler returns - ``PENDING`` or fails (retry will re-enter ``PRE_INVOKE`` → - ``INVOKING``). - * ``INVOKING`` → ``COMPLETED`` — when the execution reaches a - terminal status (handler returned ``SUCCEEDED`` / ``FAILED``, - stop requested, timeout). - """ - - PRE_INVOKE = "PRE_INVOKE" - INVOKING = "INVOKING" - COMPLETED = "COMPLETED" - - class Executor(ExecutionObserver): MAX_CONSECUTIVE_FAILED_ATTEMPTS: int = 5 RETRY_BACKOFF_SECONDS: int = 5 @@ -135,42 +121,30 @@ def __init__( checkpoint_processor: CheckpointProcessor, max_invocation_page_bytes: int | None = None, invocation_timeout_seconds: int = 900, + registry: ExecutionRegistry | None = None, ): self._store = store self._scheduler = scheduler self._invoker = invoker self._checkpoint_processor = checkpoint_processor + self._registry = ( + registry if registry is not None else ExecutionRegistry(store, scheduler) + ) self._invocation_timeout_seconds = invocation_timeout_seconds self._dispatcher = CheckpointRequestDispatcher() - self._notifier = ExecutionNotifier() - # Self-observe so processors called via our own dispatcher can - # schedule wait / retry / callback timers through the normal - # observer path. - self._notifier.add_observer(self) self._max_invocation_page_bytes = ( max_invocation_page_bytes if max_invocation_page_bytes is not None else DEFAULT_MAX_INVOCATION_PAGE_BYTES ) - # Per-ARN lock infrastructure. Locks live on the - # Executor, not the Execution, because the SQLite / Filesystem - # stores return a freshly deserialised Execution on every - # load — an instance-level lock would protect nothing across - # the load/save window. Entries are deliberately leaked on - # execution termination; removing them would create a - # lock-identity race. - self._arn_locks: dict[str, threading.Lock] = {} - self._arn_locks_supervisor: threading.Lock = threading.Lock() - # Per-ARN handler-invocation state machine. Gates - # _invoke_execution so at most one handler invocation is in - # flight per execution. Never persisted: an absent - # entry is PRE_INVOKE by convention. - self._invocation_state: dict[str, InvocationState] = {} - # Single earliest-pending wake-up timer per execution. - # Replaces the per-op call_later pattern for wait timers and - # step retries so concurrent async completions all fire on - # their earliest aggregate moment rather than chaining. - # Entries cancelled on terminal transitions. Never persisted. + # Guards the pending-wakeup map below so concurrent callers + # re-arming the earliest-pending timer can't leak a Future + # into the dict. + self._pending_wakeup_lock: threading.Lock = threading.Lock() + # Single earliest-pending wake-up timer per execution: wait + # timers and step retries share one timer that fires at their + # earliest aggregate moment. Entries are cancelled on terminal + # transitions and never persisted. self._pending_wakeup: dict[str, Future] = {} self._completion_events: dict[str, Event] = {} self._callback_timeouts: dict[str, Future] = {} @@ -409,29 +383,58 @@ def stop_execution( Raises: ResourceNotFoundException: If execution does not exist """ - with self._lock_for(execution_arn): - execution = self.get_execution(execution_arn) - - if execution.is_complete: - # Idempotent: return the existing stop timestamp - execution_op = execution.get_operation_execution_started() - stop_timestamp = execution_op.end_timestamp or datetime.now(UTC) - return StopDurableExecutionResponse(stop_timestamp=stop_timestamp) - - # Use provided error or create a default one - stop_error = error or ErrorObject.from_message( - "Execution stopped by user request" - ) + return ( + self._registry.get_or_create(execution_arn) + .submit(CallableTask(lambda: self._apply_stop(execution_arn, error))) + .result() + ) - # Stop sets TERMINATED close status (different from fail) - logger.exception("[%s] Stopping execution.", execution_arn) - execution.complete_stopped(error=stop_error) # Sets CloseStatus.TERMINATED - self._store.update(execution) - self._complete_events(execution_arn=execution_arn) + def _apply_stop( + self, execution_arn: str, error: ErrorObject | None + ) -> StopDurableExecutionResponse: + execution = self.get_execution(execution_arn) + + if execution.is_complete: + # Idempotent: return the existing stop timestamp + execution_op = execution.get_operation_execution_started() + stop_timestamp = execution_op.end_timestamp or datetime.now(UTC) + return StopDurableExecutionResponse(stop_timestamp=stop_timestamp) + + # Use provided error or create a default one + stop_error = error or ErrorObject.from_message( + "Execution stopped by user request" + ) + + # Stop sets TERMINATED close status (different from fail) + logger.exception("[%s] Stopping execution.", execution_arn) + execution.complete_stopped(error=stop_error) # Sets CloseStatus.TERMINATED + self._store.update(execution) + self._complete_events(execution_arn=execution_arn) return StopDurableExecutionResponse(stop_timestamp=datetime.now(UTC)) def get_execution_state( + self, + execution_arn: str, + checkpoint_token: str | None = None, + marker: str | None = None, + max_items: int | None = None, + ) -> GetDurableExecutionStateResponse: + """Return a page of operations, serialized on the execution's worker.""" + worker = self._registry.get_or_create(execution_arn) + return ( + self._registry.get_or_create(execution_arn) + .submit( + CallableTask( + lambda: self._get_execution_state( + execution_arn, checkpoint_token, marker, max_items + ) + ) + ) + .result() + ) + + def _get_execution_state( self, execution_arn: str, checkpoint_token: str | None = None, @@ -441,8 +444,8 @@ def get_execution_state( """Return a page of operations from the pinned snapshot. Valid only while the execution is ``INVOKING``. The - call is a pure read: no ``handler_seen_seq`` advance, no - ``token_sequence`` bump, no idempotency mutation. + call is a pure read: no ``handler_seen_seq`` advance, + no ``token_sequence`` bump, no idempotency mutation. Raises: ResourceNotFoundException: execution does not exist. @@ -450,24 +453,20 @@ def get_execution_state( is not ``INVOKING``, the token is stale, or the marker does not resolve against the pinned sequence. """ - with self._lock_for(execution_arn): - execution = self.get_execution(execution_arn) + execution = self.get_execution(execution_arn) - # the relevant invariant gate check. - if ( - self._invocation_state.get(execution_arn, InvocationState.PRE_INVOKE) - is not InvocationState.INVOKING - ): - msg = "Invalid checkpoint token" - raise InvalidParameterValueException(msg) + # the relevant invariant gate check. + if self._invocation_gate(execution_arn) is not InvocationState.INVOKING: + msg = "Invalid checkpoint token" + raise InvalidParameterValueException(msg) - # the relevant invariant token check. - if not self._is_current_token(execution, checkpoint_token): - msg = "Invalid checkpoint token" - raise InvalidParameterValueException(msg) + # the relevant invariant token check. + if not self._is_current_token(execution, checkpoint_token): + msg = "Invalid checkpoint token" + raise InvalidParameterValueException(msg) - paginator = OperationPaginatorState.pin(execution) - ops, next_marker = paginator.page(marker, self._max_invocation_page_bytes) + paginator = OperationPaginatorState.pin(execution) + ops, next_marker = paginator.page(marker, self._max_invocation_page_bytes) return GetDurableExecutionStateResponse(operations=ops, next_marker=next_marker) @@ -513,7 +512,6 @@ def get_execution_history( # Generate events from update history (one event per update). # Each checkpoint update increments the history event counter, - # so every recorded transition becomes a separate history event. updates: list[OperationUpdate] = execution.updates timestamps: list[datetime] = execution.update_timestamps # Track cumulative step_details per operation for retry details @@ -790,6 +788,27 @@ def checkpoint_execution( checkpoint_token: str, updates: list[OperationUpdate] | None = None, client_token: str | None = None, + ) -> CheckpointDurableExecutionResponse: + """Process a checkpoint, serializing it on the execution's worker. + + Routes through the per-execution worker so checkpoints for one + execution never overlap. + """ + worker = self._registry.get_or_create(execution_arn) + return worker.submit( + CallableTask( + lambda: self._checkpoint_execution( + execution_arn, checkpoint_token, updates, client_token + ) + ) + ).result() + + def _checkpoint_execution( + self, + execution_arn: str, + checkpoint_token: str, + updates: list[OperationUpdate] | None = None, + client_token: str | None = None, ) -> CheckpointDurableExecutionResponse: """Process a checkpoint request for an execution. @@ -798,9 +817,9 @@ def checkpoint_execution( truncates to one page, and advances ``handler_seen_seq`` only for the ops actually returned. - The full load-mutate-save sequence runs under the per-ARN - lock so concurrent callers against the same execution - can never read a half-applied mutation. + Runs on the execution's worker lane, so concurrent callers + against the same execution can never read a half-applied + mutation. Raises: ResourceNotFoundException: If the execution does not exist. @@ -808,92 +827,55 @@ def checkpoint_execution( invalid (wrong ``token_sequence`` or the execution is already complete). """ - with self._lock_for(execution_arn): - execution = self.get_execution(execution_arn) - - # the relevant invariant: idempotency first. A caller retrying a previously- - # successful checkpoint is entitled to the cached response - # even if the execution has since moved on (though in - # practice it won't — the SDK waits for a response before - # advancing). Idempotency check runs before the token check. - cached = self._maybe_replay_cached( - execution, checkpoint_token, client_token - ) - if cached is not None: - return cached + execution = self.get_execution(execution_arn) - # the relevant invariant: invocation-state gate. Rejects stale checkpoints - # from handler processes that died or timed out and came - # back after we tore down. - if ( - self._invocation_state.get(execution_arn, InvocationState.PRE_INVOKE) - is not InvocationState.INVOKING - ): - msg = "Invalid checkpoint token" - raise InvalidParameterValueException(msg) - - if not self._is_current_token(execution, checkpoint_token): - msg = "Invalid checkpoint token" - raise InvalidParameterValueException(msg) - - if updates: - CheckpointValidator.validate_input(updates, execution) - self._dispatcher.apply_updates( - execution=execution, - updates=updates, - client_token=client_token, - notifier=self._notifier, - touch=execution.touch_operation, - ) + # the relevant invariant: idempotency first. A caller retrying a previously- + # successful checkpoint is entitled to the cached response + # even if the execution has since moved on (though in + # practice it won't — the SDK waits for a response before + # advancing). Idempotency check runs before the token check. + cached = self._maybe_replay_cached(execution, checkpoint_token, client_token) + if cached is not None: + return cached + + # Invocation-state gate: reject stale checkpoints from + # handler processes that died or timed out and came back after + # the worker tore down. + if self._invocation_gate(execution_arn) is not InvocationState.INVOKING: + msg = "Invalid checkpoint token" + raise InvalidParameterValueException(msg) - new_token_sequence = execution.advance_token_sequence() + if not self._is_current_token(execution, checkpoint_token): + msg = "Invalid checkpoint token" + raise InvalidParameterValueException(msg) - paginator = OperationPaginatorState.pin(execution) - delta = paginator.unseen_operations() - response_ops, _ = paginator.page_subset( - delta, self._max_invocation_page_bytes - ) - # the relevant invariant: checkpoint response is always single-page. If the - # delta is larger than one page, advance handler_seen_seq - # only to cover what we actually returned; the rest stays - # in the next delta. - if response_ops: - highest_delivered_seq = max( - execution.operation_last_touched_seq[op.operation_id] - for op in response_ops - ) - paginator.advance_handler_seen(highest_delivered_seq) - - new_token = CheckpointToken( - execution_arn=execution.durable_execution_arn, - token_sequence=new_token_sequence, - invocation_id=execution.current_invocation_id, - ).to_str() - - # the relevant invariant: cache the response so a retried call with the - # same (client_token, inbound_checkpoint_token) gets a - # byte-identical replay. - execution.last_checkpoint = CheckpointIdempotencyRecord( - client_token=client_token or "", - inbound_checkpoint_token=checkpoint_token, - outbound_checkpoint_token=new_token, - operations=list(response_ops), - next_marker=None, - ) + result = CheckpointCore.apply( + execution, + checkpoint_token, + updates or [], + client_token, + self._dispatcher, + self._max_invocation_page_bytes, + ) - self._store.update(execution) + self._store.update(execution) - response = CheckpointDurableExecutionResponse( - checkpoint_token=new_token, - new_execution_state=CheckpointUpdatedExecutionState( - operations=response_ops, - next_marker=None, - ), - ) + response = CheckpointDurableExecutionResponse( + checkpoint_token=result.checkpoint_token, + new_execution_state=CheckpointUpdatedExecutionState( + operations=result.operations, + next_marker=None, + ), + ) + + # Apply lifecycle effects after the write, not mid-apply: + # completion and failure submit follow-up work onto this + # execution's lane, which a running lane task cannot re-enter. + apply_effects(result.effects, self) - # the relevant invariant: re-arm the earliest-pending wake-up outside the - # lock so the next scheduler-driven completion fires at the - # minimum pending timestamp across Wait / Step ops. + # Re-arm the earliest-pending wake-up so the next + # scheduler-driven completion fires at the minimum pending + # timestamp across Wait / Step ops. self._schedule_earliest_pending(execution_arn) return response @@ -906,7 +888,9 @@ def _maybe_replay_cached( ) -> CheckpointDurableExecutionResponse | None: """Replay the cached response for a retried checkpoint call. - Returns ``None`` when there is no idempotency match. + Returns ``None`` unless the last checkpoint matches the incoming + ``(client_token, checkpoint_token)`` pair. A match is a + byte-identical replay; no state is mutated. """ cached = CheckpointCore.match_cached(execution, checkpoint_token, client_token) if cached is None: @@ -939,39 +923,32 @@ def _is_current_token( and parsed.invocation_id == execution.current_invocation_id ) - def _lock_for(self, arn: str) -> threading.Lock: - """Return the per-ARN ``threading.Lock`` for an execution. + def _invocation_gate(self, arn: str) -> InvocationState: + """Read the invocation gate from the execution's worker. - Lazily created on first access. Entries are deliberately - leaked — see :class:`Executor` docstring and the design - for why removal would introduce a lock-identity race. + Absent worker means no invocation has started, i.e. PRE_INVOKE. """ - with self._arn_locks_supervisor: - lock = self._arn_locks.get(arn) - if lock is None: - lock = threading.Lock() - self._arn_locks[arn] = lock - return lock + worker = self._registry.get(arn) + return worker.status if worker is not None else InvocationState.PRE_INVOKE + + def _set_invocation_gate(self, arn: str, status: InvocationState) -> None: + """Set the invocation gate on the execution's worker.""" + self._registry.get_or_create(arn).set_status(status) def _release_gate(self, arn: str, to: InvocationState) -> None: - """Transition ``_invocation_state[arn]`` out of ``INVOKING``. + """Transition the invocation gate out of ``INVOKING``. Called from ``_on_invoke_finished_ok`` (PENDING → PRE_INVOKE or terminal → COMPLETED) and ``_handle_invoke_failure`` - (FAILED → PRE_INVOKE so a retry can re-enter). Caller MUST - hold ``_lock_for(arn)`` around the state read that justified - the transition. + (FAILED → PRE_INVOKE so a retry can re-enter). """ - self._invocation_state[arn] = to + self._set_invocation_gate(arn, to) def _cleanup_execution_state(self, arn: str) -> None: """Tear down per-execution Executor-side state when an - execution reaches a terminal status. Removes entries from - ``_callback_timeouts`` and ``_callback_heartbeats`` and cancels - pending futures. **Does NOT remove ``_arn_locks[arn]``** - : removal would let two threads acquire different lock - objects for the same ARN across cleanup boundaries, breaking - the serialisation guarantee. + execution reaches a terminal status. Removes entries from + ``_callback_timeouts`` and ``_callback_heartbeats`` and cancels + pending futures. """ # Cancel and drop callback timers. for callback_id in list(self._callback_timeouts): @@ -983,7 +960,6 @@ def _cleanup_execution_state(self, arn: str) -> None: if future is not None and not future.done(): future.cancel() self._completion_events.pop(arn, None) - self._invocation_state.pop(arn, None) # Cancel the earliest-pending wake-up timer on terminal # transitions — no further re-invokes needed. pending = self._pending_wakeup.pop(arn, None) @@ -997,8 +973,10 @@ def _earliest_pending_timestamp(execution: Execution) -> datetime | None: Sources: - * ``wait.scheduled_end_timestamp`` for ``STARTED`` Wait ops. - * ``step.next_attempt_timestamp`` for ``PENDING`` Step ops. + * ``wait.scheduled_end_timestamp`` for ``STARTED`` Wait ops + + * ``step.next_attempt_timestamp`` for ``PENDING`` Step ops + Callback timeouts are deliberately excluded: they keep their existing per-callback timers and would @@ -1042,7 +1020,7 @@ def _schedule_earliest_pending(self, execution_arn: str) -> None: if execution.is_complete: # Drop any stale wake-up without re-arming. - with self._arn_locks_supervisor: + with self._pending_wakeup_lock: existing = self._pending_wakeup.pop(execution_arn, None) if existing is not None and not existing.done(): existing.cancel() @@ -1050,7 +1028,7 @@ def _schedule_earliest_pending(self, execution_arn: str) -> None: earliest = self._earliest_pending_timestamp(execution) if earliest is None: - with self._arn_locks_supervisor: + with self._pending_wakeup_lock: existing = self._pending_wakeup.pop(execution_arn, None) if existing is not None and not existing.done(): existing.cancel() @@ -1062,7 +1040,7 @@ def _schedule_earliest_pending(self, execution_arn: str) -> None: completion_event = self._completion_events.get(execution_arn) # Cancel-then-arm atomically under the supervisor lock so # concurrent callers can't leave orphan Futures in the dict. - with self._arn_locks_supervisor: + with self._pending_wakeup_lock: existing = self._pending_wakeup.pop(execution_arn, None) future = self._scheduler.call_later( self._fire_due_and_invoke_handler(execution_arn), @@ -1073,6 +1051,59 @@ def _schedule_earliest_pending(self, execution_arn: str) -> None: if existing is not None and not existing.done(): existing.cancel() + def _fire_due_operations(self, execution_arn: str) -> bool: + """Complete any Wait/Step ops whose scheduled moment has passed. + + Returns True if at least one operation completed (so the caller + re-invokes the handler). + """ + try: + execution = self._store.load(execution_arn) + except Exception: # noqa: BLE001 + return False + + if execution.is_complete: + return False + + now = datetime.now(UTC) + completed_any = False + for op in list(execution.operations): + if ( + op.operation_type == OperationType.WAIT + and op.status == OperationStatus.STARTED + and op.wait_details is not None + and op.wait_details.scheduled_end_timestamp is not None + and op.wait_details.scheduled_end_timestamp <= now + ): + try: + execution.complete_wait(op.operation_id) + completed_any = True + except Exception: # noqa: BLE001 + logger.exception( + "[%s] earliest-pending: complete_wait failed for %s", + execution_arn, + op.operation_id, + ) + elif ( + op.operation_type == OperationType.STEP + and op.status == OperationStatus.PENDING + and op.step_details is not None + and op.step_details.next_attempt_timestamp is not None + and op.step_details.next_attempt_timestamp <= now + ): + try: + execution.complete_retry(op.operation_id) + completed_any = True + except Exception: # noqa: BLE001 + logger.exception( + "[%s] earliest-pending: complete_retry failed for %s", + execution_arn, + op.operation_id, + ) + if completed_any: + self._store.update(execution) + return completed_any + def _fire_due_and_invoke_handler( self, execution_arn: str ) -> Callable[[], Awaitable[None]]: @@ -1087,53 +1118,12 @@ def _fire_due_and_invoke_handler( """ async def fire_due() -> None: - with self._lock_for(execution_arn): - try: - execution = self._store.load(execution_arn) - except Exception: # noqa: BLE001 - return - if execution.is_complete: - return - - now = datetime.now(UTC) - completed_any = False - for op in list(execution.operations): - if ( - op.operation_type == OperationType.WAIT - and op.status == OperationStatus.STARTED - and op.wait_details is not None - and op.wait_details.scheduled_end_timestamp is not None - and op.wait_details.scheduled_end_timestamp <= now - ): - try: - execution.complete_wait(op.operation_id) - completed_any = True - except Exception: # noqa: BLE001 - logger.exception( - "[%s] earliest-pending: complete_wait failed for %s", - execution_arn, - op.operation_id, - ) - elif ( - op.operation_type == OperationType.STEP - and op.status == OperationStatus.PENDING - and op.step_details is not None - and op.step_details.next_attempt_timestamp is not None - and op.step_details.next_attempt_timestamp <= now - ): - try: - execution.complete_retry(op.operation_id) - completed_any = True - except Exception: # noqa: BLE001 - logger.exception( - "[%s] earliest-pending: complete_retry failed for %s", - execution_arn, - op.operation_id, - ) - if completed_any: - self._store.update(execution) - - # Outside the lock: invoke and arm the next wake-up. + completed_any = await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit(CallableTask(lambda: self._fire_due_operations(execution_arn))) + .result() + ) + # Outside the lane: invoke and arm the next wake-up. if completed_any: self._invoke_execution(execution_arn) self._schedule_earliest_pending(execution_arn) @@ -1164,12 +1154,13 @@ def send_callback_success( try: callback_token = CallbackToken.from_str(callback_id) - with self._lock_for(callback_token.execution_arn): - execution = self.get_execution(callback_token.execution_arn) - execution.complete_callback_success(callback_id, result) - self._store.update(execution) - self._cleanup_callback_timeouts(callback_id) - self._invoke_execution(callback_token.execution_arn) + arn = callback_token.execution_arn + self._registry.get_or_create(arn).submit( + CallableTask( + lambda: self._complete_callback_success(callback_id, result) + ) + ).result() + self._invoke_execution(arn) logger.info("Callback success completed for callback_id: %s", callback_id) except Exception as e: msg = f"Failed to process callback success: {e}" @@ -1177,6 +1168,15 @@ def send_callback_success( return SendDurableExecutionCallbackSuccessResponse() + def _complete_callback_success( + self, callback_id: str, result: bytes | None + ) -> None: + callback_token = CallbackToken.from_str(callback_id) + execution = self.get_execution(callback_token.execution_arn) + execution.complete_callback_success(callback_id, result) + self._store.update(execution) + self._cleanup_callback_timeouts(callback_id) + def send_callback_failure( self, callback_id: str, @@ -1200,8 +1200,8 @@ def send_callback_failure( raise InvalidParameterValueException(msg) # FR: a callback failure sent with no error payload must leave the - # error fields absent, not synthesize an empty ErrorMessage. - # ErrorObject.to_dict() omits None fields, so an all-None + # error fields absent, not synthesize an empty + # ErrorMessage. ErrorObject.to_dict() omits None fields, so an all-None # error serializes to {} and the SDK surfaces undefined fields. callback_error: ErrorObject = error or ErrorObject( message=None, type=None, data=None, stack_trace=None @@ -1209,12 +1209,13 @@ def send_callback_failure( try: callback_token: CallbackToken = CallbackToken.from_str(callback_id) - with self._lock_for(callback_token.execution_arn): - execution: Execution = self.get_execution(callback_token.execution_arn) - execution.complete_callback_failure(callback_id, callback_error) - self._store.update(execution) - self._cleanup_callback_timeouts(callback_id) - self._invoke_execution(callback_token.execution_arn) + arn = callback_token.execution_arn + self._registry.get_or_create(arn).submit( + CallableTask( + lambda: self._complete_callback_failure(callback_id, callback_error) + ) + ).result() + self._invoke_execution(arn) logger.info("Callback failure completed for callback_id: %s", callback_id) except Exception as e: msg = f"Failed to process callback failure: {e}" @@ -1222,6 +1223,15 @@ def send_callback_failure( return SendDurableExecutionCallbackFailureResponse() + def _complete_callback_failure( + self, callback_id: str, callback_error: ErrorObject + ) -> None: + callback_token = CallbackToken.from_str(callback_id) + execution = self.get_execution(callback_token.execution_arn) + execution.complete_callback_failure(callback_id, callback_error) + self._store.update(execution) + self._cleanup_callback_timeouts(callback_id) + def send_callback_heartbeat( self, callback_id: str ) -> SendDurableExecutionCallbackHeartbeatResponse: @@ -1243,19 +1253,10 @@ def send_callback_heartbeat( try: callback_token: CallbackToken = CallbackToken.from_str(callback_id) - with self._lock_for(callback_token.execution_arn): - execution: Execution = self.get_execution(callback_token.execution_arn) - - # Find callback operation to verify it exists and is active - _, operation = execution.find_callback_operation(callback_id) - if operation.status != OperationStatus.STARTED: - msg = f"Callback {callback_id} is not active" - raise ResourceNotFoundException(msg) - - # Reset heartbeat timeout if configured - self._reset_callback_heartbeat_timeout( - callback_id, execution.durable_execution_arn - ) + arn = callback_token.execution_arn + self._registry.get_or_create(arn).submit( + CallableTask(lambda: self._process_callback_heartbeat(callback_id)) + ).result() logger.info("Callback heartbeat processed for callback_id: %s", callback_id) except Exception as e: msg = f"Failed to process callback heartbeat: {e}" @@ -1263,6 +1264,20 @@ def send_callback_heartbeat( return SendDurableExecutionCallbackHeartbeatResponse() + def _process_callback_heartbeat(self, callback_id: str) -> None: + callback_token = CallbackToken.from_str(callback_id) + execution = self.get_execution(callback_token.execution_arn) + + # Verify the callback operation exists and is active. + _, operation = execution.find_callback_operation(callback_id) + if operation.status != OperationStatus.STARTED: + msg = f"Callback {callback_id} is not active" + raise ResourceNotFoundException(msg) + + self._reset_callback_heartbeat_timeout( + callback_id, execution.durable_execution_arn + ) + def _validate_invocation_response_and_store( self, execution_arn: str, @@ -1322,68 +1337,176 @@ def _validate_invocation_response_and_store( ) raise IllegalStateException(msg_unexpected_status) - def _invoke_handler(self, execution_arn: str) -> Callable[[], Awaitable[None]]: - """Create a parameterless callable that captures execution arn for the scheduler.""" + def _begin_invocation( + self, execution_arn: str + ) -> tuple[Execution, DurableExecutionInvocationInput] | None: + """Claim the invocation gate and build the handler input. - async def invoke() -> None: - # Under the per-ARN lock, claim the invocation gate and - # snapshot the execution for input construction. Release - # the lock BEFORE the blocking Lambda invoke — otherwise - # the customer handler's HTTP callbacks (which land on - # other threads and also acquire this lock) would block - # forever, deadlocking the execution. - execution: Execution - invocation_input: DurableExecutionInvocationInput - try: - with self._lock_for(execution_arn): - execution = self._store.load(execution_arn) + Returns the execution and its invocation input when this call + claims the gate (PRE_INVOKE -> INVOKING); returns None when the + execution is already complete, the gate is COMPLETED, or another + invocation is in flight (recording a deferred re-invoke). + """ + execution = self._store.load(execution_arn) - if execution.is_complete: - logger.info( - "[%s] Execution already completed, ignoring invoke", - execution_arn, - ) - self._invocation_state[execution_arn] = ( - InvocationState.COMPLETED - ) - return + if execution.is_complete: + logger.info( + "[%s] Execution already completed, ignoring invoke", + execution_arn, + ) + self._set_invocation_gate(execution_arn, InvocationState.COMPLETED) + return None - current_state = self._invocation_state.get( - execution_arn, InvocationState.PRE_INVOKE - ) - if current_state is InvocationState.COMPLETED: - logger.info( - "[%s] Invocation state already COMPLETED, not re-invoking", - execution_arn, - ) - return - if current_state is InvocationState.INVOKING: - # Another invocation is already in flight - # . Record that we wanted to re-invoke; - # the in-flight handler's post-invoke hook will - # consult ``needs_reinvoke`` and schedule a - # follow-up if set. - execution.needs_reinvoke = True - self._store.save(execution) - logger.info( - "[%s] Handler already INVOKING; deferring re-invoke", - execution_arn, - ) - return + current_state = self._invocation_gate(execution_arn) + if current_state is InvocationState.COMPLETED: + logger.info( + "[%s] Invocation state already COMPLETED, not re-invoking", + execution_arn, + ) + return None + if current_state is InvocationState.INVOKING: + # Another invocation is in flight. Record the + # re-invoke intent; the in-flight handler's post-invoke + # hook consults needs_reinvoke and schedules a follow-up. + execution.needs_reinvoke = True + self._store.save(execution) + logger.info( + "[%s] Handler already INVOKING; deferring re-invoke", + execution_arn, + ) + return None - # Claim the gate. the relevant invariant: at most one handler - # invocation per execution in flight. - self._invocation_state[execution_arn] = InvocationState.INVOKING - execution.begin_new_invocation() + # Claim the gate. the relevant invariant: at most one handler invocation per + # execution in flight. + self._set_invocation_gate(execution_arn, InvocationState.INVOKING) + execution.begin_new_invocation() + invocation_input = self._invoker.create_invocation_input(execution=execution) + self._store.save(execution) + return execution, invocation_input - invocation_input = self._invoker.create_invocation_input( - execution=execution - ) - self._store.save(execution) + def _finish_invocation( + self, + execution_arn: str, + invoke_response: InvokeResponse, + invocation_start: datetime, + invocation_end: datetime, + ) -> None: + """Apply a completed handler invocation. + + Records the attempt, validates the response (completing or + retrying), resets the retry counter on a clean run, and releases + the gate to PRE_INVOKE (re-invoking if a trigger was deferred) or + COMPLETED. + """ + execution = self._store.load(execution_arn) + + execution.record_invocation_completion( + invocation_start, + invocation_end, + invoke_response.request_id, + ) + self._store.save(execution) + + if execution.is_complete: + # A stop / timeout landed while we were invoking — the + # terminal state is authoritative; the handler's + # response does not overwrite it. + logger.info( + "[%s] Execution completed during invocation, ignoring result", + execution_arn, + ) + self._set_invocation_gate(execution_arn, InvocationState.COMPLETED) + return + + response = invoke_response.invocation_output + try: + self._validate_invocation_response_and_store( + execution_arn, response, execution + ) + except (InvalidParameterValueException, IllegalStateException) as e: + logger.warning( + "[%s] Lambda output validation failure: %s", execution_arn, e + ) + error_obj = ErrorObject.from_exception(e) + # Release the gate before retry scheduling so the retry's + # invoke finds PRE_INVOKE. + self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE) + self._retry_invocation(execution, error_obj) + return + + # the relevant invariant: a clean invocation resets the retry counter. + reset_target = self._store.load(execution_arn) + reset_target.consecutive_failed_invocation_attempts = 0 + self._store.save(reset_target) + + # Release the gate: terminal -> COMPLETED, otherwise PRE_INVOKE + # (re-invoking if a trigger was deferred mid-invocation, the relevant invariant). + reloaded = self._store.load(execution_arn) + if reloaded.is_complete: + self._set_invocation_gate(execution_arn, InvocationState.COMPLETED) + should_reinvoke = False + else: + self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE) + should_reinvoke = reloaded.needs_reinvoke + if should_reinvoke: + reloaded.needs_reinvoke = False + self._store.save(reloaded) + + if should_reinvoke: + self._invoke_execution(execution_arn) + else: + # the relevant invariant: arm the earliest-pending wake-up. + self._schedule_earliest_pending(execution_arn) + + def _fail_invocation_not_found( + self, execution_arn: str, error: ErrorObject + ) -> None: + self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE) + self._fail_workflow(execution_arn, error) + + def _retry_after_timeout( + self, + execution_arn: str, + error: ErrorObject, + invocation_start: datetime, + invocation_end: datetime, + ) -> None: + execution = self._store.load(execution_arn) + execution.record_invocation_completion( + invocation_start, invocation_end, str(uuid.uuid4()) + ) + self._store.save(execution) + self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE) + self._retry_invocation(execution, error) + + def _retry_after_error(self, execution_arn: str, error: ErrorObject) -> None: + self._set_invocation_gate(execution_arn, InvocationState.PRE_INVOKE) + try: + execution = self._store.load(execution_arn) + except Exception: # noqa: BLE001 + return + self._retry_invocation(execution, error) + + def _invoke_handler(self, execution_arn: str) -> Callable[[], Awaitable[None]]: + """Create a parameterless callable that captures execution arn for the scheduler.""" + + async def invoke() -> None: + # The gate claim, the blocking Lambda invoke, and the + # post-invoke processing run as separate steps: the two + # state-mutation windows run on the execution's worker lane + # (serialized with checkpoints), while the blocking invoke + # runs off the lane so the handler's checkpoints can run on + # it. + try: + claim = await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit(CallableTask(lambda: self._begin_invocation(execution_arn))) + .result() + ) + if claim is None: + return + execution, invocation_input = claim - # Lock released. Blocking Lambda call happens here so - # HTTP callback threads can acquire the lock for their - # own load-mutate-save windows. invocation_start = datetime.now(UTC) invoke_response = await asyncio.wait_for( asyncio.to_thread( @@ -1395,110 +1518,35 @@ async def invoke() -> None: timeout=self._invocation_timeout_seconds, ) invocation_end = datetime.now(UTC) - - # Re-acquire the lock for post-invoke state updates. - # While we were blocked, the Execution may have been - # mutated by HTTP callback threads; reload fresh. - with self._lock_for(execution_arn): - execution = self._store.load(execution_arn) - - execution.record_invocation_completion( - invocation_start, - invocation_end, - invoke_response.request_id, - ) - self._store.save(execution) - - if execution.is_complete: - # A stop_execution / timeout landed while we - # were invoking — the terminal state is - # authoritative; don't let the handler's - # response overwrite it. - logger.info( - "[%s] Execution completed during invocation, ignoring result", - execution_arn, - ) - self._invocation_state[execution_arn] = ( - InvocationState.COMPLETED + await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit( + CallableTask( + lambda: self._finish_invocation( + execution_arn, + invoke_response, + invocation_start, + invocation_end, + ) ) - return - - response = invoke_response.invocation_output - try: - self._validate_invocation_response_and_store( - execution_arn, response, execution - ) - except (InvalidParameterValueException, IllegalStateException) as e: - logger.warning( - "[%s] Lambda output validation failure: %s", - execution_arn, - e, ) - error_obj = ErrorObject.from_exception(e) - # Release gate before retry scheduling so the - # retry's _invoke_execution call finds PRE_INVOKE. - # Under the per-ARN lock to block concurrent - # triggers. - with self._lock_for(execution_arn): - self._invocation_state[execution_arn] = ( - InvocationState.PRE_INVOKE - ) - self._retry_invocation(execution, error_obj) - return - - # the relevant invariant: a clean invocation (no validation failure, - # no exception) resets the retry counter to zero. The - # prior implementation only ever grew this counter. - with self._lock_for(execution_arn): - reset_target = self._store.load(execution_arn) - reset_target.consecutive_failed_invocation_attempts = 0 - self._store.save(reset_target) - - # Clean return from handler. If the response was - # terminal, _complete_workflow has already set - # is_complete and the next _invoke_execution call - # (if any) will see COMPLETED. If PENDING, release - # the gate to PRE_INVOKE so the next trigger can - # start a fresh invocation. - with self._lock_for(execution_arn): - reloaded = self._store.load(execution_arn) - if reloaded.is_complete: - self._invocation_state[execution_arn] = ( - InvocationState.COMPLETED - ) - should_reinvoke = False - else: - self._invocation_state[execution_arn] = ( - InvocationState.PRE_INVOKE - ) - # the relevant invariant: re-invoke if another trigger - # deferred during this invocation. Clear the - # flag before scheduling so a concurrent - # trigger doesn't set it a second time after - # we've already picked it up. - should_reinvoke = reloaded.needs_reinvoke - if should_reinvoke: - reloaded.needs_reinvoke = False - self._store.save(reloaded) - - if should_reinvoke: - self._invoke_execution(execution_arn) - else: - # the relevant invariant: arm the earliest-pending wake-up so any - # scheduler-driven completion (wait timer, step - # retry) fires on its earliest aggregate moment. - self._schedule_earliest_pending(execution_arn) + .result() + ) except ResourceNotFoundException: logger.warning("[%s] Function No longer exists", execution_arn) error_obj = ErrorObject.from_message(message="Function not found") - # Release gate — _fail_workflow will set COMPLETED - # after writing the terminal state. Transition under - # the lock so a concurrent _invoke_execution can't - # see PRE_INVOKE and race us. - with self._lock_for(execution_arn): - self._invocation_state[execution_arn] = InvocationState.PRE_INVOKE - self._fail_workflow(execution_arn, error_obj) + await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit( + CallableTask( + lambda: self._fail_invocation_not_found( + execution_arn, error_obj + ) + ) + ) + .result() + ) except asyncio.TimeoutError: # Invocation killed by Lambda timeout. Step operations @@ -1510,34 +1558,37 @@ async def invoke() -> None: execution_arn, self._invocation_timeout_seconds, ) - with self._lock_for(execution_arn): - execution = self._store.load(execution_arn) - execution.record_invocation_completion( - invocation_start, - invocation_end, - str(uuid.uuid4()), - ) - self._store.save(execution) - self._invocation_state[execution_arn] = InvocationState.PRE_INVOKE error_obj = ErrorObject.from_message( message=f"Function timed out after {self._invocation_timeout_seconds} seconds" ) - self._retry_invocation(execution, error_obj) + await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit( + CallableTask( + lambda: self._retry_after_timeout( + execution_arn, + error_obj, + invocation_start, + invocation_end, + ) + ) + ) + .result() + ) except Exception as e: # noqa: BLE001 # Handle invocation errors (network, function not found, etc.) logger.warning("[%s] Invocation failed: %s", execution_arn, e) error_obj = ErrorObject.from_exception(e) - # Transition gate + reload execution under a single - # lock acquisition so callers never see a half-updated - # state. - with self._lock_for(execution_arn): - self._invocation_state[execution_arn] = InvocationState.PRE_INVOKE - try: - execution = self._store.load(execution_arn) - except Exception: # noqa: BLE001 - return - self._retry_invocation(execution, error_obj) + await asyncio.to_thread( + lambda: self._registry.get_or_create(execution_arn) + .submit( + CallableTask( + lambda: self._retry_after_error(execution_arn, error_obj) + ) + ) + .result() + ) return invoke @@ -1638,26 +1689,24 @@ def wait_until_complete( def complete_execution(self, execution_arn: str, result: str | None = None) -> None: """Complete execution successfully (COMPLETE_WORKFLOW_EXECUTION decision).""" logger.debug("[%s] Completing execution with result: %s", execution_arn, result) - with self._lock_for(execution_arn): - execution: Execution = self._store.load(execution_arn=execution_arn) - execution.complete_success(result=result) # Sets CloseStatus.COMPLETED - self._store.update(execution) - if execution.result is None: - msg: str = "Execution result is required" - raise IllegalStateException(msg) + execution: Execution = self._store.load(execution_arn=execution_arn) + execution.complete_success(result=result) # Sets CloseStatus.COMPLETED + self._store.update(execution) + if execution.result is None: + msg: str = "Execution result is required" + raise IllegalStateException(msg) self._complete_events(execution_arn=execution_arn) def fail_execution(self, execution_arn: str, error: ErrorObject) -> None: """Fail execution with error (FAIL_WORKFLOW_EXECUTION decision).""" logger.error("[%s] Completing execution with error: %s", execution_arn, error) - with self._lock_for(execution_arn): - execution: Execution = self._store.load(execution_arn=execution_arn) - execution.complete_fail(error=error) # Sets CloseStatus.FAILED - self._store.update(execution) - # set by complete_fail - if execution.result is None: - msg: str = "Execution result is required" - raise IllegalStateException(msg) + execution: Execution = self._store.load(execution_arn=execution_arn) + execution.complete_fail(error=error) # Sets CloseStatus.FAILED + self._store.update(execution) + # set by complete_fail + if execution.result is None: + msg: str = "Execution result is required" + raise IllegalStateException(msg) self._complete_events(execution_arn=execution_arn) # region ExecutionObserver @@ -1670,16 +1719,20 @@ def on_failed(self, execution_arn: str, error: ErrorObject) -> None: self.fail_execution(execution_arn, error) def on_timed_out(self, execution_arn: str, error: ErrorObject) -> None: - """Handle execution timeout (workflow timeout). Observer method triggered by notifier.""" + """Handle execution timeout (workflow timeout).""" logger.exception("[%s] Execution timed out.", execution_arn) - with self._lock_for(execution_arn): - execution: Execution = self._store.load(execution_arn=execution_arn) - execution.complete_timeout(error=error) # Sets CloseStatus.TIMED_OUT - self._store.update(execution) + self._registry.get_or_create(execution_arn).submit( + CallableTask(lambda: self._apply_timeout(execution_arn, error)) + ).result() self._complete_events(execution_arn=execution_arn) + def _apply_timeout(self, execution_arn: str, error: ErrorObject) -> None: + execution = self._store.load(execution_arn=execution_arn) + execution.complete_timeout(error=error) # Sets CloseStatus.TIMED_OUT + self._store.update(execution) + def on_stopped(self, execution_arn: str, error: ErrorObject) -> None: - """Handle execution stop. Observer method triggered by notifier.""" + """Handle execution stop.""" # This should not be called directly - stop_execution handles termination self.fail_execution(execution_arn, error) @@ -1805,27 +1858,32 @@ def _cleanup_callback_timeouts(self, callback_id: str) -> None: if heartbeat_future := self._callback_heartbeats.pop(callback_id, None): heartbeat_future.cancel() + def _apply_callback_timeout(self, callback_id: str, error: ErrorObject) -> None: + callback_token = CallbackToken.from_str(callback_id) + execution = self.get_execution(callback_token.execution_arn) + if execution.is_complete: + return + execution.complete_callback_timeout(callback_id, error) + self._store.update(execution) + def _on_callback_timeout(self, execution_arn: str, callback_id: str) -> None: """Handle callback timeout.""" try: callback_token = CallbackToken.from_str(callback_id) - with self._lock_for(callback_token.execution_arn): - execution = self.get_execution(callback_token.execution_arn) - - if execution.is_complete: - return - - # Fail the callback with timeout error - timeout_error = ErrorObject( - message="Callback timed out", - type=CallbackTimeoutType.TIMEOUT.value, - data=None, - stack_trace=None, + arn = callback_token.execution_arn + timeout_error = ErrorObject( + message="Callback timed out", + type=CallbackTimeoutType.TIMEOUT.value, + data=None, + stack_trace=None, + ) + self._registry.get_or_create(arn).submit( + CallableTask( + lambda: self._apply_callback_timeout(callback_id, timeout_error) ) - execution.complete_callback_timeout(callback_id, timeout_error) - self._store.update(execution) - logger.warning("[%s] Callback %s timed out", execution_arn, callback_id) - self._invoke_execution(callback_token.execution_arn) + ).result() + logger.warning("[%s] Callback %s timed out", execution_arn, callback_id) + self._invoke_execution(arn) except Exception: logger.exception( "[%s] Error processing callback timeout for %s", @@ -1839,26 +1897,22 @@ def _on_callback_heartbeat_timeout( """Handle callback heartbeat timeout.""" try: callback_token = CallbackToken.from_str(callback_id) - with self._lock_for(callback_token.execution_arn): - execution = self.get_execution(callback_token.execution_arn) - - if execution.is_complete: - return - - # Fail the callback with heartbeat timeout error - - heartbeat_error = ErrorObject( - message="Callback timed out on heartbeat", - type=CallbackTimeoutType.HEARTBEAT.value, - data=None, - stack_trace=None, - ) - execution.complete_callback_timeout(callback_id, heartbeat_error) - self._store.update(execution) - logger.warning( - "[%s] Callback %s heartbeat timed out", execution_arn, callback_id + arn = callback_token.execution_arn + heartbeat_error = ErrorObject( + message="Callback timed out on heartbeat", + type=CallbackTimeoutType.HEARTBEAT.value, + data=None, + stack_trace=None, + ) + self._registry.get_or_create(arn).submit( + CallableTask( + lambda: self._apply_callback_timeout(callback_id, heartbeat_error) ) - self._invoke_execution(callback_token.execution_arn) + ).result() + logger.warning( + "[%s] Callback %s heartbeat timed out", execution_arn, callback_id + ) + self._invoke_execution(arn) except Exception: logger.exception( "[%s] Error processing callback heartbeat timeout for %s", diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/observer.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/observer.py index e6f21c84..e1ea7435 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/observer.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/observer.py @@ -1,24 +1,40 @@ -"""Checkpoint processors can notify the Execution of notable event state changes. Observer pattern.""" +"""Execution lifecycle effects: collection and application. + +While a batch of checkpoint updates is applied, the per-type processors +record lifecycle effects (completion, failure, callback creation) on an +:class:`ExecutionNotifier`. The checkpoint orchestrator then hands the +collected effects to :func:`apply_effects`, which drives them onto an +:class:`ExecutionObserver` after the write completes. Collecting first +and applying afterwards keeps effect handling out of the write itself. +""" from __future__ import annotations -import threading from abc import ABC, abstractmethod from typing import TYPE_CHECKING -from aws_durable_execution_sdk_python_testing.token import CallbackToken +from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CallbackCreated, + Completed, + Failed, +) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Iterable from aws_durable_execution_sdk_python.lambda_service import ( - ErrorObject, CallbackOptions, + ErrorObject, ) + from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CheckpointEffect, + ) + from aws_durable_execution_sdk_python_testing.token import CallbackToken + class ExecutionObserver(ABC): - """Observer for execution lifecycle events.""" + """Receives execution lifecycle events.""" @abstractmethod def on_completed(self, execution_arn: str, result: str | None = None) -> None: @@ -44,52 +60,26 @@ def on_callback_created( callback_options: CallbackOptions | None, callback_token: CallbackToken, ) -> None: - """Called when callback is created.""" + """Called when a callback is created.""" class ExecutionNotifier: - """Notifies observers about execution events. Thread-safe.""" + """Collects lifecycle effects raised while applying checkpoint updates. + + Processors record an effect here instead of acting on it; the + orchestrator reads :attr:`effects` once the updates are applied. + """ def __init__(self) -> None: - self._observers: list[ExecutionObserver] = [] - self._lock = threading.RLock() - - def add_observer(self, observer: ExecutionObserver) -> None: - """Add an observer to be notified of execution events.""" - with self._lock: - self._observers.append(observer) - - def _notify_observers(self, method: Callable, *args, **kwargs) -> None: - """Notify all observers by calling the specified method.""" - with self._lock: - observers = self._observers.copy() - for observer in observers: - getattr(observer, method.__name__)(*args, **kwargs) - - # region event emitters + self.effects: list[CheckpointEffect] = [] + def notify_completed(self, execution_arn: str, result: str | None = None) -> None: - """Notify observers about execution completion.""" - self._notify_observers( - ExecutionObserver.on_completed, execution_arn=execution_arn, result=result - ) + """Record that the execution completed successfully.""" + self.effects.append(Completed(execution_arn=execution_arn, result=result)) def notify_failed(self, execution_arn: str, error: ErrorObject) -> None: - """Notify observers about execution failure.""" - self._notify_observers( - ExecutionObserver.on_failed, execution_arn=execution_arn, error=error - ) - - def notify_timed_out(self, execution_arn: str, error: ErrorObject) -> None: - """Notify observers about execution timeout.""" - self._notify_observers( - ExecutionObserver.on_timed_out, execution_arn=execution_arn, error=error - ) - - def notify_stopped(self, execution_arn: str, error: ErrorObject) -> None: - """Notify observers about execution being stopped.""" - self._notify_observers( - ExecutionObserver.on_stopped, execution_arn=execution_arn, error=error - ) + """Record that the execution failed.""" + self.effects.append(Failed(execution_arn=execution_arn, error=error)) def notify_callback_created( self, @@ -98,13 +88,30 @@ def notify_callback_created( callback_options: CallbackOptions | None, callback_token: CallbackToken, ) -> None: - """Notify observers about callback creation.""" - self._notify_observers( - ExecutionObserver.on_callback_created, - execution_arn=execution_arn, - operation_id=operation_id, - callback_options=callback_options, - callback_token=callback_token, + """Record that a callback was created.""" + self.effects.append( + CallbackCreated( + execution_arn=execution_arn, + operation_id=operation_id, + callback_options=callback_options, + callback_token=callback_token, + ) ) - # endregion event emitters + +def apply_effects( + effects: Iterable[CheckpointEffect], observer: ExecutionObserver +) -> None: + """Drive each collected effect onto ``observer``.""" + for effect in effects: + if isinstance(effect, Completed): + observer.on_completed(effect.execution_arn, effect.result) + elif isinstance(effect, Failed): + observer.on_failed(effect.execution_arn, effect.error) + elif isinstance(effect, CallbackCreated): + observer.on_callback_created( + effect.execution_arn, + effect.operation_id, + effect.callback_options, + effect.callback_token, + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/runner.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/runner.py index 2d52e4d2..ba75cc17 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/runner.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/runner.py @@ -38,6 +38,7 @@ DEFAULT_MAX_INVOCATION_PAGE_BYTES, ) from aws_durable_execution_sdk_python_testing.client import InMemoryServiceClient +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry from aws_durable_execution_sdk_python_testing.exceptions import ( DurableFunctionsLocalRunnerError, DurableFunctionsTestError, @@ -113,7 +114,7 @@ class WebRunnerConfig: invocation_timeout_seconds: int = 900 # Invocation-input pagination cap (bytes). Handler receives pages - # up to this size. Larger state is served via + # up to this size; larger state is served via # GetDurableExecutionState. None falls back to Executor default # (5 MB). max_invocation_page_bytes: int | None = None @@ -603,6 +604,7 @@ def __init__( max_invocation_page_bytes: int | None = None, execution_timeout: int = 300, invocation_timeout: int = 900, + store: ExecutionStore | None = None, ): self._execution_timeout = execution_timeout self._invocation_timeout = invocation_timeout @@ -613,14 +615,21 @@ def __init__( ) self._scheduler: Scheduler = Scheduler() self._scheduler.start() - self._store = InMemoryExecutionStore() + # Defaults to the in-memory store. Pass a filesystem or sqlite + # store to run the in-process lifecycle on a disk-backed store. + self._store: ExecutionStore = ( + store if store is not None else InMemoryExecutionStore() + ) self.poll_interval = poll_interval self._checkpoint_processor = CheckpointProcessor( store=self._store, scheduler=self._scheduler, max_page_bytes=self._max_invocation_page_bytes, ) - self._service_client = InMemoryServiceClient(self._checkpoint_processor) + self._registry = ExecutionRegistry(self._store, self._scheduler) + self._service_client = InMemoryServiceClient( + self._checkpoint_processor, self._registry + ) self._invoker = InProcessInvoker( handler, self._service_client, @@ -633,6 +642,7 @@ def __init__( checkpoint_processor=self._checkpoint_processor, max_invocation_page_bytes=self._max_invocation_page_bytes, invocation_timeout_seconds=invocation_timeout, + registry=self._registry, ) # Wire up observer pattern - CheckpointProcessor uses this to notify executor of state changes @@ -863,6 +873,7 @@ def start(self) -> None: checkpoint_processor = CheckpointProcessor( self._store, self._scheduler, max_page_bytes=resolved_max_page_bytes ) + registry = ExecutionRegistry(self._store, self._scheduler) # Create executor with all dependencies including checkpoint processor self._executor = Executor( @@ -872,6 +883,7 @@ def start(self) -> None: checkpoint_processor=checkpoint_processor, max_invocation_page_bytes=resolved_max_page_bytes, invocation_timeout_seconds=self._config.invocation_timeout_seconds, + registry=registry, ) # Add executor as observer to the checkpoint processor @@ -948,7 +960,7 @@ def _create_boto3_client(self) -> Any: class DurableFunctionCloudTestRunner: - """Test runner that executes durable functions against actual AWS Lambda backend. + """Test runner that executes durable functions against the deployed service. This runner invokes deployed Lambda functions and polls for execution completion, providing the same interface as DurableFunctionTestRunner for seamless test diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/stores/filesystem.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/stores/filesystem.py index f0f3154c..eca42df4 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/stores/filesystem.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/stores/filesystem.py @@ -4,6 +4,7 @@ import json import logging +import os from pathlib import Path from aws_durable_execution_sdk_python_testing.exceptions import ( @@ -46,8 +47,15 @@ def save(self, execution: Execution) -> None: file_path = self._get_file_path(execution.durable_execution_arn) data = execution.to_json_dict() - with open(file_path, "w", encoding="utf-8") as f: + # Write to a temporary file in the same directory, then replace + # the target in one atomic step. A concurrent reader then sees + # either the old file or the new file, never a partial or empty + # one. Saves for one execution run on its worker lane, so two + # writes never target the same temp path at once. + tmp_path = file_path.with_name(f"{file_path.name}.tmp") + with open(tmp_path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) + os.replace(tmp_path, file_path) def load(self, execution_arn: str) -> Execution: """Load execution from file system.""" diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/web/server.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/web/server.py index 89cb219b..68cefedc 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/web/server.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/web/server.py @@ -1,4 +1,4 @@ -"""Local testing web server for AWS Lambda Durable Functions that mimics the actual Lambda backend services.""" +"""Local testing web server for AWS Lambda Durable Functions for local testing.""" from __future__ import annotations diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/__init__.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/__init__.py new file mode 100644 index 00000000..4dc4214d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/__init__.py @@ -0,0 +1,9 @@ +"""Per-execution serial actor model for the local runner. + +Each execution is served by a single worker that runs all of its +operations (checkpoint, get-state, invoke, timer, callback, completion) +one at a time on its own lane. Serializing the operations of one +execution keeps its state consistent when several callers act on it at +once; different executions are served by different workers and run in +parallel. +""" diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/checkpoint_tasks.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/checkpoint_tasks.py new file mode 100644 index 00000000..060a133d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/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) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/execution_worker.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/execution_worker.py new file mode 100644 index 00000000..f7459fad --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/execution_worker.py @@ -0,0 +1,106 @@ +"""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 _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, so stop without joining (a join here + # would wait on the current thread). The stop flag takes effect + # at once, rejecting any later submit for a completed execution. + self._registry.remove(self._arn) + self._lane.stop(wait=False) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/lane.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/lane.py new file mode 100644 index 00000000..fb6930c4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/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() diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/registry.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/registry.py new file mode 100644 index 00000000..8aeb46ff --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/registry.py @@ -0,0 +1,57 @@ +"""The registry of per-execution workers.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python_testing.worker.execution_worker import ( + ExecutionWorker, +) + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python_testing.scheduler import Scheduler + from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore + + +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() + + def get_or_create(self, execution_arn: str) -> ExecutionWorker: + """Return the worker for ``execution_arn``, creating it if absent.""" + with self._lock: + 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 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) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/status.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/status.py new file mode 100644 index 00000000..96b6341a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/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" diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/task.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/worker/task.py new file mode 100644 index 00000000..0fa40b03 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/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.""" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py index bff22620..bf1ee527 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py @@ -46,23 +46,17 @@ def test_init(): processor.add_execution_observer(observer) # Should not raise an exception -@patch( - "aws_durable_execution_sdk_python_testing.checkpoint.processor.ExecutionNotifier" -) -def test_add_execution_observer(mock_notifier_class): +def test_add_execution_observer(): """Test adding execution observer.""" store = Mock(spec=ExecutionStore) scheduler = Mock(spec=Scheduler) - mock_notifier_instance = Mock() - mock_notifier_class.return_value = mock_notifier_instance processor = CheckpointProcessor(store, scheduler) observer = Mock() processor.add_execution_observer(observer) - # Verify observer was added through the notifier's public method - mock_notifier_instance.add_observer.assert_called_once_with(observer) + assert observer in processor._observers # noqa: SLF001 def test_process_checkpoint_success(): @@ -112,9 +106,7 @@ def test_process_checkpoint_success(): ) -@patch( - "aws_durable_execution_sdk_python_testing.checkpoint.processor.CheckpointValidator" -) +@patch("aws_durable_execution_sdk_python_testing.checkpoint.core.CheckpointValidator") def test_process_checkpoint_invalid_token_complete_execution(mock_validator): """Test checkpoint processing with complete execution.""" store = Mock(spec=ExecutionStore) @@ -144,9 +136,7 @@ def test_process_checkpoint_invalid_token_complete_execution(mock_validator): processor.process_checkpoint(checkpoint_token, updates, "client-token") -@patch( - "aws_durable_execution_sdk_python_testing.checkpoint.processor.CheckpointValidator" -) +@patch("aws_durable_execution_sdk_python_testing.checkpoint.core.CheckpointValidator") def test_process_checkpoint_invalid_token_sequence(mock_validator): """Test checkpoint processing with invalid token sequence.""" store = Mock(spec=ExecutionStore) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/transformer_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/transformer_test.py index fbb7114d..3d433cdf 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/transformer_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/transformer_test.py @@ -1,11 +1,9 @@ """Unit tests for CheckpointRequestDispatcher. -Covers the new ``apply_updates(execution, updates, client_token, -notifier, touch)`` API introduced in of -. The dispatcher mutates -``execution.operations`` in place, records per-op size in -``execution.operation_size_bytes``, and calls ``touch`` once per -accepted update. +Covers ``apply_updates(execution, updates, client_token, touch)``: it +mutates ``execution.operations`` in place, records per-op size in +``execution.operation_size_bytes``, calls ``touch`` once per accepted +update, and returns the lifecycle effects raised by the updates. """ from __future__ import annotations @@ -20,12 +18,14 @@ OperationUpdate, ) +from aws_durable_execution_sdk_python_testing.checkpoint.effects import Completed from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import ( OperationProcessor, ) from aws_durable_execution_sdk_python_testing.checkpoint.transformer import ( CheckpointRequestDispatcher, ) +from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier from aws_durable_execution_sdk_python_testing.exceptions import ( InvalidParameterValueException, ) @@ -85,7 +85,6 @@ def test_apply_updates_with_empty_list_is_a_noop(): execution=execution, updates=[], client_token=None, - notifier=Mock(), touch=touched.append, ) @@ -113,7 +112,6 @@ def test_apply_updates_unknown_type_raises(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=lambda _: None, ) @@ -135,7 +133,6 @@ def test_apply_updates_skips_ops_when_processor_returns_none(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=touched.append, ) @@ -164,7 +161,6 @@ def test_apply_updates_appends_new_operation_and_touches(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=touched.append, ) @@ -194,7 +190,6 @@ def test_apply_updates_replaces_existing_operation_in_place(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=lambda _: None, ) @@ -233,7 +228,6 @@ def test_apply_updates_preserves_order_across_multiple_updates(): ), ], client_token=None, - notifier=Mock(), touch=touched.append, ) assert execution.operations == [op1, updated_op2, op3] @@ -250,7 +244,6 @@ def test_apply_updates_preserves_order_across_multiple_updates(): ), ], client_token=None, - notifier=Mock(), touch=touched.append, ) assert execution.operations == [op1, updated_op2, op3, new_op4] @@ -289,7 +282,6 @@ def test_apply_updates_dispatches_by_operation_type(): ), ], client_token=None, - notifier=Mock(), touch=lambda _: None, ) @@ -306,7 +298,6 @@ def test_apply_updates_forwards_arn_notifier_and_current_op_to_processor(): processors={OperationType.STEP: processor}, ) execution = _make_execution([existing]) - notifier = Mock() update = OperationUpdate( operation_id="id", @@ -318,7 +309,6 @@ def test_apply_updates_forwards_arn_notifier_and_current_op_to_processor(): execution=execution, updates=[update], client_token=None, - notifier=notifier, touch=lambda _: None, ) @@ -327,10 +317,60 @@ def test_apply_updates_forwards_arn_notifier_and_current_op_to_processor(): ) assert forwarded_update == update assert forwarded_current_op == existing - assert forwarded_notifier is notifier + assert isinstance(forwarded_notifier, ExecutionNotifier) assert forwarded_arn == execution.durable_execution_arn +def test_apply_updates_returns_completion_effect(): + """An EXECUTION SUCCEED update surfaces a Completed effect for the + caller to apply after the write.""" + dispatcher = CheckpointRequestDispatcher() + execution = _make_execution() + + update = OperationUpdate( + operation_id="exec-op", + operation_type=OperationType.EXECUTION, + action=OperationAction.SUCCEED, + payload="final-result", + ) + + effects = dispatcher.apply_updates( + execution=execution, + updates=[update], + client_token=None, + touch=lambda _: None, + ) + + assert effects == [ + Completed(execution_arn=execution.durable_execution_arn, result="final-result") + ] + + +def test_apply_updates_returns_no_effects_for_plain_step(): + """A STEP START produces an operation but no lifecycle effect.""" + new_op = Mock() + new_op.operation_id = "step-op" + dispatcher = CheckpointRequestDispatcher( + processors={OperationType.STEP: _MockProcessor(return_value=new_op)}, + ) + execution = _make_execution() + + effects = dispatcher.apply_updates( + execution=execution, + updates=[ + OperationUpdate( + operation_id="step-op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ) + ], + client_token=None, + touch=lambda _: None, + ) + + assert effects == [] + + def test_apply_updates_records_payload_size_for_paging(): """The sidecar dict feeds OperationPaginatorState._size_for.""" new_op = Mock() @@ -352,7 +392,6 @@ def test_apply_updates_records_payload_size_for_paging(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=lambda _: None, ) @@ -382,7 +421,6 @@ def test_apply_updates_records_size_for_error_payload(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=lambda _: None, ) @@ -412,7 +450,6 @@ def test_apply_updates_records_size_for_bytes_payload(): execution=execution, updates=[update], client_token=None, - notifier=Mock(), touch=lambda _: None, ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/client_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/client_test.py index 15aa3665..d0da1ba2 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/client_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/client_test.py @@ -12,6 +12,12 @@ ) from aws_durable_execution_sdk_python_testing.client import InMemoryServiceClient +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry + + +def _client(processor: Mock) -> InMemoryServiceClient: + """Build a client whose registry uses a stub store/scheduler.""" + return InMemoryServiceClient(processor, ExecutionRegistry(Mock(), Mock())) def test_checkpoint(): @@ -23,7 +29,7 @@ def test_checkpoint(): ) processor.process_checkpoint.return_value = expected_output - client = InMemoryServiceClient(processor) + client = _client(processor) updates = [ OperationUpdate( @@ -52,7 +58,7 @@ def test_get_execution_state(): expected_output = StateOutput(operations=[], next_marker="marker") processor.get_execution_state.return_value = expected_output - client = InMemoryServiceClient(processor) + client = _client(processor) result = client.get_execution_state( "arn:aws:lambda:us-east-1:123456789012:function:test", "token", "marker", 500 @@ -68,7 +74,7 @@ def test_get_execution_state_default_max_items(): expected_output = StateOutput(operations=[], next_marker="marker") processor.get_execution_state.return_value = expected_output - client = InMemoryServiceClient(processor) + client = _client(processor) result = client.get_execution_state( "arn:aws:lambda:us-east-1:123456789012:function:test", "token", "marker" @@ -81,7 +87,7 @@ def test_get_execution_state_default_max_items(): def test_stop(): """Test stop method returns current datetime.""" processor = Mock() - client = InMemoryServiceClient(processor) + client = _client(processor) before = datetime.datetime.now(tz=datetime.UTC) result = client.stop( @@ -96,7 +102,7 @@ def test_stop(): def test_stop_with_none_payload(): """Test stop method with None payload.""" processor = Mock() - client = InMemoryServiceClient(processor) + client = _client(processor) result = client.stop("arn:aws:states:us-east-1:123456789012:execution:test", None) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/store_backed_lifecycle_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/store_backed_lifecycle_test.py new file mode 100644 index 00000000..e8556432 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/store_backed_lifecycle_test.py @@ -0,0 +1,123 @@ +"""End-to-end lifecycle on each store backend. + +The in-memory store returns the same object on every load, so an in +place mutation of a loaded execution survives even with no save. The +filesystem and sqlite stores deserialize a fresh object on every load, +so any mutation the lifecycle fails to persist is lost on the next load. + +Every other test runs on the in-memory store, which hides that class of +bug. This test drives one full durable execution through the executor, +the worker, and the checkpoint path on each store backend. The +execution uses steps, a wait, a child context, and completion, so it +exercises several checkpoints and the completion effect. A mutation that +is not saved would change the result or stall the run on a disk store. +""" + +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import ( + DurableContext, + durable_step, + durable_with_child_context, +) +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.types import StepContext + +from aws_durable_execution_sdk_python_testing.runner import ( + ContextOperation, + DurableFunctionTestResult, + DurableFunctionTestRunner, + StepOperation, +) +from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.sqlite import ( + SQLiteExecutionStore, +) + + +StoreFactory = Callable[[Path], ExecutionStore] + + +def _memory_store(_tmp: Path) -> ExecutionStore: + return InMemoryExecutionStore() + + +def _filesystem_store(tmp: Path) -> ExecutionStore: + return FileSystemExecutionStore.create(tmp / "fs-store") + + +def _sqlite_store(tmp: Path) -> ExecutionStore: + return SQLiteExecutionStore.create_and_initialize(str(tmp / "exec.db")) + + +@pytest.mark.parametrize( + "store_factory", + [_memory_store, _filesystem_store, _sqlite_store], + ids=["memory", "filesystem", "sqlite"], +) +def test_full_lifecycle_persists_on_each_store( + store_factory: StoreFactory, tmp_path: Path +) -> None: + @durable_step + def one(step_context: StepContext, a: int, b: int) -> str: + return f"{a} {b}" + + @durable_step + def two_1(step_context: StepContext, a: int, b: int) -> str: + return f"{a} {b}" + + @durable_step + def two_2(step_context: StepContext, a: int, b: int) -> str: + return f"{b} {a}" + + @durable_with_child_context + def two(ctx: DurableContext, a: int, b: int) -> str: + two_1_result: str = ctx.step(two_1(a, b)) + two_2_result: str = ctx.step(two_2(a, b)) + return f"{two_1_result} {two_2_result}" + + @durable_step + def three(step_context: StepContext, a: int, b: int) -> str: + return f"{a} {b}" + + @durable_execution + def function_under_test(event: Any, context: DurableContext) -> list[str]: + results: list[str] = [] + results.append(context.step(one(1, 2))) + context.wait(Duration.from_seconds(1)) + results.append(context.run_in_child_context(two(3, 4))) + results.append(context.step(three(5, 6))) + return results + + store: ExecutionStore = store_factory(tmp_path) + with DurableFunctionTestRunner( + handler=function_under_test, store=store, execution_timeout=10 + ) as runner: + result: DurableFunctionTestResult = runner.run(input="input str") + + assert result.status is InvocationStatus.SUCCEEDED + assert result.result == json.dumps(["1 2", "3 4 4 3", "5 6"]) + + one_op: StepOperation = result.get_step("one") + assert one_op.result == json.dumps("1 2") + + two_op: ContextOperation = result.get_context("two") + assert two_op.result == json.dumps("3 4 4 3") + + three_op: StepOperation = result.get_step("three") + assert three_op.result == json.dumps("5 6") diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py index 2c014bfa..6771895d 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py @@ -79,8 +79,8 @@ def _make_executor_with_started_execution() -> tuple[ # Simulate an active invocation so the gate lets checkpoint # requests through. _invoke_handler will do this # automatically; tests here pre-prime it. - executor._invocation_state[execution.durable_execution_arn] = ( # noqa: SLF001 - InvocationState.INVOKING + executor._set_invocation_gate( # noqa: SLF001 + execution.durable_execution_arn, InvocationState.INVOKING ) initial_token = CheckpointToken( @@ -434,7 +434,7 @@ def test_invocation_gate_defers_concurrent_trigger(): executor = Executor(store, scheduler, invoker, checkpoint_processor) # Simulate "handler is mid-flight" — gate is INVOKING. - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 # Grab the coroutine _invoke_handler builds and run it directly. # Under the gate check it must NOT call the mocked invoker. @@ -480,7 +480,7 @@ def test_failed_invocation_releases_the_gate(): # in a previous attempt and this failure tips it over). But never # INVOKING. assert ( - executor._invocation_state.get(arn, InvocationState.PRE_INVOKE) # noqa: SLF001 + executor._invocation_gate(arn) # noqa: SLF001 is not InvocationState.INVOKING ) @@ -681,7 +681,7 @@ def test_get_execution_state_gate_and_marker_contracts(): ) # Outside INVOKING: flip the gate to PRE_INVOKE and expect rejection. - executor._invocation_state[arn] = InvocationState.PRE_INVOKE # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.PRE_INVOKE) # noqa: SLF001 with pytest.raises(InvalidParameterValueException): executor.get_execution_state( execution_arn=arn, checkpoint_token=r1.checkpoint_token diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_invariants_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_invariants_test.py index 8114da19..f6ff227b 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_invariants_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_invariants_test.py @@ -85,7 +85,7 @@ def _make_executor() -> tuple[Executor, InMemoryExecutionStore, str, str]: store.save(execution) executor = Executor(store, Mock(), Mock(), Mock()) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 token = CheckpointToken(execution_arn=arn, token_sequence=0).to_str() return executor, store, arn, token @@ -300,14 +300,14 @@ def test_invariant_gate_release_after_transient_failures(): # Drive a few failed attempts (still under MAX=5). for _ in range(3): asyncio.run(executor._invoke_handler(arn)()) # noqa: SLF001 - state = executor._invocation_state.get(arn, InvocationState.PRE_INVOKE) # noqa: SLF001 + state = executor._invocation_gate(arn) # noqa: SLF001 assert state is not InvocationState.INVOKING # After the budget is exhausted the execution is terminal; state # becomes COMPLETED (also not INVOKING). for _ in range(Executor.MAX_CONSECUTIVE_FAILED_ATTEMPTS): asyncio.run(executor._invoke_handler(arn)()) # noqa: SLF001 - state = executor._invocation_state.get(arn, InvocationState.PRE_INVOKE) # noqa: SLF001 + state = executor._invocation_gate(arn) # noqa: SLF001 assert state is not InvocationState.INVOKING @@ -675,7 +675,7 @@ def test_invoke_handler_bails_out_on_completed_gate(): execution.start() store.save(execution) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.COMPLETED # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.COMPLETED) # noqa: SLF001 asyncio.run(executor._invoke_handler(arn)()) # noqa: SLF001 @@ -684,9 +684,8 @@ def test_invoke_handler_bails_out_on_completed_gate(): def test_invoke_handler_bails_out_on_is_complete(): - """Similarly, if the Execution is already is_complete when the - handler starts, the handler transitions state to COMPLETED and - exits without calling the invoker.""" + """If the Execution is already is_complete when the handler starts, + the handler exits without calling the invoker.""" store = InMemoryExecutionStore() invoker = Mock() executor = Executor(store, Mock(), invoker, Mock()) @@ -710,11 +709,10 @@ def test_invoke_handler_bails_out_on_is_complete(): asyncio.run(executor._invoke_handler(arn)()) # noqa: SLF001 assert invoker.invoke.call_count == 0 - # Gate transitioned to COMPLETED. - assert ( - executor._invocation_state.get(arn) # noqa: SLF001 - is InvocationState.COMPLETED - ) + # Execution stays terminal; the worker tears down on COMPLETED, so + # the gate has no live worker and reads PRE_INVOKE. + assert store.load(arn).is_complete + assert executor._invocation_gate(arn) is InvocationState.PRE_INVOKE # noqa: SLF001 def test_is_current_token_returns_false_for_none_token(): @@ -738,7 +736,7 @@ def test_is_current_token_returns_false_for_none_token(): execution.start() store.save(execution) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 with pytest.raises( InvalidParameterValueException, match="Invalid checkpoint token" @@ -766,7 +764,7 @@ def test_is_current_token_returns_false_for_malformed_token(): execution.start() store.save(execution) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 with pytest.raises( InvalidParameterValueException, match="Invalid checkpoint token" @@ -936,7 +934,7 @@ def test_delta_truncation_advances_watermark_partially(): execution.start() store.save(execution) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 token = CheckpointToken(execution_arn=arn, token_sequence=0).to_str() diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py index 65887778..855ff9ab 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py @@ -14,12 +14,10 @@ CallbackOptions, OperationUpdate, OperationAction, - OperationSubType, OperationType, Operation, OperationStatus, CallbackDetails, - StepDetails, ) from aws_durable_execution_sdk_python.lambda_service import ( ErrorObject, @@ -38,7 +36,6 @@ from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState from aws_durable_execution_sdk_python_testing.invoker import InvokeResponse from aws_durable_execution_sdk_python_testing.model import ( - InvocationCompletedDetails, ListDurableExecutionsResponse, SendDurableExecutionCallbackFailureResponse, SendDurableExecutionCallbackHeartbeatResponse, @@ -47,7 +44,6 @@ StopDurableExecutionResponse, ) from aws_durable_execution_sdk_python_testing.observer import ( - ExecutionNotifier, ExecutionObserver, ) from aws_durable_execution_sdk_python_testing.stores.memory import ( @@ -265,7 +261,7 @@ def test_start_execution_with_provided_invocation_id( result = executor.get_execution("test-arn") - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") assert result == mock_execution @@ -863,7 +859,7 @@ def test_should_complete_workflow_successfully_through_public_api( executor.complete_execution("test-arn", "result") # Assert - Verify final execution status and stored results - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_execution.complete_success.assert_called_once_with(result="result") mock_store.update.assert_called_once_with(mock_execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") @@ -883,7 +879,7 @@ def test_should_complete_workflow_with_failure_through_public_api( executor.fail_execution("test-arn", error) # Assert - Verify final execution status and stored error - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_execution.complete_fail.assert_called_once_with(error=error) mock_store.update.assert_called_once_with(mock_execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") @@ -902,7 +898,7 @@ def test_should_handle_workflow_completion_state_through_public_api( executor.complete_execution("test-arn", "result") # Assert - Verify completion was processed and observer notifications sent - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_execution.complete_success.assert_called_once_with(result="result") mock_store.update.assert_called_once_with(mock_execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") @@ -1252,7 +1248,7 @@ def test_complete_execution(executor, mock_store, mock_execution): with patch.object(executor, "_complete_events") as mock_complete_events: executor.complete_execution("test-arn", "result") - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_execution.complete_success.assert_called_once_with(result="result") mock_store.update.assert_called_once_with(mock_execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") @@ -1266,7 +1262,7 @@ def test_fail_execution(executor, mock_store, mock_execution): with patch.object(executor, "_complete_events") as mock_complete_events: executor.fail_execution("test-arn", error) - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_execution.complete_fail.assert_called_once_with(error=error) mock_store.update.assert_called_once_with(mock_execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") @@ -1571,7 +1567,7 @@ def test_get_execution_details(executor, mock_store): assert result.status == "SUCCEEDED" assert result.result == "test-result" assert result.error is None - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") def test_get_execution_details_not_found(executor, mock_store): @@ -1784,7 +1780,7 @@ def test_stop_execution(executor, mock_store): result = executor.stop_execution("test-arn") - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") mock_store.update.assert_called_once_with(execution) assert result.stop_timestamp is not None assert execution.is_complete is True @@ -1829,7 +1825,7 @@ def test_stop_execution_with_custom_error(executor, mock_store): executor.stop_execution("test-arn", error=custom_error) - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") mock_store.update.assert_called_once_with(execution) assert execution.is_complete is True assert execution.close_status == ExecutionStatus.STOPPED @@ -1889,7 +1885,7 @@ def test_get_execution_state(mock_scheduler, mock_invoker, mock_checkpoint_proce store.update(execution) # Gate open + valid token. - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 token = CheckpointToken(execution_arn=arn, token_sequence=0).to_str() result = executor.get_execution_state(arn, checkpoint_token=token) @@ -1921,7 +1917,7 @@ def test_get_execution_state_invalid_token( execution.start() store.save(execution) arn = execution.durable_execution_arn - executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 with pytest.raises( InvalidParameterValueException, match="Invalid checkpoint token" @@ -1946,11 +1942,12 @@ def test_get_execution_history(executor, mock_store): assert result.events == [] assert result.next_marker is None - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") def test_get_execution_history_with_events(executor, mock_store): """Test get_execution_history with actual events.""" + from aws_durable_execution_sdk_python.lambda_service import StepDetails # Create operations that will generate events op1 = Operation( @@ -1979,53 +1976,6 @@ def test_get_execution_history_with_events(executor, mock_store): assert result.events[1].event_type == "StepSucceeded" -def test_get_execution_history_with_invocation_completions_and_updates( - executor, mock_store -): - """Exercise the update-history and invocation-completions branches.""" - now: datetime = datetime.now(UTC) - op: Operation = Operation( - operation_id="op-1", - operation_type=OperationType.STEP, - status=OperationStatus.SUCCEEDED, - start_timestamp=now, - end_timestamp=now, - sub_type=OperationSubType.STEP, - step_details=StepDetails(result='"ok"', attempt=1), - ) - update: OperationUpdate = OperationUpdate( - operation_id="op-1", - operation_type=OperationType.STEP, - action=OperationAction.START, - sub_type=OperationSubType.STEP, - name="my-step", - ) - completion: InvocationCompletedDetails = InvocationCompletedDetails( - start_timestamp=now, - end_timestamp=now, - request_id="req-1", - ) - - mock_execution = Mock() - mock_execution.operations = [op] - mock_execution.updates = [update] - mock_execution.update_timestamps = [now] - mock_execution.invocation_completions = [completion] - mock_execution.durable_execution_arn = "arn:test" - mock_execution.start_input = Mock() - mock_execution.start_input.execution_timeout_seconds = 60 - mock_execution.result = None - mock_store.load.return_value = mock_execution - - result = executor.get_execution_history("arn:test", include_execution_data=True) - - # At minimum we expect: ExecutionStarted, the step event(s) from - # the update, and an InvocationCompleted event from the completion. - event_types = [e.event_type for e in result.events] - assert "InvocationCompleted" in event_types - assert "StepStarted" in event_types - - def test_get_execution_history_reverse_order(executor, mock_store): """Test get_execution_history with reverse order.""" op1 = Operation( @@ -2163,7 +2113,7 @@ def test_checkpoint_execution(mock_scheduler, mock_invoker, mock_checkpoint_proc # Simulate an active invocation so the invocation-state gate lets # the checkpoint through. _invoke_handler wires this # automatically. - real_executor._invocation_state[arn] = InvocationState.INVOKING # noqa: SLF001 + real_executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 initial_token = CheckpointToken(execution_arn=arn, token_sequence=0).to_str() @@ -2229,7 +2179,7 @@ def test_send_callback_success(executor, mock_store): result = executor.send_callback_success(callback_id, b"success-result") assert isinstance(result, SendDurableExecutionCallbackSuccessResponse) - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") mock_execution.complete_callback_success.assert_called_once_with( callback_id, b"success-result" ) @@ -2293,7 +2243,7 @@ def test_send_callback_failure(executor, mock_store): result = executor.send_callback_failure(callback_id) assert isinstance(result, SendDurableExecutionCallbackFailureResponse) - mock_store.load.assert_called_once_with("test-arn") + mock_store.load.assert_any_call("test-arn") mock_store.update.assert_called_once_with(mock_execution) # Verify execution is invoked after callback failure mock_invoke.assert_called_once_with("test-arn") @@ -2353,7 +2303,9 @@ def test_send_callback_heartbeat(executor, mock_store): assert isinstance(result, SendDurableExecutionCallbackHeartbeatResponse) # Called twice: once in get_execution, once in _reset_callback_heartbeat_timeout - assert mock_store.load.call_count == 2 + # get_execution + _reset_callback_heartbeat_timeout, plus the + # worker's completion-reload after the lane task. + assert mock_store.load.call_count >= 2 mock_execution.find_callback_operation.assert_called_once_with(callback_id) @@ -2716,7 +2668,7 @@ def test_on_timed_out(executor, mock_store): with patch.object(executor, "_complete_events") as mock_complete_events: executor.on_timed_out("test-arn", error) - mock_store.load.assert_called_once_with(execution_arn="test-arn") + mock_store.load.assert_any_call(execution_arn="test-arn") mock_store.update.assert_called_once_with(execution) mock_complete_events.assert_called_once_with(execution_arn="test-arn") assert execution.is_complete is True @@ -2732,27 +2684,3 @@ def test_on_stopped(executor): executor.on_stopped("test-arn", error) mock_fail.assert_called_once_with("test-arn", error) - - -def test_notify_timed_out(): - """Test notify_timed_out method.""" - notifier = ExecutionNotifier() - observer = Mock() - notifier.add_observer(observer) - - error = ErrorObject.from_message("Timeout error") - notifier.notify_timed_out("test-arn", error) - - observer.on_timed_out.assert_called_once_with(execution_arn="test-arn", error=error) - - -def test_notify_stopped(): - """Test notify_stopped method.""" - notifier = ExecutionNotifier() - observer = Mock() - notifier.add_observer(observer) - - error = ErrorObject.from_message("Stop error") - notifier.notify_stopped("test-arn", error) - - observer.on_stopped.assert_called_once_with(execution_arn="test-arn", error=error) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/observer_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/observer_test.py index 8eb40127..6c95de24 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/observer_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/observer_test.py @@ -1,15 +1,19 @@ -"""Tests for observer module.""" +"""Tests for the observer module: effect collection and application.""" import inspect -import threading -from unittest.mock import Mock import pytest -from aws_durable_execution_sdk_python.lambda_service import ErrorObject, CallbackOptions +from aws_durable_execution_sdk_python.lambda_service import CallbackOptions, ErrorObject +from aws_durable_execution_sdk_python_testing.checkpoint.effects import ( + CallbackCreated, + Completed, + Failed, +) from aws_durable_execution_sdk_python_testing.observer import ( ExecutionNotifier, ExecutionObserver, + apply_effects, ) from aws_durable_execution_sdk_python_testing.token import CallbackToken @@ -48,249 +52,155 @@ def on_callback_created( ) -def test_execution_notifier_init(): - """Test ExecutionNotifier initialization.""" - notifier = ExecutionNotifier() - - assert notifier._observers == [] # noqa: SLF001 - assert notifier._lock is not None # noqa: SLF001 - +# region ExecutionNotifier collects effects -def test_execution_notifier_add_observer(): - """Test adding an observer to ExecutionNotifier.""" - notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) +def test_execution_notifier_init_has_no_effects(): + assert ExecutionNotifier().effects == [] - assert len(notifier._observers) == 1 # noqa: SLF001 - assert notifier._observers[0] is observer # noqa: SLF001 - -def test_execution_notifier_add_multiple_observers(): - """Test adding multiple observers to ExecutionNotifier.""" +def test_notify_completed_records_completed_effect(): notifier = ExecutionNotifier() - observer1 = MockExecutionObserver() - observer2 = MockExecutionObserver() - - notifier.add_observer(observer1) - notifier.add_observer(observer2) + notifier.notify_completed("test-arn", "test-result") + assert notifier.effects == [ + Completed(execution_arn="test-arn", result="test-result") + ] - assert len(notifier._observers) == 2 # noqa: SLF001 - assert observer1 in notifier._observers # noqa: SLF001 - assert observer2 in notifier._observers # noqa: SLF001 - -def test_execution_notifier_notify_completed(): - """Test notifying observers about execution completion.""" +def test_notify_completed_no_result_records_none_result(): notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) + notifier.notify_completed("test-arn") + assert notifier.effects == [Completed(execution_arn="test-arn", result=None)] - execution_arn = "test-arn" - result = "test-result" - notifier.notify_completed(execution_arn, result) - - assert len(observer.on_completed_calls) == 1 - assert observer.on_completed_calls[0] == (execution_arn, result) - - -def test_execution_notifier_notify_completed_no_result(): - """Test notifying observers about execution completion with no result.""" +def test_notify_failed_records_failed_effect(): notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) + error = ErrorObject("TestError", "Test error message", "test-data", ["trace"]) + notifier.notify_failed("test-arn", error) + assert notifier.effects == [Failed(execution_arn="test-arn", error=error)] - execution_arn = "test-arn" - notifier.notify_completed(execution_arn) - - assert len(observer.on_completed_calls) == 1 - assert observer.on_completed_calls[0] == (execution_arn, None) - - -def test_execution_notifier_notify_failed(): - """Test notifying observers about execution failure.""" +def test_notify_callback_created_records_callback_effect(): notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) - - execution_arn = "test-arn" - error = ErrorObject( - "TestError", "Test error message", "test-data", ["stack", "trace"] + token = CallbackToken(execution_arn="test-arn", operation_id="op-1") + options = CallbackOptions() + notifier.notify_callback_created( + execution_arn="test-arn", + operation_id="op-1", + callback_options=options, + callback_token=token, ) + assert notifier.effects == [ + CallbackCreated( + execution_arn="test-arn", + operation_id="op-1", + callback_options=options, + callback_token=token, + ) + ] - notifier.notify_failed(execution_arn, error) - assert len(observer.on_failed_calls) == 1 - assert observer.on_failed_calls[0] == (execution_arn, error) +def test_notify_accumulates_in_order(): + notifier = ExecutionNotifier() + error = ErrorObject.from_message("boom") + notifier.notify_completed("arn-1", "r1") + notifier.notify_failed("arn-2", error) + assert notifier.effects == [ + Completed(execution_arn="arn-1", result="r1"), + Failed(execution_arn="arn-2", error=error), + ] -def test_execution_notifier_multiple_observers_all_notified(): - """Test that all observers are notified when multiple are registered.""" - notifier = ExecutionNotifier() - observer1 = MockExecutionObserver() - observer2 = MockExecutionObserver() +# endregion ExecutionNotifier collects effects - notifier.add_observer(observer1) - notifier.add_observer(observer2) +# region apply_effects drives effects onto an observer - execution_arn = "test-arn" - result = "test-result" - notifier.notify_completed(execution_arn, result) +def test_apply_effects_dispatches_completed(): + observer = MockExecutionObserver() + apply_effects([Completed(execution_arn="arn", result="ok")], observer) + assert observer.on_completed_calls == [("arn", "ok")] - # Both observers should be notified - assert len(observer1.on_completed_calls) == 1 - assert observer1.on_completed_calls[0] == (execution_arn, result) - assert len(observer2.on_completed_calls) == 1 - assert observer2.on_completed_calls[0] == (execution_arn, result) +def test_apply_effects_dispatches_failed(): + observer = MockExecutionObserver() + error = ErrorObject.from_message("nope") + apply_effects([Failed(execution_arn="arn", error=error)], observer) + assert observer.on_failed_calls == [("arn", error)] -def test_execution_notifier_no_observers(): - """Test that notifications work even with no observers.""" - notifier = ExecutionNotifier() - # Should not raise any exceptions - notifier.notify_completed("test-arn", "result") - notifier.notify_failed( - "test-arn", ErrorObject("Error", "Message", "data", ["trace"]) +def test_apply_effects_dispatches_callback_created(): + observer = MockExecutionObserver() + token = CallbackToken(execution_arn="arn", operation_id="op-1") + options = CallbackOptions() + apply_effects( + [ + CallbackCreated( + execution_arn="arn", + operation_id="op-1", + callback_options=options, + callback_token=token, + ) + ], + observer, ) + assert observer.on_callback_created_calls == [("arn", "op-1", options, token)] -def test_execution_notifier_thread_safety(): - """Test that ExecutionNotifier is thread-safe.""" - notifier = ExecutionNotifier() +def test_apply_effects_preserves_order_across_observer(): observer = MockExecutionObserver() - notifier.add_observer(observer) - - # Test concurrent access - def add_observer_thread(): - new_observer = MockExecutionObserver() - notifier.add_observer(new_observer) + error = ErrorObject.from_message("boom") + apply_effects( + [ + Completed(execution_arn="arn-1", result="r1"), + Failed(execution_arn="arn-2", error=error), + ], + observer, + ) + assert observer.on_completed_calls == [("arn-1", "r1")] + assert observer.on_failed_calls == [("arn-2", error)] - def notify_thread(): - notifier.notify_completed("test-arn", "result") - threads = [] - for _ in range(5): - threads.append(threading.Thread(target=add_observer_thread)) - threads.append(threading.Thread(target=notify_thread)) +def test_apply_effects_empty_is_a_noop(): + observer = MockExecutionObserver() + apply_effects([], observer) + assert observer.on_completed_calls == [] + assert observer.on_failed_calls == [] + assert observer.on_callback_created_calls == [] - for thread in threads: - thread.start() - for thread in threads: - thread.join() +# endregion apply_effects - # Should have original observer plus 5 more - assert len(notifier._observers) == 6 # noqa: SLF001 - # Original observer should have been notified multiple times - assert len(observer.on_completed_calls) >= 1 +# region ExecutionObserver interface -def test_execution_observer_abstract_methods(): - """Test that ExecutionObserver is abstract and cannot be instantiated.""" +def test_execution_observer_is_abstract(): with pytest.raises(TypeError): ExecutionObserver() -def test_mock_execution_observer_implementation(): - """Test that MockExecutionObserver properly implements all abstract methods.""" +def test_mock_execution_observer_implements_all_methods(): observer = MockExecutionObserver() - - # Test all methods can be called error = ErrorObject("Error", "Message", "data", ["trace"]) observer.on_completed("arn", "result") observer.on_failed("arn", error) observer.on_timed_out("arn", error) observer.on_stopped("arn", error) - # Verify calls were recorded assert len(observer.on_completed_calls) == 1 assert len(observer.on_failed_calls) == 1 assert len(observer.on_timed_out_calls) == 1 assert len(observer.on_stopped_calls) == 1 -def test_execution_notifier_notify_observers_with_exception(): - """Test that exceptions in one observer don't affect others.""" - notifier = ExecutionNotifier() - - # Create a mock observer that raises an exception - failing_observer = Mock(spec=ExecutionObserver) - failing_observer.on_completed.side_effect = ValueError("Test exception") - - # Create a normal observer - normal_observer = MockExecutionObserver() - - notifier.add_observer(failing_observer) - notifier.add_observer(normal_observer) - - # This should raise an exception from the failing observer - with pytest.raises(ValueError, match="Test exception"): - notifier.notify_completed("test-arn", "result") - - # The normal observer should still have been called before the exception - failing_observer.on_completed.assert_called_once_with( - execution_arn="test-arn", result="result" - ) - - -def test_execution_observer_abstract_method_coverage(): - """Test coverage of abstract methods in ExecutionObserver.""" - # This test ensures we cover the abstract method definitions - # by checking they exist and have the correct signatures - +def test_execution_observer_declares_expected_methods(): methods = inspect.getmembers(ExecutionObserver, predicate=inspect.isfunction) method_names = [name for name, _ in methods] - assert "on_completed" in method_names assert "on_failed" in method_names assert "on_timed_out" in method_names assert "on_stopped" in method_names - # on_wait_timer_scheduled and on_step_retry_scheduled were - # removed in — timer scheduling no longer dispatched - # via the observer pattern. assert "on_callback_created" in method_names -def test_execution_notifier_notify_observers_internal(): - """Test the internal _notify_observers method behavior.""" - notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) - - # Test that _notify_observers correctly calls the method on observers - notifier._notify_observers( # noqa: SLF001 - ExecutionObserver.on_completed, execution_arn="test", result="success" - ) - - assert len(observer.on_completed_calls) == 1 - assert observer.on_completed_calls[0] == ("test", "success") - - -def test_execution_notifier_all_notification_methods(): - """Test all notification methods with various parameter combinations.""" - notifier = ExecutionNotifier() - observer = MockExecutionObserver() - notifier.add_observer(observer) - - # Test notify_completed with positional args - notifier.notify_completed("arn1", "result1") - assert observer.on_completed_calls[-1] == ("arn1", "result1") - - # Test notify_completed with keyword args - notifier.notify_completed(execution_arn="arn2", result="result2") - assert observer.on_completed_calls[-1] == ("arn2", "result2") - - # Test notify_failed - error = ErrorObject("TestError", "Message", "data", ["trace"]) - notifier.notify_failed("arn3", error) - assert observer.on_failed_calls[-1] == ("arn3", error) - - # notify_wait_timer_scheduled and notify_step_retry_scheduled - # were removed in () — - # timer scheduling centralised in - # Executor._schedule_earliest_pending. +# endregion ExecutionObserver interface diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/worker/__init__.py b/packages/aws-durable-execution-sdk-python-testing/tests/worker/__init__.py new file mode 100644 index 00000000..a7ffdfd9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/worker/__init__.py @@ -0,0 +1 @@ +"""Worker actor-model tests.""" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/worker/checkpoint_tasks_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/worker/checkpoint_tasks_test.py new file mode 100644 index 00000000..82d73f4c --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/worker/checkpoint_tasks_test.py @@ -0,0 +1,122 @@ +"""Tests for checkpoint/get-state worker tasks and serialized client calls.""" + +from __future__ import annotations + +import threading +import time +from typing import TYPE_CHECKING, cast +from unittest.mock import Mock + +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + StateOutput, +) + +from aws_durable_execution_sdk_python_testing.client import InMemoryServiceClient +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( + CheckpointProcessor, + ) + from aws_durable_execution_sdk_python_testing.scheduler import Scheduler + from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore + +ARN = "arn:aws:lambda:us-west-2:123456789012:function:fn/durable-execution/exec/inv" + + +class _FakeExecution: + def __init__(self, *, is_complete: bool = False) -> None: + self.is_complete = is_complete + + +class _FakeStore: + def __init__(self, execution: _FakeExecution) -> None: + self._execution = execution + + def load(self, execution_arn: str) -> _FakeExecution: # noqa: ARG002 + return self._execution + + +def _client(processor: Mock, execution: _FakeExecution) -> InMemoryServiceClient: + registry = ExecutionRegistry( + cast("ExecutionStore", _FakeStore(execution)), cast("Scheduler", object()) + ) + return InMemoryServiceClient(cast("CheckpointProcessor", processor), registry) + + +def test_checkpoint_delegates_and_returns_output(): + execution = _FakeExecution(is_complete=False) + processor = Mock() + output = CheckpointOutput(checkpoint_token="t2", new_execution_state=Mock()) # noqa: S106 + processor.process_checkpoint.return_value = output + + client = _client(processor, execution) + result = client.checkpoint(ARN, "t1", [], "ct") + + assert result is output + processor.process_checkpoint.assert_called_once_with("t1", [], "ct") + + +def test_get_state_delegates_and_returns_output(): + execution = _FakeExecution(is_complete=False) + processor = Mock() + output = StateOutput(operations=[], next_marker=None) + processor.get_execution_state.return_value = output + + client = _client(processor, execution) + result = client.get_execution_state(ARN, "t1", "marker", 50) + + assert result is output + processor.get_execution_state.assert_called_once_with("t1", "marker", 50) + + +def test_completing_execution_tears_down_its_worker(): + execution = _FakeExecution(is_complete=True) + processor = Mock() + processor.process_checkpoint.return_value = CheckpointOutput( + checkpoint_token="t2", # noqa: S106 + new_execution_state=Mock(), + ) + registry = ExecutionRegistry( + cast("ExecutionStore", _FakeStore(execution)), cast("Scheduler", object()) + ) + client = InMemoryServiceClient(cast("CheckpointProcessor", processor), registry) + + client.checkpoint(ARN, "t1", [], None) + + assert registry.get(ARN) is None + + +def test_concurrent_checkpoints_same_arn_are_serialized(): + execution = _FakeExecution(is_complete=False) + active = 0 + max_active = 0 + guard = threading.Lock() + + def process_checkpoint(token, updates, client_token): # noqa: ANN001, ARG001 + nonlocal active, max_active + with guard: + active += 1 + max_active = max(max_active, active) + time.sleep(0.01) + with guard: + active -= 1 + return CheckpointOutput(checkpoint_token="t", new_execution_state=Mock()) # noqa: S106 + + processor = Mock() + processor.process_checkpoint.side_effect = process_checkpoint + + client = _client(processor, execution) + + def call() -> None: + client.checkpoint(ARN, "t1", [], None) + + threads = [threading.Thread(target=call) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert max_active == 1 + assert processor.process_checkpoint.call_count == 8 diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/worker/concurrency_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/worker/concurrency_test.py new file mode 100644 index 00000000..c3b573d6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/worker/concurrency_test.py @@ -0,0 +1,152 @@ +"""Concurrency invariants of the per-execution worker model. + +These prove the properties the design relies on: a completing task tears +its own lane down without deadlocking, one execution's tasks never +overlap, and different executions run in parallel. +""" + +from __future__ import annotations + +import threading +import time +from typing import TYPE_CHECKING, cast + +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry +from aws_durable_execution_sdk_python_testing.worker.status import InvocationState +from aws_durable_execution_sdk_python_testing.worker.task import ( + ExecutionTask, + TaskOutcome, +) + +if TYPE_CHECKING: + 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 WorkerContext + + +class _FakeExecution: + def __init__(self, *, is_complete: bool = False) -> None: + self.is_complete = is_complete + + +class _FakeStore: + def __init__(self, execution: _FakeExecution) -> None: + self._execution = execution + + def load(self, execution_arn: str) -> _FakeExecution: # noqa: ARG002 + return self._execution + + +def _registry(execution: _FakeExecution) -> ExecutionRegistry: + return ExecutionRegistry( + cast("ExecutionStore", _FakeStore(execution)), cast("Scheduler", object()) + ) + + +class _CompletingTask(ExecutionTask[str]): + """Marks the execution complete and returns COMPLETED, which makes + the worker tear down its own lane from inside the running task.""" + + def __init__(self, execution: _FakeExecution) -> None: + self._execution = execution + + def execute( + self, + status: InvocationState, # noqa: ARG002 + worker: WorkerContext, # noqa: ARG002 + ) -> TaskOutcome[str]: + self._execution.is_complete = True + return TaskOutcome(InvocationState.COMPLETED, "done") + + +def test_completing_task_tears_down_without_deadlock() -> None: + execution = _FakeExecution(is_complete=False) + registry = _registry(execution) + worker = registry.get_or_create("arn-complete") + + future = worker.submit(_CompletingTask(execution)) + + # A self-join during teardown would hang here; the bounded wait is + # the deadlock guard. + assert future.result(timeout=5) == "done" + # The worker removed itself on completion. + assert registry.get("arn-complete") is None + + +class _OverlapTask(ExecutionTask[None]): + """Records the peak number of concurrently-running instances.""" + + def __init__(self, state: _Concurrency) -> None: + self._state = state + + def execute( + self, + status: InvocationState, + worker: WorkerContext, # noqa: ARG002 + ) -> TaskOutcome[None]: + self._state.enter() + time.sleep(0.01) + self._state.leave() + return TaskOutcome(status, None) + + +class _Concurrency: + def __init__(self) -> None: + self._lock = threading.Lock() + self.active = 0 + self.peak = 0 + + def enter(self) -> None: + with self._lock: + self.active += 1 + self.peak = max(self.peak, self.active) + + def leave(self) -> None: + with self._lock: + self.active -= 1 + + +def test_same_execution_tasks_never_overlap() -> None: + execution = _FakeExecution(is_complete=False) + registry = _registry(execution) + worker = registry.get_or_create("arn-serial") + state = _Concurrency() + + futures = [worker.submit(_OverlapTask(state)) for _ in range(10)] + for f in futures: + f.result(timeout=5) + + assert state.peak == 1 + + +class _BarrierTask(ExecutionTask[None]): + """Blocks on a shared barrier; only completes if every lane runs at + the same time, proving cross-execution parallelism.""" + + def __init__(self, barrier: threading.Barrier) -> None: + self._barrier = barrier + + def execute( + self, + status: InvocationState, + worker: WorkerContext, # noqa: ARG002 + ) -> TaskOutcome[None]: + self._barrier.wait(timeout=5) + return TaskOutcome(status, None) + + +def test_different_executions_run_in_parallel() -> None: + execution = _FakeExecution(is_complete=False) + registry = _registry(execution) + parties = 4 + barrier = threading.Barrier(parties) + + # Distinct ARNs -> distinct workers -> distinct lanes. If the lanes + # did not run in parallel the barrier would never trip and result() + # would raise BrokenBarrierError. + futures = [ + registry.get_or_create(f"arn-{i}").submit(_BarrierTask(barrier)) + for i in range(parties) + ] + for f in futures: + f.result(timeout=10) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/worker/lane_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/worker/lane_test.py new file mode 100644 index 00000000..dea70070 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/worker/lane_test.py @@ -0,0 +1,155 @@ +"""Unit tests for :class:`SerialTaskLane`.""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable + +import pytest + +from aws_durable_execution_sdk_python_testing.worker.lane import SerialTaskLane + + +def _appender(target: list[int], value: int) -> Callable[[], None]: + """Return a zero-arg task that appends ``value`` to ``target``.""" + return lambda: target.append(value) + + +def _double(value: int) -> Callable[[], int]: + """Return a zero-arg task that yields ``value * 2``.""" + return lambda: value * 2 + + +def test_runs_tasks_in_submit_order() -> None: + lane = SerialTaskLane.create() + try: + order: list[int] = [] + futures = [lane.submit(_appender(order, i)) for i in range(20)] + for f in futures: + f.result(timeout=5) + assert order == list(range(20)) + finally: + lane.stop() + + +def test_runs_one_at_a_time() -> None: + lane = SerialTaskLane.create() + try: + active = 0 + max_active = 0 + guard = threading.Lock() + + def task() -> None: + nonlocal active, max_active + with guard: + active += 1 + max_active = max(max_active, active) + time.sleep(0.01) + with guard: + active -= 1 + + futures = [lane.submit(task) for _ in range(10)] + for f in futures: + f.result(timeout=5) + assert max_active == 1 + finally: + lane.stop() + + +def test_exception_does_not_break_lane() -> None: + lane = SerialTaskLane.create() + try: + + def boom() -> None: + raise ValueError("boom") + + failing = lane.submit(boom) + with pytest.raises(ValueError, match="boom"): + failing.result(timeout=5) + + # The lane keeps serving after a task raises. + ok = lane.submit(lambda: 42) + assert ok.result(timeout=5) == 42 + finally: + lane.stop() + + +def test_concurrent_submitters_get_correct_results() -> None: + lane = SerialTaskLane.create() + try: + results: dict[int, int] = {} + results_lock = threading.Lock() + start = threading.Barrier(8) + + def submit_many(base: int) -> None: + start.wait() + local = {n: lane.submit(_double(n)) for n in range(base, base + 25)} + for n, fut in local.items(): + with results_lock: + results[n] = fut.result(timeout=5) + + threads = [ + threading.Thread(target=submit_many, args=(base,)) + for base in range(0, 200, 25) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert results == {n: n * 2 for n in range(200)} + finally: + lane.stop() + + +def test_stop_drains_queued_tasks_then_joins() -> None: + lane = SerialTaskLane.create() + done: list[int] = [] + futures = [lane.submit(_appender(done, i)) for i in range(50)] + + lane.stop() # must drain all queued work before the thread exits + + assert lane.is_running() is False + assert all(f.done() for f in futures) + assert done == list(range(50)) + + +def test_submit_after_stop_raises() -> None: + lane = SerialTaskLane.create() + lane.stop() + with pytest.raises(RuntimeError, match="stopped"): + lane.submit(lambda: None) + + +def test_stop_is_idempotent() -> None: + lane = SerialTaskLane.create() + lane.stop() + lane.stop() # second call is a no-op, does not raise + assert lane.is_running() is False + + +def test_cancelled_future_is_skipped() -> None: + lane = SerialTaskLane.create() + try: + started = threading.Event() + release = threading.Event() + + def blocker() -> None: + started.set() + release.wait(timeout=5) + + blocking = lane.submit(blocker) + started.wait(timeout=5) + + # Queue a second task and cancel it while the lane is still busy + # with the blocker, so the consumer skips it when it dequeues. + skipped = lane.submit(lambda: None) + assert skipped.cancel() + + release.set() + blocking.result(timeout=5) + + assert skipped.cancelled() + finally: + lane.stop() diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/worker/worker_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/worker/worker_test.py new file mode 100644 index 00000000..d918317d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/worker/worker_test.py @@ -0,0 +1,125 @@ +"""Unit tests for the execution worker, registry, and task contract.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import pytest + +from aws_durable_execution_sdk_python_testing.worker.registry import ExecutionRegistry +from aws_durable_execution_sdk_python_testing.worker.status import InvocationState +from aws_durable_execution_sdk_python_testing.worker.task import ( + ExecutionTask, + TaskOutcome, +) + +if TYPE_CHECKING: + 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 WorkerContext + +ARN = "arn:aws:lambda:us-west-2:123456789012:function:fn/durable-execution/exec/inv" + + +class _FakeStore: + """Minimal store stub: ``load`` returns a sentinel the tasks ignore.""" + + def load(self, execution_arn: str) -> object: # noqa: ARG002 + return object() + + +def _registry() -> ExecutionRegistry: + return ExecutionRegistry( + cast("ExecutionStore", _FakeStore()), cast("Scheduler", object()) + ) + + +class _ReturnTask(ExecutionTask[str]): + """Trivial task: returns a fixed value and next status.""" + + def __init__(self, next_status: InvocationState, value: str) -> None: + self._next = next_status + self._value = value + + def execute( + self, + status: InvocationState, # noqa: ARG002 + worker: WorkerContext, # noqa: ARG002 + ) -> TaskOutcome[str]: + return TaskOutcome(self._next, self._value) + + +class _RecordStatusTask(ExecutionTask[InvocationState]): + """Records the status it observed and returns it as the value.""" + + def __init__(self, next_status: InvocationState) -> None: + self._next = next_status + + def execute( + self, + status: InvocationState, + worker: WorkerContext, # noqa: ARG002 + ) -> TaskOutcome[InvocationState]: + return TaskOutcome(self._next, status) + + +def test_worker_runs_task_and_returns_value() -> None: + registry = _registry() + worker = registry.get_or_create(ARN) + future = worker.submit(_ReturnTask(InvocationState.PRE_INVOKE, "hello")) + assert future.result(timeout=5) == "hello" + + +def test_worker_threads_status_between_tasks() -> None: + registry = _registry() + worker = registry.get_or_create(ARN) + + # First task sees the initial PRE_INVOKE and advances to INVOKING. + first = worker.submit(_RecordStatusTask(InvocationState.INVOKING)) + # Second task must observe the INVOKING carried over from the first. + second = worker.submit(_RecordStatusTask(InvocationState.INVOKING)) + + assert first.result(timeout=5) is InvocationState.PRE_INVOKE + assert second.result(timeout=5) is InvocationState.INVOKING + assert worker.status is InvocationState.INVOKING + + +def test_worker_tears_down_on_completed() -> None: + registry = _registry() + worker = registry.get_or_create(ARN) + + future = worker.submit(_ReturnTask(InvocationState.COMPLETED, "done")) + assert future.result(timeout=5) == "done" + + # Completion removes the worker from the registry and stops its lane. + assert registry.get(ARN) is None + assert registry.active_count() == 0 + with pytest.raises(RuntimeError, match="stopped"): + worker.submit(_ReturnTask(InvocationState.PRE_INVOKE, "late")) + + +def test_registry_get_or_create_is_idempotent_per_arn() -> None: + registry = _registry() + first = registry.get_or_create(ARN) + second = registry.get_or_create(ARN) + assert first is second + assert registry.get(ARN) is first + assert registry.active_count() == 1 + + +def test_registry_distinct_workers_per_arn() -> None: + registry = _registry() + worker_a = registry.get_or_create(ARN + "-a") + worker_b = registry.get_or_create(ARN + "-b") + assert worker_a is not worker_b + assert registry.active_count() == 2 + + +def test_registry_remove_drops_worker() -> None: + registry = _registry() + registry.get_or_create(ARN) + registry.remove(ARN) + assert registry.get(ARN) is None + assert registry.active_count() == 0 + # Removing an unknown ARN is a no-op. + registry.remove(ARN)