From cfaae2bdd77bc0bae9fafea2ec1483e652402a61 Mon Sep 17 00:00:00 2001 From: yaythomas Date: Sat, 24 Oct 2026 03:29:45 +0000 Subject: [PATCH] feat(sdk): run plugin hooks off checkpoint thread The end and change hooks for a checkpoint response ran on the checkpoint thread, the thread that applies every response for every branch, so a slow hook delayed every checkpoint. They now run elsewhere, in an order plugins can rely on. The checkpoint thread records each terminal operation's end hook before it publishes the response. The thread that then observes the result runs the hook before the operation returns: the operation's own thread after a synchronous checkpoint, or the thread that reads a result another response completed, such as a branch resumed after a refresh. Another thread that observes the result meanwhile waits for the hook; a read made from inside any plugin hook neither runs nor waits for an end hook, so hooks cannot wait on each other. So on a logical path an end hook finishes before the next operation starts, which the OpenTelemetry plugin needs for its span order. End hooks no thread claims run on the handler thread before on_invocation_end, descendants before ancestors, and an operation ends at most once per invocation. on_operation_change runs on a dispatcher thread. Notifications that arrive while a plugin is still in the hook merge into the pending notification: updated_operations carries the newest state of every operation changed since the previous delivery, operations the newest map. So the backlog is at most one pending notification. The plugin contract docstring states where each hook runs and what order it may assume. --- .../execution_plugin.py | 6 +- .../invocation_plugin.py | 6 +- .../log_filter.py | 9 +- .../plugin.py | 7 +- .../context.py | 18 +- .../plugin.py | 780 +++++++- .../aws_durable_execution_sdk_python/state.py | 105 +- .../tests/context_test.py | 2 +- .../handler_context_propagation_int_test.py | 61 +- .../e2e/plugin_hook_contract_int_test.py | 383 ++++ .../tests/plugin_hook_registry_test.py | 1704 +++++++++++++++++ 11 files changed, 2978 insertions(+), 103 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/plugin_hook_contract_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/plugin_hook_registry_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 2cc841a35..64f55caf3 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -805,9 +805,9 @@ def _release_invocation_scope(self) -> None: on such a thread is still correlated to this finished invocation's spans. * ``_tracing_enabled`` is cleared so a hook that arrives after the - invocation end -- one dispatched off the checkpointing path, for - instance -- cannot start a span after the invocation span was ended and - the provider flushed. + invocation end cannot start a span after the invocation span was ended + and the provider flushed. The SDK dispatches no operation hook after + the end hook; this guard covers a caller that does. """ self._detach_remaining_contexts() unbind_invocation(self) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index d6ab7f792..a861c76b6 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -893,9 +893,9 @@ def _release_invocation_scope(self) -> None: on such a thread is still correlated to this finished invocation's spans. * ``_tracing_enabled`` is cleared so a hook that arrives after the - invocation end -- one dispatched off the checkpointing path, for - instance -- cannot start a span after the invocation span was ended and - the provider flushed. + invocation end cannot start a span after the invocation span was ended + and the provider flushed. The SDK dispatches no operation hook after + the end hook; this guard covers a caller that does. """ self._detach_remaining_contexts() unbind_invocation(self) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py index 40e55414b..48f3e764b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py @@ -41,10 +41,11 @@ coordinator's Context, carrying the invocation claim. Plugins also reaffirm the claim from hooks that run before user code on that thread. - A record emitted on a thread that carries no live claim is left unstamped, - whatever the number of open invocations. Two threads carry no claim: the - SDK's background checkpointing thread, and a thread customer code starts - itself, since Python does not copy context into a new thread. Correlation - is lost for those records. + whatever the number of open invocations. Three kinds of thread carry no + claim: the SDK's background checkpointing thread, the SDK's plugin change + thread, which runs ``on_operation_change``, and a thread customer code + starts itself, since Python does not copy context into a new thread. + Correlation is lost for those records. Resolving an unclaimed record against the single open invocation, when exactly one is open, would be wrong in two orderings. A thread claimed by invocation A diff --git a/packages/aws-durable-execution-sdk-python-telemetry-exporter/src/aws_durable_execution_sdk_python_telemetry_exporter/plugin.py b/packages/aws-durable-execution-sdk-python-telemetry-exporter/src/aws_durable_execution_sdk_python_telemetry_exporter/plugin.py index ba31546ad..41cb5f8cc 100644 --- a/packages/aws-durable-execution-sdk-python-telemetry-exporter/src/aws_durable_execution_sdk_python_telemetry_exporter/plugin.py +++ b/packages/aws-durable-execution-sdk-python-telemetry-exporter/src/aws_durable_execution_sdk_python_telemetry_exporter/plugin.py @@ -496,9 +496,10 @@ def __init__( # operation_id -> OperationInfo, adopted verbatim from the SDK's # authoritative snapshot (invocation start/end and operation-change). self._operations: dict[str, OperationInfo] = {} - # Set once this invocation has ended. A hook that arrives afterwards (an - # operation-change for a checkpoint that completed just before the end) - # must emit nothing (mirrors the Java ExecutionState.closed flag). + # Set once this invocation has ended. A hook that arrives afterwards + # must emit nothing (mirrors the Java ExecutionState.closed flag). The + # SDK delivers every pending operation-change before the end hook; the + # guard covers a caller that does not. self._closed = False # A duplicate End must wait for the admitted closing build before its # drain captures scheduler state. Nested callbacks defer that wait to diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py index afc816dda..9ac766eb5 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py @@ -509,8 +509,15 @@ def _peek_next_operation_id(self) -> str: ) def _peek_next_checkpoint(self) -> CheckpointedResult: - """Return the checkpoint for this context's next operation, without consuming it.""" - return self.state.get_checkpoint_result(self._peek_next_operation_id()) + """Return the checkpoint for this context's next operation, without consuming it. + + The read takes no plugin end hook: the hook of an operation a + checkpoint response completed during this invocation belongs to the + read that consumes the result, the operation's own. + """ + return self.state.get_checkpoint_result( + self._peek_next_operation_id(), claim_end_hook=False + ) def _next_operation_exists(self) -> bool: """True if a checkpoint exists for this context's next operation.""" @@ -548,9 +555,14 @@ def _replay_aware( # Only look up a checkpoint when replaying. Operation call sites pass # their atomically allocated identifier so concurrent callers inspect # the checkpoint for their own ID rather than peeking at shared state. + # Both reads inspect: neither consumes the result, so neither takes + # the plugin end hook a checkpoint response may have recorded for the + # operation. The operation's own read, which consumes the result, does. next_checkpoint: CheckpointedResult | None = ( ( - self.state.get_checkpoint_result(operation_identifier.operation_id) + self.state.get_checkpoint_result( + operation_identifier.operation_id, claim_end_hook=False + ) if operation_identifier is not None else self._peek_next_checkpoint() ) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index 33780eb23..fd72dae91 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -6,7 +6,8 @@ import functools import inspect import logging -from collections.abc import Iterator, Mapping, Sequence +import threading +from collections.abc import Collection, Iterator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, MutableMapping, Protocol, cast @@ -393,7 +394,106 @@ class PropagationMetadata: class DurableInstrumentationPlugin: - """Base class for plugins. Override only the methods you need.""" + """Base class for plugins. Override only the methods you need. + + **Every hook is a synchronous call on the thread named below, and that + thread waits for it to return.** No hook runs on the checkpoint thread, + the one thread that sends every checkpoint of the invocation and applies + every response. So a slow hook delays the thread it runs on, and never a + checkpoint of another thread. A plugin that needs slow work, such as a + network export, does it on its own thread and finishes it in + ``on_invocation_end``. + + **Where each hook runs.** + + - ``on_invocation_start``, ``on_invocation_end``: the handler thread, + before the handler starts, and after every operation thread has ended + and every other hook has returned. + - ``on_operation_start``, ``on_user_function_start``, + ``on_user_function_end``: the thread that runs the operation. + - ``on_operation_end``: the thread that observes the operation's terminal + result, before the operation returns to user code. For an operation this + invocation completes, that is the thread that ran it, after its terminal + checkpoint is confirmed. For an operation completed outside this + invocation (a wait that elapsed, a callback, a child execution), it is + the thread that replays the operation when the completion arrived before + this invocation, or the thread that reads the result when the completion + arrives in a checkpoint response during it: a branch resumed after a + refresh, a callback result read. A second thread that observes the same + result while the hook runs waits for it, so no thread acts on a result + whose end hook has not returned. The one exception is a read made from + inside any plugin hook, on the hook's own thread: it neither runs nor + waits for an end hook. So a result read cannot make the hook that + makes it wait for another end hook, and a lock a plugin holds across a + hook cannot be needed by an end hook that read starts. A child context + that writes no terminal checkpoint (a virtual context, or one in + ReplayChildren mode) gets it on its thread when its body returns. + - ``on_operation_end`` for a terminal operation no thread reads in this + invocation, such as a completion that arrives after the thread that + would have read it stopped: the handler thread, before + ``on_invocation_end``. A completed descendant no thread read, such as a + callback created and never awaited, runs with its ancestor's end hook, + before it, on the thread that observes the ancestor's result. Within one + invocation an operation's end hook runs at most once, whichever of + these threads runs it. + - ``on_operation_change``: the invocation's plugin change thread. + + **Order.** For one operation, on its thread: ``on_operation_start`` + returns before ``on_user_function_start`` is called; + ``on_user_function_end`` returns before the terminal checkpoint is sent; + ``on_operation_end`` returns before the operation returns to user code. So + on one thread an operation's end hook returns before the next operation's + start hook is called, and a child's end hook returns before its parent's + is called. Threads are not ordered with each other: parallel and map + branches run their hooks concurrently, so a plugin guards state shared + between hooks with a lock. ``on_operation_change`` is not ordered with the + end hooks of the same response. + + Exception to child before parent: a map or parallel with an early + completion policy completes while a branch still runs. A FLAT branch + whose body then returns runs its end hook on its own thread, after the + parent's. A branch that attempts another durable operation after the + parent completed is stopped there and gets no end hook; so is a NESTED + branch, whose terminal checkpoint is rejected. + + **Change notifications.** ``on_operation_change`` runs on the plugin + change thread, in response order. One call may cover several checkpoint + responses: while the plugins are still in the hook, the responses that + arrive merge into one pending notification. ``updated_operations`` then + holds the newest state of every operation whose status changed since the + previous call, and ``operations`` the newest full map. Intermediate + states and the number of calls are not guaranteed. The pending + notification is delivered before ``on_invocation_end``. The single pending + slot bounds the backlog: whatever a plugin's speed, one notification is + being delivered and at most one is pending. A slow plugin also delays the + plugins after it in the list for that call, because the thread runs the + plugins one at a time. If the process cannot start the thread, the SDK + logs a warning and delivers the one merged notification on the handler + thread before ``on_invocation_end``. + + **Repeats.** An operation whose result is replayed from its terminal + checkpoint gets one end hook, in the first invocation that sees the + terminal checkpoint, and is not reported in a later invocation that + replays it. The service decides which operations an invocation sees as + updated; a retried invocation sees the same ones again. A child context + whose body runs again on replay reports its end on every run: a + ReplayChildren context (its checkpoint holds a summary, not the result, so + the body re-runs) with ``is_replayed`` true; a virtual context (no + checkpoint at all) with ``is_replayed`` false, because it has no + checkpoint to derive it from. + + **Failures.** An exception a hook raises is logged and ignored, so a + plugin cannot fail the execution. ``KeyboardInterrupt``, ``SystemExit`` + and ``GeneratorExit`` are the exceptions. On an operation thread and on + the handler thread, the fallback delivery of a change notification + included, they propagate from the SDK call that ran the hook. One raised + by an end hook the handler thread runs before ``on_invocation_end`` does + not stop the other such hooks; ``on_invocation_end`` is dispatched to + every plugin before it re-raises. On the plugin change thread they are + logged, the other plugins still receive the notification, and the next + notification is still delivered; nothing is carried to another thread. + No hook runs while the SDK holds a lock. + """ def provide_propagation_metadata( self, @@ -427,8 +527,8 @@ def on_operation_start(self, info: OperationStartInfo) -> None: """ Called before an operation's START checkpoint is queued, or when a prior non-terminal operation is replayed. This guarantees that it - strictly precedes ``on_user_function_start``. This is called NOT within - the thread that runs operation. + strictly precedes ``on_user_function_start``. This is called within + the thread that runs the operation. Args: info: Information about the operation. @@ -438,9 +538,18 @@ def on_operation_start(self, info: OperationStartInfo) -> None: def on_operation_end(self, info: OperationEndInfo) -> None: """ - Called when an operation reaches a terminal status. Terminal operations - are not emitted again during replay. Child contexts without a terminal - checkpoint may emit this from the thread that runs the operation. + Called when an operation reaches a terminal status, or when a child + context without a terminal checkpoint returns. An operation with a + terminal checkpoint is not reported again on replay; a child context + without one is reported on every replay. + + Runs on the thread that observes the terminal result, before the + operation returns to user code: the thread that ran the operation, or + the thread that replays, or reads the result of, an operation + completed outside this invocation. So a slow hook delays that thread + only. A terminal operation no thread reads in this invocation gets it + on the handler thread before ``on_invocation_end``. See the class + docstring. Args: info: Information about the operation. @@ -449,8 +558,14 @@ def on_operation_end(self, info: OperationEndInfo) -> None: def on_operation_change(self, info: OperationChangeInfo) -> None: """ - Called when checkpointed operations change after a checkpoint response is merged. - This is called NOT within the thread that runs operation. + Called when checkpointed operations change after a checkpoint response + is merged. Runs on the invocation's plugin change thread, in response + order, never on the checkpoint thread, so it delays no checkpoint. One + call may cover several responses: ``updated_operations`` holds the + newest state of every operation changed since the previous call, and + ``operations`` the newest full map. It is not ordered with the + ``on_operation_end`` hooks of the same response. See the class + docstring. Args: info: Updated operations and the full operation map for the invocation. @@ -646,6 +761,181 @@ def _contain_plugin_failure( return control +class _RecordedEndHook: + """One recorded end hook and who may run it. + + ``done`` is None while no thread owns the hook. It is an event once a + thread owns it: the synchronous waiter the checkpoint loop reserved it + for (``reserved`` True until that waiter arrives), or the thread running + it. The event completes + when the hook has run, or when the entry is retired unrun after a + thread-control exception in an earlier hook of the same claim, so another + thread that observes the same result waits for it instead of running it + again or going on without it. + + Internal to the SDK. Not part of the public API. + """ + + __slots__ = ("done", "info", "reserved") + + def __init__(self, info: OperationEndInfo, *, reserved: bool = False) -> None: + self.info = info + self.reserved = reserved + self.done: threading.Event | None = threading.Event() if reserved else None + + +class _EndHookClaim: + """What a thread must do after it observed a terminal result: run or wait. + + Built under the registry lock by :meth:`PluginExecutor.take_end_hook`, + acted on by :meth:`run` outside every lock. ``to_run`` holds the entries + this thread now owns, descendants first; ``to_wait`` the event of a hook + another thread owns. Both empty means there is nothing to do. + + Internal to the SDK. Not part of the public API. + """ + + __slots__ = ("_executor", "to_run", "to_wait") + + def __init__( + self, + executor: PluginExecutor, + to_run: list[_RecordedEndHook], + to_wait: threading.Event | None, + ) -> None: + self._executor = executor + self.to_run = to_run + self.to_wait = to_wait + + def run(self) -> None: + if self.to_wait is not None: + self.to_wait.wait() + if self.to_run: + self._executor._run_claimed(self.to_run) # noqa: SLF001 - same module + + +def _merge_change_infos( + pending: OperationChangeInfo, latest: OperationChangeInfo +) -> OperationChangeInfo: + """Merge a newer change notification into the pending one. + + The result describes both checkpoint responses as one. ``updated_operations`` + holds every operation either notification reported, each in the state of + the newest full map, so a plugin sees the latest state and never an older + one after a newer. ``operations`` is the newest full map. The map only + grows, so it holds every operation the pending notification reported; the + pending state is used only for an operation the map lacks. + """ + updated_operations = { + operation_id: latest.operations.get(operation_id, info) + for reported in (pending.updated_operations, latest.updated_operations) + for operation_id, info in reported.items() + } + return OperationChangeInfo( + execution_arn=latest.execution_arn, + updated_operations=updated_operations, + operations=latest.operations, + ) + + +class _ChangeDispatcher: + """One invocation's thread for ``on_operation_change``, with one pending slot. + + The checkpoint thread submits one notification per response and runs none. + This thread delivers them in order. A submission that arrives while a + delivery is in progress merges into the single pending notification + (:func:`_merge_change_infos`), so a slow plugin holds at most two + notifications in memory, the one being delivered and the pending one, + rather than a growing queue. A thread-control exception a hook raises + here is logged, the other plugins still receive the notification, and the + next notification is still delivered: this thread has no caller to stop, + and the SDK has no channel that carries an exception from this thread to + the handler thread. + + Built by :meth:`start`, which also starts the thread, because starting a + thread can fail. An instance built directly has no running thread: + :meth:`join` then delivers the pending notification on the calling thread. + That is the fallback when the process cannot start a thread. + + Internal to the SDK. Not part of the public API. + """ + + def __init__(self, executor: PluginExecutor) -> None: + self._executor = executor + self._pending: OperationChangeInfo | None = None + self._stopping = False + self._condition = threading.Condition() + self._thread = threading.Thread( + target=self._run, name="plugin-change", daemon=True + ) + self._thread_started = False + + @classmethod + def start(cls, executor: PluginExecutor) -> _ChangeDispatcher: + """Create a dispatcher and start its thread. Raises RuntimeError when it cannot start.""" + dispatcher = cls(executor) + dispatcher._thread.start() + dispatcher._thread_started = True + return dispatcher + + def submit(self, info: OperationChangeInfo) -> None: + """Queue a notification; merge it into the pending one if there is one.""" + with self._condition: + self._pending = ( + info + if self._pending is None + else _merge_change_infos(self._pending, info) + ) + self._condition.notify() + + def join(self) -> None: + """Deliver what is pending, stop the thread, return once nothing is pending. + + Without a running thread, the pending notification is delivered here, + on the calling thread, and a thread-control exception a hook raises + propagates to the caller. + """ + with self._condition: + self._stopping = True + self._condition.notify() + if self._thread_started: + self._thread.join() + info = self._take() + if info is not None: + # On the calling thread, the handler's: a thread-control exception + # propagates here, as every hook on that thread does. + self._executor.execute_plugins(info, finish_fan_out=True) + + def _take(self) -> OperationChangeInfo | None: + with self._condition: + info, self._pending = self._pending, None + return info + + def _run(self) -> None: + while True: + with self._condition: + while self._pending is None and not self._stopping: + self._condition.wait() + info, self._pending = self._pending, None + if info is None: + # Stopping, and nothing left to deliver. + return + self._deliver(info) + + def _deliver(self, info: OperationChangeInfo) -> None: + try: + # finish_fan_out: the other plugins still receive this notification. + self._executor.execute_plugins(info, finish_fan_out=True) + except BaseException: # noqa: BLE001 - only thread control reaches here + # execute_plugins contains every other failure. Nothing on this + # thread can be stopped by the interrupt, so it is logged and the + # next notification is still delivered. + logger.exception( + "Plugin on_operation_change hook raised on the plugin change " + "thread; ignored" + ) + + class PluginExecutor: """One invocation's plugin instances, metadata and dispatch. @@ -677,6 +967,32 @@ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None self._run_entered = False + # End hooks recorded by the checkpoint thread and not yet finished, + # keyed by operation id, in recording order. The thread that observes + # an operation's terminal result claims and runs its hook; a second + # thread that observes the result meanwhile waits for that hook. What + # no thread claims runs on the handler thread before on_invocation_end. + self._end_hooks: dict[str, _RecordedEndHook] = {} + # Operations whose end hook this invocation has run or retired, by any + # path. The replay path reports a terminal operation the invocation + # input carried, and a checkpoint response can list that operation + # too; the set keeps the end hook to one per operation per invocation. + self._ended: set[str] = set() + # Held while recording end hooks and while deciding which thread runs + # one. A reader holds it across its read of the operation state and + # its decision, so no other thread's claim comes between the two. + self.claim_lock = threading.Lock() + # Per thread: how many plugin hooks the thread is inside right now. + # execute_plugins and provide_propagation_metadata, the two calls + # into a plugin's hook methods, count their call. A claim decision + # made by a thread inside a hook returns at once, so a read a plugin + # makes from any hook neither runs nor waits for an end hook. See + # take_end_hook. + self._hook_depth = threading.local() + # The thread for on_operation_change. Created by the first change + # notification, joined before on_invocation_end. + self._dispatcher: _ChangeDispatcher | None = None + self._dispatcher_lock = threading.Lock() def provide_propagation_metadata( self, info: PropagationInput @@ -703,7 +1019,8 @@ def provide_propagation_metadata( except Exception: pass try: - metadata = plugin.provide_propagation_metadata(info) + with self._in_hook_scope(): + metadata = plugin.provide_propagation_metadata(info) if metadata is None: continue if inspect.iscoroutine(metadata): @@ -780,14 +1097,296 @@ class exists to prevent; failing loudly here keeps the bug from try: yield finally: + # on_invocation_end has run the leftover hooks on every path that + # reaches it. This covers a scope that ends before it is called. + self._finish_hooks() self._invocation_status = None self._operations_provider = None # Drop this invocation's plugin instances: nothing outlives the - # invocation. Every dispatch is synchronous, so there is no queued - # work still holding one. + # invocation, and the change thread has exited. self._plugins = [] self._started = [] + def record_end_hooks( + self, operations: Sequence[Operation], waited_for: Collection[str] = () + ) -> None: + """Record the end hook of every terminal operation a checkpoint response reports. + + The checkpoint thread calls this right before the response's + operations become readable, and runs no hook. + Recording before publishing is what makes a later claim reliable: a + thread that reads a terminal operation finds its hook recorded. + + ``waited_for`` names the operations whose synchronous waiter the + response is about to wake. Their hooks are reserved for that waiter, + which claims them in ``create_checkpoint`` before the operation + returns; a thread that reads such a result meanwhile waits for the + hook, and an ancestor's claim leaves it alone. Every other hook is + unowned: the thread that reads the result claims it, or the handler + thread at the end of the invocation (:meth:`drain_end_hooks`). + + An operation reported again while its hook is unowned keeps one entry, + with the newest state; one reported while its hook is owned, or after + it ran, is not recorded again. + """ + if not self._plugin_factories: + return + infos = [ + self._end_info(operation) + for operation in operations + if self._is_terminal_status(operation.status) + ] + if not infos: + return + with self.claim_lock: + for info in infos: + if info.operation_id in self._ended: + continue + recorded = self._end_hooks.get(info.operation_id) + if recorded is None: + self._end_hooks[info.operation_id] = _RecordedEndHook( + info, reserved=info.operation_id in waited_for + ) + elif recorded.done is None: + recorded.info = info + + def claim_end_hook(self, operation_id: str, *, waiter: bool = False) -> None: + """Run the recorded end hook of one operation on the calling thread, if any. + + Called by a thread that observes the operation's terminal result, + before it acts on the result: ``waiter`` is True for the synchronous + waiter in ``create_checkpoint``, False for a reader. An unowned hook + runs here; so does one reserved for this waiter. A hook another + thread owns is waited for, so no thread acts on the result before the + hook has returned. A call from inside any plugin hook returns at + once, running and waiting for nothing: a read a plugin makes from a + hook inspects the result and takes no part in the ordering, so a + hook that reads its own operation does not wait for itself, and + hooks cannot wait on each other. No entry means nothing happens: + that is every read of an operation completed before this invocation, + and of one whose hook already ran. + + Unowned terminal descendants of the operation run first, deepest + first. A response can carry a completed descendant no thread read, + such as a callback created and not awaited, together with its + parent's completion; running it here keeps a child's end hook before + its parent's, as on the thread that reads each result. + """ + with self.claim_lock: + claim = self.take_end_hook(operation_id, waiter=waiter) + if claim is not None: + claim.run() + + def take_end_hook( + self, operation_id: str, *, waiter: bool = False + ) -> _EndHookClaim | None: + """Decide, under ``claim_lock``, what the calling thread does for an operation. + + The caller holds ``claim_lock`` and calls :meth:`_EndHookClaim.run` + after releasing it. ``get_checkpoint_result`` holds the lock from its + read of the operation state through this decision, and every other + claim is made under the same lock, so no claim is decided between a + read and the decision that follows it. Returns None when there is + nothing to do. + + None also for a thread that is inside any plugin hook. Such a + thread must not wait: an end hook completes its event only when it + returns, so two hooks that read each other's operations would wait + on each other forever, and a change hook that waits for an end hook + while holding a lock that hook needs would do the same. It must not + run a hook either, which would nest one hook in another. The hook + stays recorded for the read that consumes the result, or for the + drain. + """ + if not self._plugin_factories or self._in_hook(): + return None + recorded = self._end_hooks.get(operation_id) + if recorded is None: + return None + if recorded.done is not None and not (waiter and recorded.reserved): + return _EndHookClaim(self, [], recorded.done) + return _EndHookClaim(self, self._claim_tree_locked(recorded), None) + + def _in_hook(self) -> bool: + """True when the calling thread is inside a plugin hook this executor dispatched.""" + return getattr(self._hook_depth, "value", 0) > 0 + + @contextlib.contextmanager + def _in_hook_scope(self) -> Iterator[None]: + """Mark the calling thread as inside a plugin hook for the duration.""" + self._hook_depth.value = getattr(self._hook_depth, "value", 0) + 1 + try: + yield + finally: + self._hook_depth.value -= 1 + + def drain_end_hooks(self) -> None: + """Run every recorded end hook no thread ran, here, on the calling thread. + + Called on the handler thread once every operation thread has ended + and the checkpoint thread has stopped, so nothing is recorded after + it and no owner is still coming. These are terminal operations no + thread observed in this invocation: a completion that arrived after + the thread that would have read it stopped, or a branch stopped by + its parent's early completion. The earliest recorded entry goes + first, with its unowned descendants before it, then the next. + + A thread-control exception from one hook is held, the other trees + still run (the entries of its own tree not reached are retired, see + :meth:`_run_claimed`), and the first one raised is re-raised once + nothing is left. Stopping at the first would leave the other trees + recorded, and the executor's scope exit would then run them after + ``on_invocation_end``. + """ + held_control: BaseException | None = None + while True: + with self.claim_lock: + recorded = next( + ( + entry + for entry in self._end_hooks.values() + if entry.done is None or entry.reserved + ), + None, + ) + if recorded is None: + break + claimed = self._claim_tree_locked(recorded) + try: + self._run_claimed(claimed) + except BaseException as control: # noqa: BLE001 - execute_plugins contains the rest + if held_control is None: + held_control = control + if held_control is not None: + raise held_control + + def _claim_tree_locked(self, root: _RecordedEndHook) -> list[_RecordedEndHook]: + """Take an entry and its unowned descendants for this thread. Caller holds the lock. + + Returns them in the order to run: every descendant before its + ancestor, the root last. Descendants are found through the recorded + parent ids, so only a descendant whose recorded ancestors are all + unowned is reached. A descendant another thread owns is left to that + thread: the two threads are not ordered with each other. + """ + unowned = [entry for entry in self._end_hooks.values() if entry.done is None] + ordered = [root] + taken = {root.info.operation_id} + index = 0 + while index < len(ordered): + parent_id = ordered[index].info.operation_id + for entry in unowned: + if ( + entry.info.parent_id == parent_id + and entry.info.operation_id not in taken + ): + taken.add(entry.info.operation_id) + ordered.append(entry) + index += 1 + for entry in ordered: + if entry.done is None: + entry.done = threading.Event() + entry.reserved = False + ordered.reverse() + return ordered + + def _run_claimed(self, claimed: list[_RecordedEndHook]) -> None: + """Run the hooks this thread took, then retire each entry and release its waiters. + + A thread-control exception from one hook propagates at once; the + entries not reached are retired without running, so no thread waits + on them forever, and they are not reported later either. + """ + for index, entry in enumerate(claimed): + try: + self.execute_plugins(entry.info) + except BaseException: + for skipped in claimed[index + 1 :]: + self._retire(skipped) + raise + finally: + self._retire(entry) + + def _retire(self, entry: _RecordedEndHook) -> None: + with self.claim_lock: + self._end_hooks.pop(entry.info.operation_id, None) + self._ended.add(entry.info.operation_id) + if entry.done is not None: + entry.done.set() + + def _run_end_hook_once(self, info: OperationEndInfo) -> None: + """Run an end hook on the calling thread unless this invocation already ran one for the operation. + + The replay path uses this: it reports an operation the invocation + input listed as updated, which a checkpoint response may also have + recorded. Both paths go through the registry, so the operation ends + once, and a thread on either path waits for a hook the other is + running. + """ + with self.claim_lock: + if info.operation_id in self._ended: + return + if info.operation_id not in self._end_hooks: + self._end_hooks[info.operation_id] = _RecordedEndHook(info) + claim = self.take_end_hook(info.operation_id) + if claim is not None: + claim.run() + + def _submit_change(self, info: OperationChangeInfo) -> None: + """Queue a change notification for the plugin change thread. Runs nothing here.""" + with self._dispatcher_lock: + dispatcher = self._dispatcher_locked() + dispatcher.submit(info) + + def record_change( + self, + updated_operations: Sequence[Operation], + operations: Mapping[str, Operation], + previous_operations: Mapping[str, Operation], + ) -> None: + """Queue one checkpoint response's change hook for the plugin change thread. + + The checkpoint thread calls this after the response's operations are + published and before it wakes the waiters, and runs no hook. Nothing + is queued when no operation changed status. + """ + info = self._change_info(updated_operations, operations, previous_operations) + if info is not None: + self._submit_change(info) + + def _dispatcher_locked(self) -> _ChangeDispatcher: + """The invocation's change dispatcher, created on first use. Caller holds the lock.""" + if self._dispatcher is None: + try: + self._dispatcher = _ChangeDispatcher.start(self) + except RuntimeError: + # The process cannot start a thread. Observability must not + # fail the execution, so the notifications merge into one that + # on_invocation_end delivers on the handler thread. + logger.warning( + "Plugin change thread could not be started; " + "on_operation_change runs once on the handler thread " + "before on_invocation_end", + exc_info=True, + ) + self._dispatcher = _ChangeDispatcher(self) + return self._dispatcher + + def _finish_hooks(self) -> None: + """Run the unclaimed end hooks, then deliver every pending change notification. + + Called on the handler thread once the checkpoint thread has stopped, + so nothing is recorded or submitted after this returns. The change + thread is joined even if a drained hook raises. + """ + try: + self.drain_end_hooks() + finally: + with self._dispatcher_lock: + dispatcher, self._dispatcher = self._dispatcher, None + if dispatcher is not None: + dispatcher.join() + def _create_plugins(self, info: InvocationStartInfo) -> None: """Build this invocation's plugin instances from its start info. @@ -881,7 +1480,7 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: if control is not None: raise control from None - def execute_plugins(self, info): + def execute_plugins(self, info, *, finish_fan_out: bool = False): """Dispatch one hook to this invocation's plugins. A plugin receives a hook only once it has received the invocation-start @@ -902,10 +1501,12 @@ def execute_plugins(self, info): started, so cutting the loop short costs a plugin its only chance to finish: Telemetry would not drain, and OTel would leave spans unended. The first such exception is held and re-raised once every plugin has been - called, so the thread still stops and nothing is swallowed. No other hook - defers: stopping a start-hook loop early leaves later plugins with nothing - to clean up, because the pairing rule above then withholds their end hook - too. + called, so the thread still stops and nothing is swallowed. A caller + passes ``finish_fan_out=True`` to get the same treatment for another + hook; the change thread does, so that one plugin's interrupt does not + cost the other plugins a notification. No other hook defers: stopping a + start-hook loop early leaves later plugins with nothing to clean up, + because the pairing rule above then withholds their end hook too. Anything :meth:`_dispatch_plugin` raises is already a thread-control failure -- it contains everything else -- so the end path catches @@ -921,27 +1522,38 @@ def execute_plugins(self, info): raised there would be swallowed and the end-hook fan-out could not hold it. An earlier ``sync`` parameter offered the pool path; no caller ever passed it, and it is removed rather than left as a way to opt out of - those rules. + those rules. The hooks built on the checkpoint thread are not dispatched + there: the checkpoint loop records them, and the thread that claims an + end hook, or the plugin change thread, makes this call, so that thread + is the caller and the pairing and re-raise rules above hold there. + + The calling thread is marked as inside a plugin hook for the whole + fan-out, so a result read a hook makes neither runs nor waits for an + end hook (see :meth:`take_end_hook`). This and + :meth:`provide_propagation_metadata` are the two calls into a + plugin's hook methods. """ if not self._plugin_factories: return starting = isinstance(info, InvocationStartInfo) - ending = isinstance(info, InvocationEndInfo) + ending = finish_fan_out or isinstance(info, InvocationEndInfo) deferred_control: BaseException | None = None - for plugin in self._plugins if starting else self._started: - if starting: - # Count a hook before dispatch: even a failed Start can own state - # that only its End hook can release. All lifecycle hooks share - # the invocation worker's Context, including after a failure. - self._started.append(plugin) - if not ending: - self._dispatch_plugin(plugin, info) - continue - try: - self._dispatch_plugin(plugin, info) - except BaseException as control: # noqa: BLE001 - held and re-raised below - if deferred_control is None: - deferred_control = control + with self._in_hook_scope(): + for plugin in self._plugins if starting else self._started: + if starting: + # Count a hook before dispatch: even a failed Start can own + # state that only its End hook can release. All lifecycle + # hooks share the invocation worker's Context, including + # after a failure. + self._started.append(plugin) + if not ending: + self._dispatch_plugin(plugin, info) + continue + try: + self._dispatch_plugin(plugin, info) + except BaseException as control: # noqa: BLE001 - held and re-raised below + if deferred_control is None: + deferred_control = control if deferred_control is not None: raise deferred_control @@ -1059,6 +1671,17 @@ def on_invocation_end( # on_invocation_start not called, skip return + # Every other hook of this invocation has run before the end hook: the + # operation threads have ended, the end hooks no thread claimed run + # here, and the change thread is joined. A thread-control exception a + # drained hook raises on this thread is held, so every plugin still + # receives its end hook, and re-raised after it. + held_control: BaseException | None = None + try: + self._finish_hooks() + except BaseException as control: # noqa: BLE001 - execute_plugins contains the rest + held_control = control + # Re-read the operation map so the end hook sees the state as of the end # of this invocation, not the snapshot taken at its start. invocation_end_info = ( @@ -1068,7 +1691,16 @@ def on_invocation_end( operations=self._snapshot_operation_infos(self._operations_provider), ) ) - self.execute_plugins(invocation_end_info) + try: + self.execute_plugins(invocation_end_info) + except BaseException: + if held_control is not None: + # Both are thread control. The drained hook's came first; the + # end hook's is chained as its context. + raise held_control # noqa: B904 - chained below as __context__ + raise + if held_control is not None: + raise held_control def on_user_function_start( self, @@ -1181,7 +1813,14 @@ def on_operation_update( operations: Mapping[str, Operation] | None = None, previous_operations: Mapping[str, Operation] | None = None, ): - """Execute any registered plugins for operation updates. + """Run the end hooks for operation updates here, and queue the change hook. + + Called for an operation that changed while the execution was suspended, + by the thread that replays it; the end hook runs on that thread. The + change hook, built only when both operation maps are given, goes to + the plugin change thread. Checkpoint responses during the invocation + take the other path: the checkpoint thread records their hooks with + :meth:`record_end_hooks` and :meth:`record_change`, and runs none. Updates such as STARTED might be omitted because START and completion action (e.g. SUCCEED/FAIL) may be checkpointed in batch and the backend returns only the terminal status (e.g. SUCCEEDED/PENDING/FAILED). @@ -1205,33 +1844,44 @@ def on_operation_update( ) for operation in updated_operations: if self._is_terminal_status(operation.status): - self.execute_plugins( - OperationEndInfo( - operation_id=operation.operation_id, - operation_type=_to_operation_type(operation.operation_type), - sub_type=operation.sub_type, - name=operation.name, - parent_id=operation.parent_id, - start_time=operation.start_timestamp, - end_time=operation.end_timestamp, - result=_extract_result(operation), - status=operation.status, - error=self._extract_error(operation), - attempt=( - operation.step_details.attempt - if operation.step_details - else None - ), - is_replayed=False, - ), - ) + self._run_end_hook_once(self._end_info(operation)) + + change_info = self._change_info( + updated_operations, operations, previous_operations + ) + if change_info is not None: + self._submit_change(change_info) + def _end_info(self, operation: Operation) -> OperationEndInfo: + return OperationEndInfo( + operation_id=operation.operation_id, + operation_type=_to_operation_type(operation.operation_type), + sub_type=operation.sub_type, + name=operation.name, + parent_id=operation.parent_id, + start_time=operation.start_timestamp, + end_time=operation.end_timestamp, + result=_extract_result(operation), + status=operation.status, + error=self._extract_error(operation), + attempt=( + operation.step_details.attempt if operation.step_details else None + ), + is_replayed=False, + ) + + def _change_info( + self, + updated_operations: Sequence[Operation], + operations: Mapping[str, Operation] | None, + previous_operations: Mapping[str, Operation] | None, + ) -> OperationChangeInfo | None: if ( operations is None or previous_operations is None or self._invocation_status is None ): - return + return None changed_operations = [ operation @@ -1240,17 +1890,15 @@ def on_operation_update( or previous_operations[operation.operation_id].status != operation.status ] if not changed_operations: - return + return None - self.execute_plugins( - OperationChangeInfo( - execution_arn=self._invocation_status.execution_arn, - updated_operations={ - operation.operation_id: OperationInfo.from_operation(operation) - for operation in changed_operations - }, - operations=_to_operation_info_map(operations), - ), + return OperationChangeInfo( + execution_arn=self._invocation_status.execution_arn, + updated_operations={ + operation.operation_id: OperationInfo.from_operation(operation) + for operation in changed_operations + }, + operations=_to_operation_info_map(operations), ) @staticmethod diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index f77d5d156..e9a7588d2 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -49,6 +49,8 @@ import datetime from collections.abc import MutableMapping + from aws_durable_execution_sdk_python.plugin import _EndHookClaim + logger = logging.getLogger(__name__) # Longest wait on an empty queue before re-checking the shutdown signal. @@ -462,6 +464,7 @@ def fetch_paginated_operations( initial_operations: list[Operation], checkpoint_token: str, next_marker: str | None, + before_publish: Callable[[list[Operation]], None] | None = None, ) -> list[Operation]: """Add initial operations and fetch all paginated operations from the Durable Functions API. This method is thread_safe. @@ -471,6 +474,11 @@ def fetch_paginated_operations( initial_operations: initial operations to be added to ExecutionState checkpoint_token: checkpoint token used to call Durable Functions API. next_marker: a marker indicates that there are paginated operations. + before_publish: called with the fetched operations right before + they become readable through get_checkpoint_result. The + checkpoint loop records the operations' plugin end hooks here, + so a thread that reads a terminal operation always finds its + hook recorded. Returns: List of all operations fetched from the Durable Functions API @@ -501,6 +509,8 @@ def fetch_paginated_operations( finally: # Always store whatever operations we successfully fetched if all_operations: + if before_publish is not None: + before_publish(all_operations) with self._operations_lock: self._operations.update( {op.operation_id: op for op in all_operations} @@ -549,7 +559,9 @@ def has_prior_operations(self) -> bool: for op in self._operations.values() ) - def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult: + def get_checkpoint_result( + self, checkpoint_id: str, *, claim_end_hook: bool = True + ) -> CheckpointedResult: """Get checkpoint result. Note this does not invoke the Durable Functions API. It only checks @@ -560,6 +572,14 @@ def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult: Args: checkpoint_id: str - id for checkpoint to retrieve. + claim_end_hook: With plugins, a read of a terminal operation also + runs, or waits for, the end hook a checkpoint response recorded + for it during this invocation (see the comment in the body). + The default. A read that only inspects a checkpoint passes + False: DurableContext reads the next operation's checkpoint to + place the replay boundary, and an operation's identity before + it runs. Neither read consumes the result, so neither takes the + hook; the read that consumes it does. Returns: CheckpointedResult with is_succeeded True if the checkpoint exists and its @@ -567,9 +587,37 @@ def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult: SUCCEEDED, or if the checkpoint doesn't exist, then return CheckpointedResult with is_succeeded=False,result=None. """ - # checking status are deliberately under a lighter non-serialized lock - with self._operations_lock: - checkpoint = self._operations.get(checkpoint_id) + if not self._plugin_executor.has_plugins or not claim_end_hook: + # checking status are deliberately under a lighter non-serialized lock + with self._operations_lock: + checkpoint = self._operations.get(checkpoint_id) + else: + # Every thread that consumes a result it did not checkpoint itself + # reads it here: the operation executors through + # OperationExecutor._get_checkpoint_result, Callback.result, and + # the map and parallel coordinators for their branches. So this is + # where such a thread takes the end hook of an operation a + # checkpoint response completed during this invocation, to run it + # before it acts on the result, or to wait for it if another + # thread owns it. The replay-boundary reads in DurableContext pass + # claim_end_hook=False and take the branch above: they inspect the + # checkpoint without consuming the result, so the hook stays for + # the read that does, on the thread that does. The checkpoint + # loop records the hook before it publishes the operation, so a + # reader that sees the terminal state finds the hook. The read and + # the decision happen under the plugin executor's claim lock, + # which every claim takes, so no other thread's claim is decided + # between the two. Only a terminal read takes a hook: a read of + # the older, non-terminal state did not see the completion. The + # hook itself runs below, outside every lock. + claim: _EndHookClaim | None = None + with self._plugin_executor.claim_lock: + with self._operations_lock: + checkpoint = self._operations.get(checkpoint_id) + if checkpoint and checkpoint.status in _TERMINAL_OPERATION_STATUSES: + claim = self._plugin_executor.take_end_hook(checkpoint_id) + if claim is not None: + claim.run() if checkpoint: return CheckpointedResult.create_from_operation(checkpoint) @@ -860,7 +908,25 @@ def create_checkpoint( raise DurableExecutionsError(msg) # Wait for completion - will raise BackgroundThreadError if background thread fails - completion_event.wait() + try: + completion_event.wait() + finally: + # This thread now observes the operation's result, so it runs + # the operation's end hook, if the response recorded one, + # before the operation returns or raises to user code. The + # hook therefore follows this operation's user-function hooks + # and precedes this thread's next operation. A wait that + # raises was woken by a failed, revoked or settled batch. A + # revoked or settled batch records nothing. A failed batch + # records nothing either, except when the failure is in + # fetching a later page: the pages already fetched are + # recorded and published, so a hook can run here while the + # BackgroundThreadError propagates, which is the order the + # hook would have had on a successful batch. + if operation_update is not None: + self._plugin_executor.claim_end_hook( + operation_update.operation_id, waiter=True + ) else: logger.debug("Enqueued checkpoint operation for asynchronous processing") @@ -1111,13 +1177,38 @@ def checkpoint_batches_forever(self) -> None: if execution_completed else output.new_execution_state.next_marker ) + # No plugin hook runs on this thread. The response's end + # hooks are recorded before its operations become readable, + # so the thread that observes a terminal result (the waiter + # woken below, or a later reader in get_checkpoint_result) + # finds the hook and runs it before it acts on the result. + # A hook whose operation has a waiter in this batch is + # reserved for that waiter. + waited_for = { + queued_op.operation_update.operation_id + for queued_op in batch + if queued_op.operation_update is not None + and queued_op.completion_event is not None + } updated_operations = self.fetch_paginated_operations( output.new_execution_state.operations, current_checkpoint_token, fetch_marker, + before_publish=( + functools.partial( + self._plugin_executor.record_end_hooks, + waited_for=waited_for, + ) + if has_plugins + else None + ), ) - if has_plugins: - self._plugin_executor.on_operation_update( + if previous_operations is not None: + # The change hook describes the whole response, so it + # goes to the plugin change thread; it is submitted + # before the waiters wake so that it is pending before + # the invocation can end. + self._plugin_executor.record_change( updated_operations, self.operations, previous_operations, diff --git a/packages/aws-durable-execution-sdk-python/tests/context_test.py b/packages/aws-durable-execution-sdk-python/tests/context_test.py index ac996fde3..623fee84e 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -705,7 +705,7 @@ def test_operation_replay_aware_looks_up_allocated_operation_id(): assert context._step_counter.get_current() == 1 # noqa: SLF001 mock_state.get_checkpoint_result.assert_called_once_with( - operation_identifier.operation_id + operation_identifier.operation_id, claim_end_hook=False ) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py index f42008990..e79994f8c 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py @@ -77,9 +77,9 @@ def _event() -> dict: def _tracking_checkpoint(): """Checkpoint mock that accumulates operations, as the service would. - SUCCEED actions are recorded as SUCCEEDED so the SDK dispatches - operation-end hooks, which is how a hook reaches the background - checkpointing thread. + SUCCEED actions are recorded as SUCCEEDED so the checkpoint responses + carry terminal operations, which is what makes the SDK record and run + operation-end hooks. """ operations: dict[str, Operation] = {} @@ -140,19 +140,27 @@ class _ContextvarSettingPlugin(DurableInstrumentationPlugin): def __init__(self, value: str) -> None: self._value = value self.operation_end_observations: list[tuple[str, str]] = [] + self.operation_change_observations: list[tuple[str, str]] = [] self._lock = threading.Lock() def on_invocation_start(self, info) -> None: # noqa: ARG002 _probe.set(self._value) def on_operation_end(self, info) -> None: # noqa: ARG002 - # Dispatched from the background checkpointing thread, which reads back - # the terminal status of a checkpointed operation. + # Runs on the thread that waited for the operation's checkpoint. with self._lock: self.operation_end_observations.append( (threading.current_thread().name, _probe.get()) ) + def on_operation_change(self, info) -> None: # noqa: ARG002 + # Runs on the invocation's plugin change thread, after the checkpoint + # response is merged. + with self._lock: + self.operation_change_observations.append( + (threading.current_thread().name, _probe.get()) + ) + def test_handler_body_sees_a_contextvar_set_by_a_plugin_at_invocation_start(): """The invocation-start hook's contextvar reaches top-level handler code. @@ -229,13 +237,12 @@ def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 assert after == "set-by-caller" -def test_checkpointing_thread_does_not_carry_the_invocation_context(): - """The background checkpointing thread is submitted without the context. +def test_plugin_change_thread_does_not_carry_the_invocation_context(): + """The plugin change thread, like the checkpoint thread, starts from an empty context. - It runs SDK checkpointing rather than user code, so it is left starting from - an empty context. This test pins that choice: a plugin hook dispatched from - that thread sees the contextvar's default, not the value the invocation - thread set. + Both run SDK work rather than user code. This test verifies that: the + change hook, which runs on the plugin change thread, sees the contextvar's + default, not the value the invocation thread set at invocation start. """ plugin = _ContextvarSettingPlugin("claimed-by-plugin") handler_thread_names: list[str] = [] @@ -254,8 +261,36 @@ def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 assert result["Status"] == InvocationStatus.SUCCEEDED.value off_handler_thread = [ value - for thread_name, value in plugin.operation_end_observations + for thread_name, value in plugin.operation_change_observations if thread_name not in handler_thread_names ] - assert off_handler_thread, plugin.operation_end_observations + assert off_handler_thread, plugin.operation_change_observations assert set(off_handler_thread) == {UNSET} + + +def test_operation_end_hook_runs_on_the_operation_thread_with_its_context(): + """A step's end hook runs on the thread that ran the step, in its context. + + The checkpoint thread records the end hook; the thread waiting for the + step's checkpoint claims and runs it. So the hook sees what that thread + sees, including the contextvar the plugin set at invocation start. + """ + plugin = _ContextvarSettingPlugin("claimed-by-plugin") + handler_thread_names: list[str] = [] + + @durable_step + def a_step(_step_context) -> str: + return "stepped" + + @durable_execution(plugins=[plugin_factory(plugin)]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + handler_thread_names.append(threading.current_thread().name) + return context.step(a_step(), name="a-step") + + result = _run(my_handler) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert plugin.operation_end_observations + assert plugin.operation_end_observations == [ + (handler_thread_names[0], "claimed-by-plugin") + ] diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_hook_contract_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_hook_contract_int_test.py new file mode 100644 index 000000000..dc79daa83 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_hook_contract_int_test.py @@ -0,0 +1,383 @@ +"""The plugin hook contract, run through the real engine. + +The unit tests in ``plugin_hook_registry_test.py`` drive ``PluginExecutor`` +and the checkpoint loop directly. These cases drive a handler through +``DurableFunctionTestRunner``, so the hooks run on the threads the SDK really +uses: an operation's end hook on the thread that observes its result, the +change hook on the invocation's plugin change thread. +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +from aws_durable_execution_sdk_python.config import ( + CompletionConfig, + Duration, + NestingType, + ParallelConfig, +) +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + OperationChangeInfo, + OperationEndInfo, + OperationStartInfo, + OperationType, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from tests.test_helpers import plugin_factory + + +class _TimelinePlugin(DurableInstrumentationPlugin): + """Records when each hook was entered and left, by operation name, and on which thread.""" + + def __init__(self, end_hook_delay: dict[str, float] | None = None) -> None: + self.end_hook_delay = end_hook_delay or {} + self.events: list[tuple[float, str, str]] = [] + self._lock = threading.Lock() + + def _record(self, label: str) -> None: + with self._lock: + self.events.append( + (time.monotonic(), label, threading.current_thread().name) + ) + + def labels(self) -> list[str]: + with self._lock: + return [label for _, label, _ in self.events] + + def at(self, label: str) -> float: + with self._lock: + return next(stamp for stamp, found, _ in self.events if found == label) + + def thread_of(self, label: str) -> str: + with self._lock: + return next(thread for _, found, thread in self.events if found == label) + + def on_operation_start(self, info: OperationStartInfo) -> None: + self._record(f"start:{info.name}") + + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.operation_type is OperationType.EXECUTION: + return + self._record(f"end-enter:{info.name}") + delay = self.end_hook_delay.get(info.name or "") + if delay: + time.sleep(delay) + self._record(f"end-leave:{info.name}") + + def on_operation_change(self, info: OperationChangeInfo) -> None: # noqa: ARG002 + self._record("change") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: # noqa: ARG002 + self._record("invocation_end") + + +def _run(handler, plugin: _TimelinePlugin): + wrapped = durable_execution(handler, plugins=[plugin_factory(plugin)]) + with DurableFunctionTestRunner(handler=wrapped) as runner: + return runner.run(input="{}", execution_timeout=60) + + +def test_an_end_hook_completes_before_its_operation_returns() -> None: + """A step's end hook, even a slow one, has returned when the step returns. + + So on one thread the end hook of one operation precedes the start hook of + the next, which is the order the OpenTelemetry plugins and the shared + conformance suite rely on. + """ + plugin = _TimelinePlugin(end_hook_delay={"first": 0.3}) + step_returned: dict[str, float] = {} + + def handler(_event: Any, context: DurableContext) -> str: + context.step(lambda _c: "first", name="first") + step_returned["first"] = time.monotonic() + context.step(lambda _c: "second", name="second") + return "done" + + result = _run(handler, plugin) + + assert result.result == '"done"' + assert plugin.at("end-leave:first") <= step_returned["first"] + labels = plugin.labels() + assert labels.index("end-leave:first") < labels.index("start:second") + assert labels.index("end-leave:second") < labels.index("invocation_end") + + +def test_an_end_hook_runs_on_the_thread_that_ran_the_operation() -> None: + """Top-level operations: the handler thread. Branch operations: the branch thread. + + A child context runs on its caller's thread, so a parallel branch is used + for the branch case; its thread differs from the handler's. + """ + plugin = _TimelinePlugin() + threads: dict[str, str] = {} + + def branch(ctx: DurableContext) -> str: + threads["branch"] = threading.current_thread().name + return ctx.step(lambda _c: "inner", name="inner") + + def handler(_event: Any, context: DurableContext) -> str: + threads["handler"] = threading.current_thread().name + context.step(lambda _c: "top", name="top") + return context.parallel([branch], name="fan").get_results()[0] + + result = _run(handler, plugin) + + assert result.result == '"inner"' + assert threads["branch"] != threads["handler"] + assert plugin.thread_of("end-enter:top") == threads["handler"] + assert plugin.thread_of("end-enter:inner") == threads["branch"] + # The change hook runs on the plugin change thread: neither of those, and + # not the checkpoint thread either. + assert plugin.thread_of("change") == "plugin-change" + + +def test_a_slow_end_hook_delays_only_its_own_branch() -> None: + """Two parallel branches checkpoint through the one checkpoint thread. + + Branch 0's step has a slow end hook. When hooks ran on the checkpoint + thread, that hook delayed every checkpoint submitted after it, so branch + 1's step could not return until it finished. Now the hook runs on branch + 0's thread. Branch 1 checkpoints its step only once the slow hook has + entered, so its checkpoint is submitted after the slow one in every + interleaving, and it still returns while the slow hook is running. + """ + slow_hook_entered = threading.Event() + release_slow_hook = threading.Event() + + class _GatedPlugin(_TimelinePlugin): + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.name == "slow": + self._record("end-enter:slow") + slow_hook_entered.set() + # An assertion here would be contained by the SDK; a label is + # asserted on instead. + if not release_slow_hook.wait(10): + self._record("timed-out") + self._record("end-leave:slow") + return + super().on_operation_end(info) + + plugin = _GatedPlugin() + returned: dict[str, float] = {} + + def slow_branch(ctx: DurableContext) -> str: + ctx.step(lambda _c: "slow", name="slow") + returned["slow"] = time.monotonic() + return "slow" + + def quick_branch(ctx: DurableContext) -> str: + assert slow_hook_entered.wait(10) + ctx.step(lambda _c: "quick", name="quick") + returned["quick"] = time.monotonic() + # The quick branch returned while the slow hook was still blocked. + release_slow_hook.set() + return "quick" + + def handler(_event: Any, context: DurableContext) -> list[str]: + return context.parallel([slow_branch, quick_branch], name="both").get_results() + + result = _run(handler, plugin) + + assert result.result == '["slow", "quick"]' + assert "timed-out" not in plugin.labels() + assert returned["quick"] < plugin.at("end-leave:slow") + assert plugin.at("end-leave:slow") <= returned["slow"] + + +def test_nested_end_hooks_complete_in_order_child_before_parent() -> None: + """parent.start, child.start, child.end, parent.end. + + The child's terminal checkpoint is synchronous, and its end hook runs on + the child's thread before the child returns. So the child's hook has + returned before the parent can complete and run its own. + """ + plugin = _TimelinePlugin(end_hook_delay={"inner": 0.2}) + + def child(ctx: DurableContext) -> str: + return ctx.step(lambda _c: "stepped", name="inner") + + def handler(_event: Any, context: DurableContext) -> str: + return context.run_in_child_context(child, name="outer") + + result = _run(handler, plugin) + + assert result.result == '"stepped"' + labels = plugin.labels() + assert labels.index("start:outer") < labels.index("start:inner") + assert labels.index("end-leave:inner") < labels.index("end-enter:outer") + assert labels.index("end-leave:outer") < labels.index("invocation_end") + + +def test_map_item_end_hooks_complete_before_the_map_end_hook() -> None: + plugin = _TimelinePlugin(end_hook_delay={"doubles": 0.05}) + + def handler(_event: Any, context: DurableContext) -> list[int]: + return context.map( + [1, 2, 3], lambda _ctx, item, _index, _items: item * 2, name="doubles" + ).get_results() + + result = _run(handler, plugin) + + assert result.result == "[2, 4, 6]" + labels = plugin.labels() + map_end = labels.index("end-enter:doubles") + item_leaves = [ + index + for index, label in enumerate(labels) + if label.startswith("end-leave:") and label != "end-leave:doubles" + ] + assert len(item_leaves) == 3 + assert all(leave < map_end for leave in item_leaves) + assert labels.index("end-leave:doubles") < labels.index("invocation_end") + + +def test_early_completion_lets_a_branch_end_after_its_parent() -> None: + """The documented exception to child-before-parent. + + With ``first_successful`` the parallel completes when the first branch + does, while the second branch is still running. The parent's end hook + runs then, on the handler thread. The second branch is a FLAT, virtual + context: it checkpoints nothing, so when it returns it runs its own end + hook on its thread, after the parent's. The SDK joins every branch thread + before ``on_invocation_end``, so that hook still precedes it. + + The quick branch returns only once the slow branch is running. On early + completion the SDK cancels the branches the pool has not started yet, and + a branch that never ran has no end hook; so without this, a slow machine + would have no hook to find. + """ + slow_started = threading.Event() + parent_ended = threading.Event() + + class _ReleasingPlugin(_TimelinePlugin): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.name == "race": + parent_ended.set() + + plugin = _ReleasingPlugin() + + def quick(_ctx: DurableContext) -> str: + assert slow_started.wait(10) + return "quick" + + def slow(_ctx: DurableContext) -> str: + slow_started.set() + assert parent_ended.wait(10) + return "slow" + + def handler(_event: Any, context: DurableContext) -> str: + results = context.parallel( + [quick, slow], + config=ParallelConfig( + completion_config=CompletionConfig.first_successful(), + nesting_type=NestingType.FLAT, + ), + name="race", + ) + return results.get_results()[0] + + result = _run(handler, plugin) + + assert result.result == '"quick"' + labels = plugin.labels() + assert labels.index("end-enter:parallel-branch-0") < labels.index("end-enter:race") + assert labels.index("end-enter:race") < labels.index("end-enter:parallel-branch-1") + assert labels.index("end-leave:parallel-branch-1") < labels.index("invocation_end") + + +def test_a_wait_in_a_branch_resumed_by_a_refresh_ends_before_the_branches_next_step() -> ( + None +): + """The wait completes in a refresh response, which no operation thread waits for. + + The other branch stays in a step until the first branch is past its wait, + so the coordinator is still running when the wait comes due: it waits for + the refresh and resumes the branch in this invocation. The resumed branch + thread re-traverses the wait and reads its result; that read runs the end + hook, on that thread, before the branch starts its next step. So the + wait's end hook precedes the step's start hook on one thread, which is the + order the OpenTelemetry suite asserts, and the end of the invocation finds + nothing left to run for it. + """ + plugin = _TimelinePlugin() + step_threads: dict[str, str] = {} + first_branch_done = threading.Event() + + def resumed(ctx: DurableContext) -> str: + ctx.step(lambda _c: "before", name="before-wait") + ctx.wait(Duration.from_seconds(1), name="branch-wait") + step_threads["after-wait"] = threading.current_thread().name + after = ctx.step(lambda _c: "after", name="after-wait") + first_branch_done.set() + return after + + def hold(_c: Any) -> str: + # Keeps the coordinator in the invocation until the first branch has + # resumed; the timeout only bounds a failing test. + first_branch_done.wait(30) + return "other" + + def other(ctx: DurableContext) -> str: + return ctx.step(hold, name="hold-step") + + def handler(_event: Any, context: DurableContext) -> list[str]: + return context.parallel([resumed, other], name="both").get_results() + + result = _run(handler, plugin) + + assert result.result == '["after", "other"]' + labels = plugin.labels() + assert labels.count("invocation_end") == 1, ( + "the invocation suspended instead of resuming the branch in-process" + ) + assert labels.count("end-enter:branch-wait") == 1 + assert labels.index("end-leave:branch-wait") < labels.index("start:after-wait") + assert plugin.thread_of("end-enter:branch-wait") == step_threads["after-wait"] + + +def test_a_callback_result_read_ends_the_callback_before_the_next_step() -> None: + """The submitter completes the callback, so its result arrives in a response + no operation thread waits for; the thread that reads the result runs the + end hook first. When the completion arrives only after this invocation + suspended, the replaying thread runs it instead. Either way the callback's + end hook precedes the next step's start hook on one thread, once. + """ + + class _CallbackAware(_TimelinePlugin): + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.operation_type is OperationType.CALLBACK: + self._record("end-enter:the-callback") + self._record("end-leave:the-callback") + return + super().on_operation_end(info) + + plugin = _CallbackAware() + step_threads: dict[str, str] = {} + + def handler(_event: Any, context: DurableContext) -> str: + context.wait_for_callback( + lambda callback_id, _ctx: runner.send_callback_success( + callback_id, b'"completed"' + ), + name="await-external", + ) + step_threads["after-callback"] = threading.current_thread().name + return context.step(lambda _c: "after", name="after-callback") + + wrapped = durable_execution(handler, plugins=[plugin_factory(plugin)]) + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", execution_timeout=60) + + assert result.result == '"after"' + labels = plugin.labels() + assert labels.count("end-enter:the-callback") == 1 + assert labels.index("end-leave:the-callback") < labels.index("start:after-callback") + assert plugin.thread_of("end-enter:the-callback") == step_threads["after-callback"] diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_hook_registry_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_hook_registry_test.py new file mode 100644 index 000000000..c2f042605 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_hook_registry_test.py @@ -0,0 +1,1704 @@ +"""Where plugin hooks run once the checkpoint thread stops running them. + +The checkpoint thread records the hooks of one checkpoint response and runs +none. It records each terminal operation's end hook before it publishes the +response's operations, so the thread that then observes a terminal result, +the waiter of a synchronous checkpoint or a later reader, claims the hook and +runs it before it acts on the result. End hooks no thread claims run on the +handler thread before ``on_invocation_end``. The change hook goes to the +invocation's plugin change thread, which merges the notifications that arrive +while a plugin is still in the hook. + +A failing case here means a change breaks that placement or order; the class +docstring of ``DurableInstrumentationPlugin`` states it as the public contract. +""" + +from __future__ import annotations + +import datetime +import functools +import logging +import threading +import time +from collections.abc import Callable, Collection +from types import TracebackType +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python import plugin as plugin_module +from aws_durable_execution_sdk_python.exceptions import GetExecutionStateError +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + DurableExecutionInvocationOutput, + InvocationStatus as ServiceInvocationStatus, + LambdaClient, + Operation, + OperationAction, + OperationStatus, + OperationSubType, + OperationType as ServiceOperationType, + OperationUpdate, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + OperationChangeInfo, + OperationEndInfo, + OperationInfo, + PluginExecutor, + PropagationInput, + PropagationMetadata, + _merge_change_infos, +) +from aws_durable_execution_sdk_python import state as state_module +from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext +from aws_durable_execution_sdk_python.identifier import OperationIdentifier +from aws_durable_execution_sdk_python.state import ExecutionState, ReplayStatus +from aws_durable_execution_sdk_python.threading import CompletionEvent +from tests.test_helpers import plugin_factory + +START_TS = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) +SUCCEEDED = DurableExecutionInvocationOutput(status=ServiceInvocationStatus.SUCCEEDED) +TIMEOUT = 10.0 + + +def _operation( + operation_id: str, + status: OperationStatus = OperationStatus.SUCCEEDED, + parent_id: str | None = None, +) -> Operation: + return Operation( + operation_id=operation_id, + operation_type=ServiceOperationType.STEP, + status=status, + name=operation_id, + sub_type=OperationSubType.STEP, + parent_id=parent_id, + start_timestamp=START_TS, + end_timestamp=START_TS, + ) + + +def _update( + operation_id: str, action: OperationAction = OperationAction.SUCCEED +) -> OperationUpdate: + return OperationUpdate( + operation_id=operation_id, + operation_type=ServiceOperationType.STEP, + action=action, + name=operation_id, + ) + + +def _map(*operations: Operation) -> dict[str, Operation]: + return {operation.operation_id: operation for operation in operations} + + +class _Recorder(DurableInstrumentationPlugin): + """Records every hook, and any marker the test adds, with the thread it ran on.""" + + def __init__(self) -> None: + self.events: list[tuple[str, str]] = [] + self.changes: list[OperationChangeInfo] = [] + self.ends: list[OperationEndInfo] = [] + self._lock = threading.Lock() + + def record(self, label: str) -> None: + with self._lock: + self.events.append((label, threading.current_thread().name)) + + def block_on(self, event: threading.Event) -> None: + """Wait inside a hook. An assertion there would be contained by the SDK.""" + if not event.wait(TIMEOUT): + self.record("timed-out") + + def labels(self) -> list[str]: + with self._lock: + return [label for label, _ in self.events] + + def thread_of(self, label: str) -> str: + with self._lock: + return next(name for found, name in self.events if found == label) + + def on_invocation_start(self, info: InvocationStartInfo) -> None: # noqa: ARG002 + self.record("invocation_start") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: # noqa: ARG002 + self.record("invocation_end") + + def on_operation_end(self, info: OperationEndInfo) -> None: + with self._lock: + self.ends.append(info) + self.record(f"operation_end:{info.operation_id}") + + def on_operation_change(self, info: OperationChangeInfo) -> None: + with self._lock: + self.changes.append(info) + self.record("operation_change:" + ",".join(sorted(info.updated_operations))) + + +def _started_executor(plugin: DurableInstrumentationPlugin) -> PluginExecutor: + executor = PluginExecutor([plugin_factory(plugin)]) + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=None, + execution_start_time=START_TS, + is_first_invocation=True, + ) + return executor + + +def _thread(target: Callable[[], object], name: str) -> threading.Thread: + """A started daemon thread: a thread a failed test leaves blocked cannot hold pytest open.""" + thread = threading.Thread(target=target, name=name, daemon=True) + thread.start() + return thread + + +def _join(thread: threading.Thread) -> None: + thread.join(TIMEOUT) + assert not thread.is_alive() + + +def _wait_until(condition: Callable[[], bool], timeout: float = TIMEOUT) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.001) + return condition() + + +def _end_in_thread(executor: PluginExecutor, name: str = "handler") -> threading.Thread: + return _thread(lambda: executor.on_invocation_end(SUCCEEDED), name) + + +# region End hook registry + + +def test_record_runs_no_hook_and_drain_runs_one_per_terminal_operation() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + + executor.record_end_hooks( + [ + _operation("a"), + _operation("started", OperationStatus.STARTED), + _operation("b", OperationStatus.FAILED), + ] + ) + assert recorder.labels() == ["invocation_start"] + + executor.drain_end_hooks() + + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "operation_end:b", + ] + assert recorder.thread_of("operation_end:a") == threading.current_thread().name + assert [info.status for info in recorder.ends] == [ + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + ] + + +def test_claim_runs_the_hook_on_the_claiming_thread_and_only_once() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + claimer = _thread(lambda: executor.claim_end_hook("a"), "claimer") + _join(claimer) + executor.claim_end_hook("a") + executor.drain_end_hooks() + + assert recorder.labels() == ["invocation_start", "operation_end:a"] + assert recorder.thread_of("operation_end:a") == "claimer" + + +def test_concurrent_claims_of_one_operation_run_the_hook_exactly_once() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + go = threading.Event() + + def claim() -> None: + go.wait(TIMEOUT) + executor.claim_end_hook("a") + + claimers = [_thread(claim, f"claimer-{i}") for i in range(8)] + go.set() + for claimer in claimers: + _join(claimer) + + assert recorder.labels().count("operation_end:a") == 1 + + +def test_claiming_an_unrecorded_operation_does_nothing() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + executor.claim_end_hook("not-recorded") + + assert recorder.labels() == ["invocation_start"] + + +def test_the_registry_is_inert_without_plugins() -> None: + executor = PluginExecutor(None) + + executor.record_end_hooks([_operation("a")]) + executor.claim_end_hook("a") + executor.drain_end_hooks() + executor.on_invocation_end(SUCCEEDED) + + +def test_drain_runs_the_unclaimed_hooks_in_recording_order() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + executor.record_end_hooks([_operation("b"), _operation("c")]) + executor.claim_end_hook("b") + + executor.drain_end_hooks() + + assert recorder.labels() == [ + "invocation_start", + "operation_end:b", + "operation_end:a", + "operation_end:c", + ] + + +def test_an_operation_reported_again_before_its_hook_ran_keeps_one_entry() -> None: + """The repeated report replaces the entry's state; the entry keeps its position.""" + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a"), _operation("b")]) + executor.record_end_hooks([_operation("a", OperationStatus.FAILED)]) + + executor.drain_end_hooks() + + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "operation_end:b", + ] + assert recorder.ends[0].status is OperationStatus.FAILED + + +def test_an_operation_the_replay_path_reported_is_not_recorded_again() -> None: + """The invocation input carried the operation; a response lists it again.""" + recorder = _Recorder() + executor = _started_executor(recorder) + + executor.on_operation_update(_operation("a")) + executor.record_end_hooks([_operation("a")]) + executor.drain_end_hooks() + + assert recorder.labels() == ["invocation_start", "operation_end:a"] + + +def test_a_claimed_operation_is_not_reported_again_by_the_replay_path() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + executor.claim_end_hook("a") + executor.on_operation_update(_operation("a")) + + assert recorder.labels() == ["invocation_start", "operation_end:a"] + + +def test_unclaimed_end_hooks_run_on_the_handler_thread_before_invocation_end() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + _join(_end_in_thread(executor)) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "invocation_end", + ] + assert recorder.thread_of("operation_end:a") == "handler" + + +class _ObservedEvent(threading.Event): + """An event whose wait() reports that a thread has entered the wait.""" + + def __init__(self, entered: threading.Event) -> None: + super().__init__() + self._entered = entered + + def wait(self, timeout: float | None = None) -> bool: + self._entered.set() + return super().wait(timeout) + + +def test_a_second_observer_waits_for_a_hook_another_thread_is_running() -> None: + """Two threads observe one result: one runs the hook, the other waits for it. + + Without the wait the second thread would act on the result while the end + hook still runs. The readiness signal comes from inside the wait itself: + the hook's completion event is replaced by one that reports entry into + its wait, so the negative assertion is made while the second thread is in + that wait, and an implementation that does not wait never signals. + """ + hook_entered = threading.Event() + release = threading.Event() + wait_entered = threading.Event() + + class _Slow(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + hook_entered.set() + self.block_on(release) + self.record("end-left") + + recorder = _Slow() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + first = _thread(lambda: executor.claim_end_hook("a"), "first") + assert hook_entered.wait(TIMEOUT) + with executor.claim_lock: + entry = executor._end_hooks["a"] # noqa: SLF001 + assert entry.done is not None + entry.done = _ObservedEvent(wait_entered) + + def second_observer() -> None: + executor.claim_end_hook("a") + recorder.record("second-proceeded") + + second = _thread(second_observer, "second") + assert wait_entered.wait(TIMEOUT), "the second observer did not wait" + assert "second-proceeded" not in recorder.labels() + release.set() + _join(first) + _join(second) + + labels = recorder.labels() + assert "timed-out" not in labels + assert labels.count("operation_end:a") == 1 + assert labels.index("end-left") < labels.index("second-proceeded") + + +def test_a_control_exception_in_one_hook_of_a_tree_retires_the_rest() -> None: + """The entries not reached are retired, so no thread waits on them forever.""" + + class _Interrupting(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id == "child": + raise KeyboardInterrupt + + recorder = _Interrupting() + executor = _started_executor(recorder) + executor.record_end_hooks( + [_operation("parent"), _operation("child", parent_id="parent")] + ) + + with pytest.raises(KeyboardInterrupt): + executor.claim_end_hook("parent") + + # Nothing is left owned-but-unfinished: a later observer returns at once, + # and the drain finds nothing. + _join(_thread(lambda: executor.claim_end_hook("parent"), "late-observer")) + executor.drain_end_hooks() + assert recorder.labels() == ["invocation_start", "operation_end:child"] + + +def test_a_hook_that_reads_its_own_operation_does_not_wait_for_itself() -> None: + class _Reentrant(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + executor.claim_end_hook(info.operation_id) + self.record("reentered") + + recorder = _Reentrant() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + _join(_thread(lambda: executor.claim_end_hook("a"), "claimer")) + + assert recorder.labels() == ["invocation_start", "operation_end:a", "reentered"] + + +def test_a_hook_reserved_for_a_waiter_is_run_by_the_waiter_and_awaited_by_a_reader() -> ( + None +): + """The response wakes the operation's synchronous waiter: the hook is its. + + A reader that observes the result first waits for the waiter's hook, and + an ancestor's claim leaves the reserved hook alone. + """ + recorder = _Recorder() + executor = _started_executor(recorder) + wait_entered = threading.Event() + executor.record_end_hooks( + [_operation("parent"), _operation("child", parent_id="parent")], + waited_for={"child"}, + ) + with executor.claim_lock: + entry = executor._end_hooks["child"] # noqa: SLF001 + assert entry.done is not None + entry.done = _ObservedEvent(wait_entered) + + executor.claim_end_hook("parent") + assert recorder.labels() == ["invocation_start", "operation_end:parent"] + + def reader() -> None: + executor.claim_end_hook("child") + recorder.record("reader-proceeded") + + reading = _thread(reader, "reader") + assert wait_entered.wait(TIMEOUT), "the reader did not wait for the waiter" + assert "reader-proceeded" not in recorder.labels() + + waiting = _thread(lambda: executor.claim_end_hook("child", waiter=True), "waiter") + _join(waiting) + _join(reading) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:parent", + "operation_end:child", + "reader-proceeded", + ] + assert recorder.thread_of("operation_end:child") == "waiter" + + +def test_a_reader_keeps_its_descendant_hook_from_an_ancestors_waiter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The reader decided under the claim lock; the ancestor's waiter claims next. + + The reader is held inside the claim lock, after it read the child's + terminal state and took the child's hook. The waiter's claim is made + only when the lock is released, so it finds the child owned and leaves + it: the child's hook runs on the reader, the parent's on the waiter. + """ + recorder = _Recorder() + state = _state(recorder, []) + executor = state._plugin_executor # noqa: SLF001 + state.fetch_paginated_operations( + [_operation("parent"), _operation("child", parent_id="parent")], + "token", + None, + before_publish=functools.partial( + executor.record_end_hooks, waited_for={"parent"} + ), + ) + decided = threading.Event() + proceed = threading.Event() + original_take = executor.take_end_hook + + def held_take(operation_id: str, *, waiter: bool = False) -> object: + claim = original_take(operation_id, waiter=waiter) + if threading.current_thread().name == "reader": + decided.set() + proceed.wait(TIMEOUT) + return claim + + monkeypatch.setattr(executor, "take_end_hook", held_take) + + def reader() -> None: + state.get_checkpoint_result("child") + recorder.record("reader-proceeded") + + reading = _thread(reader, "reader") + assert decided.wait(TIMEOUT) + waiting = _thread(lambda: executor.claim_end_hook("parent", waiter=True), "waiter") + proceed.set() + _join(reading) + _join(waiting) + + assert recorder.thread_of("operation_end:child") == "reader" + assert recorder.thread_of("operation_end:parent") == "waiter" + labels = recorder.labels() + assert labels.index("operation_end:child") < labels.index("reader-proceeded") + assert labels.count("operation_end:child") == 1 + executor.on_invocation_end(SUCCEEDED) + + +class _ObservedLock: + """Stands in for the claim lock; reports that a thread is acquiring it.""" + + def __init__(self, lock: threading.Lock, acquiring: threading.Event) -> None: + self._lock = lock + self._acquiring = acquiring + + def __enter__(self) -> bool: + self._acquiring.set() + return self._lock.__enter__() + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self._lock.__exit__(exc_type, exc, tb) + + +def test_a_reader_takes_the_claim_lock_before_it_reads_the_state() -> None: + """What makes the handoff above hold: the read itself is under the lock. + + With the lock held by the test, the reader must reach the lock without + having read the state. A read made before the lock would let another + thread's claim come between the read and the decision. + """ + recorder = _Recorder() + state = _state(recorder, []) + executor = state._plugin_executor # noqa: SLF001 + state.fetch_paginated_operations( + [_operation("a")], "token", None, before_publish=executor.record_end_hooks + ) + readers: list[str] = [] + + class _ObservedOperations(dict): # type: ignore[type-arg] + def get(self, key: str, default: object = None) -> object: + readers.append(threading.current_thread().name) + return super().get(key, default) + + state._operations = _ObservedOperations(state._operations) # noqa: SLF001 + acquiring = threading.Event() + real_lock = executor.claim_lock + executor.claim_lock = _ObservedLock(real_lock, acquiring) # type: ignore[assignment] + + with real_lock: + reading = _thread(lambda: state.get_checkpoint_result("a"), "reader") + assert acquiring.wait(TIMEOUT) + assert "reader" not in readers, "the state was read before the claim lock" + _join(reading) + + assert readers == ["reader"] + assert recorder.thread_of("operation_end:a") == "reader" + executor.claim_lock = real_lock + executor.on_invocation_end(SUCCEEDED) + + +def test_a_waiters_claim_is_decided_under_the_claim_lock() -> None: + """The other half of the handoff: the waiter's claim takes the same lock. + + With the lock held by the test, the waiter reaches the lock without having + run any hook. A claim decided without the lock could take a descendant + between its reader's read and that reader's decision. + """ + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks( + [_operation("parent"), _operation("child", parent_id="parent")], + waited_for={"parent"}, + ) + acquiring = threading.Event() + real_lock = executor.claim_lock + executor.claim_lock = _ObservedLock(real_lock, acquiring) # type: ignore[assignment] + + with real_lock: + waiting = _thread( + lambda: executor.claim_end_hook("parent", waiter=True), "waiter" + ) + assert acquiring.wait(TIMEOUT) + assert recorder.labels() == ["invocation_start"], "a hook ran without the lock" + _join(waiting) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:child", + "operation_end:parent", + ] + + +class _OwnerTrackingLock: + """Stands in for the claim lock; remembers which thread holds it, and numbers + its acquisitions so two steps can be shown to share one.""" + + def __init__(self, lock: threading.Lock) -> None: + self._lock = lock + self.owner: int | None = None + self.acquisition = 0 + + def __enter__(self) -> bool: + held = self._lock.__enter__() + self.owner = threading.get_ident() + self.acquisition += 1 + return held + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.owner = None + self._lock.__exit__(exc_type, exc, tb) + + +def test_every_claim_decision_is_made_while_the_deciding_thread_holds_the_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reader, waiter, replay path and drain: each decides inside the lock. + + Taking and releasing the lock before the decision would satisfy a test + that only observes acquisition; this one observes, at the decision + itself, that the deciding thread holds the lock. The reader's read of + the state and its decision must also fall in the same acquisition: two + acquisitions, one per step, would let another claim come between them. + """ + recorder = _Recorder() + state = _state(recorder, []) + executor = state._plugin_executor # noqa: SLF001 + lock = _OwnerTrackingLock(executor.claim_lock) + executor.claim_lock = lock # type: ignore[assignment] + decisions: list[tuple[str, bool, int]] = [] + reads: list[tuple[str, bool, int]] = [] + original_take = executor.take_end_hook + original_tree = executor._claim_tree_locked # noqa: SLF001 + + def observed_take(operation_id: str, *, waiter: bool = False) -> object: + decisions.append( + (operation_id, lock.owner == threading.get_ident(), lock.acquisition) + ) + return original_take(operation_id, waiter=waiter) + + def observed_tree(root: plugin_module._RecordedEndHook) -> object: # noqa: SLF001 + decisions.append( + ( + root.info.operation_id, + lock.owner == threading.get_ident(), + lock.acquisition, + ) + ) + return original_tree(root) + + class _ObservedOperations(dict): # type: ignore[type-arg] + def get(self, key: str, default: object = None) -> object: + reads.append((key, lock.owner == threading.get_ident(), lock.acquisition)) + return super().get(key, default) + + monkeypatch.setattr(executor, "take_end_hook", observed_take) + monkeypatch.setattr(executor, "_claim_tree_locked", observed_tree) + state.fetch_paginated_operations( + [_operation("waited"), _operation("read"), _operation("drained")], + "token", + None, + before_publish=functools.partial( + executor.record_end_hooks, waited_for={"waited"} + ), + ) + + state._operations = _ObservedOperations(state._operations) # noqa: SLF001 + + _join(_thread(lambda: state.get_checkpoint_result("read"), "reader")) + _join(_thread(lambda: executor.claim_end_hook("waited", waiter=True), "waiter")) + executor.on_operation_update(_operation("replayed")) + executor.drain_end_hooks() + + assert {operation_id for operation_id, _, _ in decisions} == { + "read", + "waited", + "replayed", + "drained", + } + assert all(under_lock for _, under_lock, _ in decisions), decisions + # The reader read the state under the lock, in the same acquisition as + # its decision. + assert len(reads) == 1, reads + read_key, read_under_lock, read_acquisition = reads[0] + assert (read_key, read_under_lock) == ("read", True) + assert {acq for op, _, acq in decisions if op == "read"} == {read_acquisition} + assert sorted(label for label in recorder.labels() if ":" in label) == [ + "operation_end:drained", + "operation_end:read", + "operation_end:replayed", + "operation_end:waited", + ] + + +def test_hooks_that_read_each_others_operations_do_not_wait_on_each_other() -> None: + """Two threads run the end hooks of a and b at the same time. + + From inside each hook the plugin reads the other operation, as a plugin + that calls a callback's ``result()`` from a hook would. Each hook's + completion event is set only when the hook returns, so were an in-hook + read to wait for the other hook, the two would wait on each other + forever. A read from inside a hook waits for nothing: both hooks return, + both reads see the terminal result, and each hook runs once. + """ + both_in_hook = threading.Barrier(2, timeout=TIMEOUT) + seen_from_hook: dict[str, bool] = {} + holder: dict[str, ExecutionState] = {} + + class _CrossReading(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id in ("a", "b"): + both_in_hook.wait() + other = "b" if info.operation_id == "a" else "a" + seen_from_hook[info.operation_id] = ( + holder["state"].get_checkpoint_result(other).is_succeeded() + ) + self.record("hook-returned:" + info.operation_id) + + recorder = _CrossReading() + state = _state(recorder, []) + holder["state"] = state + state._plugin_executor.record_end_hooks([_operation("a"), _operation("b")]) # noqa: SLF001 + state._operations.update(_map(_operation("a"), _operation("b"))) # noqa: SLF001 + + reader_a = _thread(lambda: state.get_checkpoint_result("a"), "reader-a") + reader_b = _thread(lambda: state.get_checkpoint_result("b"), "reader-b") + _join(reader_a) + _join(reader_b) + + assert seen_from_hook == {"a": True, "b": True} + labels = recorder.labels() + assert labels.count("operation_end:a") == 1 + assert labels.count("operation_end:b") == 1 + assert recorder.thread_of("operation_end:a") == "reader-a" + assert recorder.thread_of("operation_end:b") == "reader-b" + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + assert recorder.labels().count("operation_end:a") == 1 + + +def test_a_change_hook_that_holds_a_lock_and_reads_a_result_does_not_deadlock() -> None: + """A plugin guards its hooks with one lock, as the contract tells it to. + + From on_operation_change, holding the lock, it reads a callback's result + through the call ``Callback.result()`` makes. The callback's end hook is + recorded and no thread has taken it. Were the read to claim the hook, the + end hook would run nested in the change hook, on the change thread, and + block on the lock that thread holds: a deadlock, and the read would + never complete. A read from inside any hook claims nothing, so the read + completes and the change hook returns; the end hook runs on the handler + thread at the drain, once, before the invocation-end hook. + + The invocation ends only after the read has happened, so the drain + cannot run the end hook first and leave the read nothing to claim. + """ + lock = threading.Lock() + read_done = threading.Event() + holder: dict[str, ExecutionState] = {} + seen: dict[str, bool] = {} + + class _Locking(_Recorder): + def on_operation_change(self, info: OperationChangeInfo) -> None: + super().on_operation_change(info) + with lock: + seen["cb"] = holder["state"].get_checkpoint_result("cb").is_succeeded() + read_done.set() + + def on_operation_end(self, info: OperationEndInfo) -> None: + with lock: + super().on_operation_end(info) + + recorder = _Locking() + state = _state(recorder, []) + holder["state"] = state + cb = _operation("cb") + executor = state._plugin_executor # noqa: SLF001 + executor.record_end_hooks([cb]) + state._operations.update(_map(cb)) # noqa: SLF001 + executor.record_change([cb], _map(cb), {}) + + assert read_done.wait(TIMEOUT), "the read from inside the change hook hung" + assert seen == {"cb": True} + assert "operation_end:cb" not in recorder.labels() + + _join(_end_in_thread(executor)) + + labels = recorder.labels() + assert labels.count("operation_end:cb") == 1 + assert labels.index("operation_end:cb") < labels.index("invocation_end") + assert recorder.thread_of("operation_end:cb") == "handler" + + +def test_a_propagation_metadata_provider_that_reads_a_result_runs_no_hook() -> None: + """``provide_propagation_metadata`` is the other entry into plugin code. + + A provider that reads a terminal operation whose hook is recorded and + untaken runs no hook, like a read from any hook; the hook stays for the + consuming read. + """ + holder: dict[str, ExecutionState] = {} + + class _Reading(_Recorder): + def provide_propagation_metadata( + self, info: PropagationInput + ) -> PropagationMetadata | None: + assert holder["state"].get_checkpoint_result("a").is_succeeded() + self.record("read-a-from-provider") + return None + + recorder = _Reading() + state = _state(recorder, []) + holder["state"] = state + executor = state._plugin_executor # noqa: SLF001 + executor.record_end_hooks([_operation("a")]) + state._operations.update(_map(_operation("a"))) # noqa: SLF001 + + executor.provide_propagation_metadata( + PropagationInput( + execution_arn="arn:exec", operation_id="invoke", target_function_name="f" + ) + ) + labels = recorder.labels() + assert "read-a-from-provider" in labels + assert "operation_end:a" not in labels + + assert state.get_checkpoint_result("a").is_succeeded() + assert recorder.labels().count("operation_end:a") == 1 + executor.on_invocation_end(SUCCEEDED) + assert recorder.labels().count("operation_end:a") == 1 + + +def test_a_read_from_inside_an_end_hook_runs_no_hook() -> None: + """A hook reads an operation whose recorded hook no thread has taken. + + The read inspects the result and runs nothing: nesting one hook in + another is not done. The hook stays recorded, and the thread that later + consumes the result runs it, before it acts on the result. + """ + holder: dict[str, ExecutionState] = {} + + class _ReadingOther(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id == "a": + assert holder["state"].get_checkpoint_result("c").is_succeeded() + self.record("read-c-from-hook-a") + + recorder = _ReadingOther() + state = _state(recorder, []) + holder["state"] = state + state._plugin_executor.record_end_hooks([_operation("a"), _operation("c")]) # noqa: SLF001 + state._operations.update(_map(_operation("a"), _operation("c"))) # noqa: SLF001 + + _join(_thread(lambda: state.get_checkpoint_result("a"), "reader-a")) + labels = recorder.labels() + assert "operation_end:c" not in labels + assert labels.index("read-c-from-hook-a") > labels.index("operation_end:a") + + def consumer() -> None: + assert state.get_checkpoint_result("c").is_succeeded() + recorder.record("consumed:c") + + _join(_thread(consumer, "reader-c")) + labels = recorder.labels() + assert labels.index("operation_end:c") < labels.index("consumed:c") + assert recorder.thread_of("operation_end:c") == "reader-c" + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + assert recorder.labels().count("operation_end:c") == 1 + + +def test_a_virtual_child_contexts_end_hook_reads_without_waiting() -> None: + """The end hook of a context that wrote no checkpoint is an end hook too. + + It runs on the context's thread as its body returns, through a different + dispatch than the registry's. A read from inside it of an operation whose + hook another thread is running returns at once, like any read from + inside an end hook; otherwise that hook could wait on this one. + """ + hook_a_running = threading.Event() + release_hook_a = threading.Event() + holder: dict[str, ExecutionState] = {} + + class _ReadingFromChildEnd(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id == "a": + hook_a_running.set() + assert release_hook_a.wait(TIMEOUT) + elif info.operation_id == "virtual": + assert holder["state"].get_checkpoint_result("a").is_succeeded() + self.record("read-a-from-virtual-end-hook") + + recorder = _ReadingFromChildEnd() + state = _state(recorder, []) + holder["state"] = state + state._plugin_executor.record_end_hooks([_operation("a")]) # noqa: SLF001 + state._operations.update(_map(_operation("a"))) # noqa: SLF001 + + reader_a = _thread(lambda: state.get_checkpoint_result("a"), "reader-a") + assert hook_a_running.wait(TIMEOUT) + virtual = OperationIdentifier( + operation_id="virtual", + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + parent_id=None, + name="virtual", + operation_type=ServiceOperationType.CONTEXT, + ) + _join( + _thread( + lambda: state.emit_child_context_end_hook( + virtual, OperationStatus.SUCCEEDED + ), + "virtual", + ) + ) + # The virtual context's hook returned while hook a is still running. + assert "read-a-from-virtual-end-hook" in recorder.labels() + assert reader_a.is_alive() + release_hook_a.set() + _join(reader_a) + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + assert recorder.labels().count("operation_end:a") == 1 + + +def test_a_hook_may_read_an_operation_whose_hook_its_thread_runs_next() -> None: + """Within one taken tree, the running thread owns every entry: a read of a + sibling or ancestor whose hook is about to run on this thread returns at + once instead of waiting for itself.""" + + class _ReadsParent(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id == "child": + executor.claim_end_hook("parent") + self.record("read-parent-from-child-hook") + + recorder = _ReadsParent() + executor = _started_executor(recorder) + executor.record_end_hooks( + [_operation("parent"), _operation("child", parent_id="parent")] + ) + + _join(_thread(lambda: executor.claim_end_hook("parent"), "claimer")) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:child", + "read-parent-from-child-hook", + "operation_end:parent", + ] + + +def test_a_reserved_hook_whose_waiter_never_came_is_drained() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")], waited_for={"a"}) + + executor.drain_end_hooks() + + assert recorder.labels() == ["invocation_start", "operation_end:a"] + + +def test_a_claim_runs_unread_descendants_first() -> None: + """A response completes a context and a callback inside it nobody awaited. + + The thread that observes the context's result runs the callback's end + hook first, then the context's: a child ends before its parent, as it + would have on the thread that read the callback. + """ + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks( + [ + _operation("context"), + _operation("callback", parent_id="context"), + _operation("grandchild", parent_id="callback"), + _operation("unrelated"), + ] + ) + + executor.claim_end_hook("context") + + assert recorder.labels() == [ + "invocation_start", + "operation_end:grandchild", + "operation_end:callback", + "operation_end:context", + ] + executor.drain_end_hooks() + assert recorder.labels()[-1] == "operation_end:unrelated" + + +def test_drain_runs_descendants_before_ancestors() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + executor.record_end_hooks( + [_operation("parent"), _operation("child", parent_id="parent")] + ) + + executor.drain_end_hooks() + + assert recorder.labels() == [ + "invocation_start", + "operation_end:child", + "operation_end:parent", + ] + + +# endregion + +# region Change dispatcher + + +def test_change_notifications_run_in_order_on_the_change_thread() -> None: + recorder = _Recorder() + executor = _started_executor(recorder) + a, b, c = _operation("a"), _operation("b"), _operation("c") + + executor.record_change([a], _map(a), {}) + executor.record_change([b], _map(a, b), _map(a)) + executor.record_change([c], _map(a, b, c), _map(a, b)) + _join(_end_in_thread(executor)) + + changes = recorder.changes + # Each response reports a different operation, so with or without merging + # every id is delivered once, and in response order. + delivered = [key for change in changes for key in change.updated_operations] + assert delivered == ["a", "b", "c"] + assert changes[-1].operations.keys() == {"a", "b", "c"} + for label in recorder.labels(): + if label.startswith("operation_change"): + assert recorder.thread_of(label) == "plugin-change" + assert recorder.labels()[-1] == "invocation_end" + + +def test_notifications_arriving_during_a_slow_hook_merge_into_one() -> None: + """Two responses arrive while the plugin is in the hook: one delivery covers both. + + The second response reports b (started) and d; the third reports c, and + its map holds b succeeded. The merged delivery reports b, c and d, with b + in the newest state, and carries the newest map. Replacing the pending + notification would lose d; keeping the pending states would report b as + started. + """ + entered = threading.Event() + release = threading.Event() + + class _SlowOnFirst(_Recorder): + def on_operation_change(self, info: OperationChangeInfo) -> None: + super().on_operation_change(info) + if len(self.changes) == 1: + entered.set() + self.block_on(release) + + recorder = _SlowOnFirst() + executor = _started_executor(recorder) + a, c, d = _operation("a"), _operation("c"), _operation("d") + b_started = _operation("b", OperationStatus.STARTED) + b_done = _operation("b") + + executor.record_change([a], _map(a), {}) + assert entered.wait(TIMEOUT) + executor.record_change([b_started, d], _map(a, b_started, d), _map(a)) + executor.record_change([c], _map(a, b_done, c, d), _map(a, b_started, d)) + release.set() + _join(_end_in_thread(executor)) + + assert "timed-out" not in recorder.labels() + assert [sorted(change.updated_operations) for change in recorder.changes] == [ + ["a"], + ["b", "c", "d"], + ] + merged = recorder.changes[1] + assert merged.updated_operations["b"].status is OperationStatus.SUCCEEDED + assert merged.operations.keys() == {"a", "b", "c", "d"} + assert merged.operations["b"].status is OperationStatus.SUCCEEDED + + +def test_merge_takes_the_union_of_ids_in_the_newest_state() -> None: + def info(operation: Operation) -> OperationInfo: + return OperationInfo.from_operation(operation) + + a, b_old, b_new, c = ( + _operation("a"), + _operation("b", OperationStatus.STARTED), + _operation("b"), + _operation("c"), + ) + pending = OperationChangeInfo( + execution_arn="arn:exec", + updated_operations={"a": info(a), "b": info(b_old)}, + operations={"a": info(a), "b": info(b_old)}, + ) + latest = OperationChangeInfo( + execution_arn="arn:exec", + updated_operations={"c": info(c)}, + operations={"a": info(a), "b": info(b_new), "c": info(c)}, + ) + + merged = _merge_change_infos(pending, latest) + + assert list(merged.updated_operations) == ["a", "b", "c"] + assert merged.updated_operations["b"].status is OperationStatus.SUCCEEDED + assert merged.operations is latest.operations + assert merged.execution_arn == "arn:exec" + + +class _ObservedThread: + """Stands in for the dispatcher thread; reports that join() was entered.""" + + def __init__(self, thread: threading.Thread, entered: threading.Event) -> None: + self._thread = thread + self._entered = entered + + def join(self, timeout: float | None = None) -> None: + self._entered.set() + self._thread.join(timeout) + + +def test_on_invocation_end_waits_for_a_notification_still_in_the_hook() -> None: + """The readiness signal comes from inside the dispatcher thread's join, so + the negative assertion is made while the ending thread is in that join, + and an implementation that does not join never signals. + """ + entered = threading.Event() + release = threading.Event() + join_entered = threading.Event() + + class _Slow(_Recorder): + def on_operation_change(self, info: OperationChangeInfo) -> None: + super().on_operation_change(info) + entered.set() + self.block_on(release) + self.record("change-left") + + recorder = _Slow() + executor = _started_executor(recorder) + a = _operation("a") + executor.record_change([a], _map(a), {}) + assert entered.wait(TIMEOUT) + dispatcher = executor._dispatcher # noqa: SLF001 + assert dispatcher is not None + dispatcher._thread = _ObservedThread(dispatcher._thread, join_entered) # type: ignore[assignment] # noqa: SLF001 + + ending = _end_in_thread(executor) + assert join_entered.wait(TIMEOUT), "on_invocation_end did not join the dispatcher" + assert "invocation_end" not in recorder.labels() + release.set() + _join(ending) + + assert recorder.labels() == [ + "invocation_start", + "operation_change:a", + "change-left", + "invocation_end", + ] + + +def test_no_change_thread_starts_without_a_notification( + monkeypatch: pytest.MonkeyPatch, +) -> None: + starts: list[str] = [] + original = plugin_module._ChangeDispatcher.start # noqa: SLF001 + + def counting_start(executor: PluginExecutor) -> plugin_module._ChangeDispatcher: # noqa: SLF001 + starts.append("started") + return original(executor) + + monkeypatch.setattr(plugin_module._ChangeDispatcher, "start", counting_start) # noqa: SLF001 + recorder = _Recorder() + executor = _started_executor(recorder) + a = _operation("a") + + executor.record_end_hooks([a]) + # No status changed between the maps: no notification, no thread. + executor.record_change([a], _map(a), _map(a)) + executor.on_invocation_end(SUCCEEDED) + + assert starts == [] + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "invocation_end", + ] + + +def test_notifications_run_once_on_the_handler_thread_when_no_thread_can_start( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + def cannot_start(executor: PluginExecutor) -> plugin_module._ChangeDispatcher: # noqa: ARG001, SLF001 + msg = "can't start new thread" + raise RuntimeError(msg) + + monkeypatch.setattr(plugin_module._ChangeDispatcher, "start", cannot_start) # noqa: SLF001 + recorder = _Recorder() + executor = _started_executor(recorder) + a, b = _operation("a"), _operation("b") + + with caplog.at_level(logging.WARNING, logger=plugin_module.__name__): + executor.record_change([a], _map(a), {}) + executor.record_change([b], _map(a, b), _map(a)) + assert recorder.labels() == ["invocation_start"] + _join(_end_in_thread(executor)) + + assert recorder.labels() == [ + "invocation_start", + "operation_change:a,b", + "invocation_end", + ] + assert recorder.thread_of("operation_change:a,b") == "handler" + assert "could not be started" in caplog.text + + +def test_a_control_exception_on_the_change_thread_is_logged_and_the_next_runs( + caplog: pytest.LogCaptureFixture, +) -> None: + class _Interrupting(_Recorder): + def on_operation_change(self, info: OperationChangeInfo) -> None: + super().on_operation_change(info) + if "a" in info.updated_operations: + raise KeyboardInterrupt + + recorder = _Interrupting() + executor = _started_executor(recorder) + a, b = _operation("a"), _operation("b") + + with caplog.at_level(logging.ERROR, logger=plugin_module.__name__): + executor.record_change([a], _map(a), {}) + assert _wait_until(lambda: "operation_change:a" in recorder.labels()) + executor.record_change([b], _map(a, b), _map(a)) + executor.on_invocation_end(SUCCEEDED) + + assert recorder.labels() == [ + "invocation_start", + "operation_change:a", + "operation_change:b", + "invocation_end", + ] + assert "plugin change thread; ignored" in caplog.text + + +def test_a_control_exception_from_one_plugin_does_not_cost_the_others_the_notification() -> ( + None +): + class _Interrupting(_Recorder): + def on_operation_change(self, info: OperationChangeInfo) -> None: + super().on_operation_change(info) + raise KeyboardInterrupt + + interrupting = _Interrupting() + healthy = _Recorder() + executor = PluginExecutor([plugin_factory(interrupting), plugin_factory(healthy)]) + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=None, + execution_start_time=START_TS, + is_first_invocation=True, + ) + a = _operation("a") + + executor.record_change([a], _map(a), {}) + executor.on_invocation_end(SUCCEEDED) + + assert healthy.labels() == [ + "invocation_start", + "operation_change:a", + "invocation_end", + ] + assert interrupting.labels()[1] == "operation_change:a" + + +def test_a_control_exception_from_a_drained_end_hook_is_raised_after_the_end_hook() -> ( + None +): + class _Interrupting(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + raise KeyboardInterrupt + + recorder = _Interrupting() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + with pytest.raises(KeyboardInterrupt): + executor.on_invocation_end(SUCCEEDED) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "invocation_end", + ] + + +def test_a_control_exception_from_one_drained_tree_does_not_stop_the_others() -> None: + """The other trees still run before the end hook; the exception follows it. + + Were the drain to stop at the first, the executor's scope exit would + drain the rest after on_invocation_end. + """ + + class _Interrupting(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + if info.operation_id == "a": + raise KeyboardInterrupt + + recorder = _Interrupting() + executor = PluginExecutor([plugin_factory(recorder)]) + with pytest.raises(KeyboardInterrupt), executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=None, + execution_start_time=START_TS, + is_first_invocation=True, + ) + executor.record_end_hooks([_operation("a"), _operation("b")]) + executor.on_invocation_end(SUCCEEDED) + + assert recorder.labels() == [ + "invocation_start", + "operation_end:a", + "operation_end:b", + "invocation_end", + ] + + +def test_an_ordinary_exception_in_a_claimed_hook_is_contained( + caplog: pytest.LogCaptureFixture, +) -> None: + class _Failing(_Recorder): + def on_operation_end(self, info: OperationEndInfo) -> None: + super().on_operation_end(info) + msg = "plugin bug" + raise ValueError(msg) + + recorder = _Failing() + executor = _started_executor(recorder) + executor.record_end_hooks([_operation("a")]) + + with caplog.at_level(logging.ERROR, logger=plugin_module.__name__): + executor.claim_end_hook("a") + + assert recorder.labels() == ["invocation_start", "operation_end:a"] + assert "exception ignored" in caplog.text + + +# endregion + +# region ExecutionState: record before publish, claim on observe + + +def _state( + recorder: _Recorder, + response: list[Operation], + failure: Exception | None = None, +) -> ExecutionState: + client = Mock(spec=LambdaClient) + if failure is not None: + client.checkpoint.side_effect = failure + else: + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="next", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=response, next_marker=None + ), + ) + return ExecutionState( + durable_execution_arn="arn:exec", + initial_checkpoint_token="token", # noqa: S106 + operations={}, + service_client=client, + plugin_executor=_started_executor(recorder), + ) + + +def test_fetch_records_before_it_publishes_the_operations() -> None: + recorder = _Recorder() + state = _state(recorder, []) + seen: list[tuple[list[str], bool, bool]] = [] + + def before_publish(operations: list[Operation]) -> None: + # Called before the operations lock is taken: the map is read directly, + # and the lock is free. Recording takes the plugin executor's claim + # lock, and a reader takes the claim lock before the operations lock, + # so recording under the operations lock would invert that order. + lock_free = state._operations_lock.acquire(blocking=False) # noqa: SLF001 + if lock_free: + state._operations_lock.release() # noqa: SLF001 + seen.append( + ( + [op.operation_id for op in operations], + "a" in state._operations, + lock_free, + ) # noqa: SLF001 + ) + + state.fetch_paginated_operations( + [_operation("a")], "token", None, before_publish=before_publish + ) + + assert seen == [(["a"], False, True)] + assert state.get_checkpoint_result("a").is_succeeded() + + +def _run_loop(state: ExecutionState) -> threading.Thread: + return _thread(state.checkpoint_batches_forever, "checkpoint") + + +def test_a_synchronous_waiter_runs_its_end_hook_before_create_checkpoint_returns() -> ( + None +): + """The checkpoint thread records, wakes the waiter, and runs nothing itself.""" + recorder = _Recorder() + state = _state(recorder, [_operation("a")]) + loop = _run_loop(state) + + def waiter() -> None: + state.create_checkpoint(_update("a"), is_sync=True) + recorder.record("returned:a") + + _join(_thread(waiter, "waiter")) + state.stop_checkpointing() + _join(loop) + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + + labels = recorder.labels() + assert labels.index("operation_end:a") < labels.index("returned:a") + assert recorder.thread_of("operation_end:a") == "waiter" + assert recorder.thread_of("operation_change:a") == "plugin-change" + assert all(thread != "checkpoint" for _, thread in recorder.events) + assert labels.count("operation_end:a") == 1 + + +def test_the_end_hook_is_recorded_before_the_waiter_is_signalled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Observed on the checkpoint thread, at the moment it sets the waiter's event. + + A waiter woken before the record would claim nothing, and its hook would + run late, at the end of the invocation; the timing makes that rare, so the + order is asserted at the signal itself rather than through the waiter. + """ + recorder = _Recorder() + state = _state(recorder, [_operation("a")]) + executor = state._plugin_executor # noqa: SLF001 + recorded_at_signal: list[bool] = [] + + class _ObservingEvent(CompletionEvent): + def set(self, error: BaseException | None = None) -> None: + with executor.claim_lock: + recorded_at_signal.append("a" in executor._end_hooks) # noqa: SLF001 + super().set(error) + + monkeypatch.setattr(state_module, "CompletionEvent", _ObservingEvent) + loop = _run_loop(state) + _join(_thread(lambda: state.create_checkpoint(_update("a"), True), "waiter")) + state.stop_checkpointing() + _join(loop) + + assert recorded_at_signal == [True] + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + + +def test_a_reader_of_a_result_another_response_completed_runs_its_end_hook() -> None: + """A wait that elapsed, reported in another operation's response. + + The thread that later reads the result claims the hook in + get_checkpoint_result, so the hook precedes anything that thread does + with the result, and the end of the invocation finds nothing to drain. + """ + recorder = _Recorder() + state = _state(recorder, [_operation("a"), _operation("elapsed")]) + loop = _run_loop(state) + + _join(_thread(lambda: state.create_checkpoint(_update("a"), True), "waiter")) + + def reader() -> None: + assert state.get_checkpoint_result("elapsed").is_succeeded() + recorder.record("read:elapsed") + + _join(_thread(reader, "reader")) + state.stop_checkpointing() + _join(loop) + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + + labels = recorder.labels() + assert labels.index("operation_end:elapsed") < labels.index("read:elapsed") + assert recorder.thread_of("operation_end:elapsed") == "reader" + assert recorder.thread_of("operation_end:a") == "waiter" + assert labels.count("operation_end:elapsed") == 1 + + +def test_the_replay_boundary_reads_take_no_end_hook() -> None: + """DurableContext reads a checkpoint twice without consuming its result. + + While replaying, it reads the next operation's checkpoint to place the + replay boundary, and an operation's own checkpoint to validate its + identity before the operation runs. Neither read runs the end hook a + checkpoint response recorded for the operation; the operation's own + read, which consumes the result, does. So the hook neither delays the + operation before it nor runs before the identity check. + """ + recorder = _Recorder() + context = DurableContext( + state=_state(recorder, []), + execution_context=ExecutionContext(durable_execution_arn="arn:exec"), + replay_status=ReplayStatus.REPLAY, + ) + state = context.state + # The context's logical steps 1 and 2. Step 1 completed in a prior + # invocation; step 2 completed during this one, reported by a response. + first_id = context._create_step_id() # noqa: SLF001 + second_id = context._peek_next_operation_id() # noqa: SLF001 + state._operations[first_id] = _operation(first_id) # noqa: SLF001 + state._plugin_executor.record_end_hooks([_operation(second_id)]) # noqa: SLF001 + state._operations[second_id] = _operation(second_id) # noqa: SLF001 + + # The boundary read at step 1's exit sees step 2 complete, and runs nothing. + assert context._next_operation_exists() # noqa: SLF001 + assert "operation_end:" + second_id not in recorder.labels() + + # Step 2's identity read runs nothing either. + second = OperationIdentifier( + operation_id=context._create_step_id(), # noqa: SLF001 + sub_type=OperationSubType.STEP, + parent_id=None, + name=second_id, + operation_type=ServiceOperationType.STEP, + ) + assert second.operation_id == second_id + with context._replay_aware(second): # noqa: SLF001 + assert "operation_end:" + second_id not in recorder.labels() + assert "operation_end:" + second_id not in recorder.labels() + + # The read that consumes the result runs the hook, on this thread. + assert state.get_checkpoint_result(second_id).is_succeeded() + assert recorder.labels().count("operation_end:" + second_id) == 1 + assert ( + recorder.thread_of("operation_end:" + second_id) + == threading.current_thread().name + ) + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + assert recorder.labels().count("operation_end:" + second_id) == 1 + + +def test_a_result_no_thread_reads_gets_its_end_hook_at_invocation_end() -> None: + recorder = _Recorder() + state = _state(recorder, [_operation("a"), _operation("unread")]) + loop = _run_loop(state) + + _join(_thread(lambda: state.create_checkpoint(_update("a"), True), "waiter")) + state.stop_checkpointing() + _join(loop) + _join(_end_in_thread(state._plugin_executor)) # noqa: SLF001 + + labels = recorder.labels() + assert labels.index("operation_end:unread") < labels.index("invocation_end") + assert recorder.thread_of("operation_end:unread") == "handler" + assert labels.count("operation_end:unread") == 1 + + +def test_the_loop_records_end_hooks_before_it_publishes_the_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The loop's own wiring: record_end_hooks runs while the operation is not yet readable.""" + recorder = _Recorder() + state = _state(recorder, [_operation("a")]) + executor = state._plugin_executor # noqa: SLF001 + published_at_record: list[bool] = [] + waited: list[set[str]] = [] + original = executor.record_end_hooks + + def observing_record( + operations: list[Operation], waited_for: Collection[str] = () + ) -> None: + # Called before the operations lock is taken, so the map is read directly. + published_at_record.append("a" in state._operations) # noqa: SLF001 + waited.append(set(waited_for)) + original(operations, waited_for=waited_for) + + monkeypatch.setattr(executor, "record_end_hooks", observing_record) + loop = _run_loop(state) + _join(_thread(lambda: state.create_checkpoint(_update("a"), True), "waiter")) + state.stop_checkpointing() + _join(loop) + + assert published_at_record == [False] + assert recorder.thread_of("operation_end:a") == "waiter" + assert waited == [{"a"}] + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + + +def test_a_read_of_the_older_non_terminal_state_runs_no_hook() -> None: + """The hook of a completion the reader did not see stays for the reader that does.""" + recorder = _Recorder() + state = _state(recorder, []) + state.fetch_paginated_operations( + [_operation("a", OperationStatus.STARTED)], "token", None + ) + executor = state._plugin_executor # noqa: SLF001 + # The response completing "a" was recorded but not yet published. + executor.record_end_hooks([_operation("a")]) + + assert not state.get_checkpoint_result("a").is_terminal() + assert recorder.labels() == ["invocation_start"] + + state.fetch_paginated_operations([_operation("a")], "token", None) + assert state.get_checkpoint_result("a").is_succeeded() + assert recorder.labels() == ["invocation_start", "operation_end:a"] + + +def test_a_waiter_woken_by_a_failed_later_page_still_runs_its_recorded_hook() -> None: + """The first page completes the operation; fetching the next page fails. + + The pages already fetched are recorded and published, then the waiter is + woken with BackgroundThreadError. The claim is in a finally, so the hook + runs on the waiter thread before the error propagates, in the order a + successful batch would have given it. + """ + recorder = _Recorder() + client = Mock(spec=LambdaClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="next", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=[_operation("a")], next_marker="page-2" + ), + ) + client.get_execution_state.side_effect = GetExecutionStateError("page fetch failed") + state = ExecutionState( + durable_execution_arn="arn:exec", + initial_checkpoint_token="token", # noqa: S106 + operations={}, + service_client=client, + plugin_executor=_started_executor(recorder), + ) + loop = _run_loop(state) + outcome: list[str] = [] + + def waiter() -> None: + try: + state.create_checkpoint(_update("a"), is_sync=True) + except BaseException as error: # noqa: BLE001 - BackgroundThreadError is a BaseException + outcome.append(type(error).__name__) + recorder.record("raised:a") + + _join(_thread(waiter, "waiter")) + _join(loop) + + assert outcome == ["BackgroundThreadError"] + labels = recorder.labels() + assert labels.index("operation_end:a") < labels.index("raised:a") + assert recorder.thread_of("operation_end:a") == "waiter" + state._plugin_executor.on_invocation_end(SUCCEEDED) # noqa: SLF001 + + +def test_a_waiter_whose_checkpoint_failed_runs_no_hook() -> None: + recorder = _Recorder() + state = _state(recorder, [], failure=RuntimeError("service down")) + loop = _run_loop(state) + outcome: list[str] = [] + + def waiter() -> None: + try: + state.create_checkpoint(_update("a"), is_sync=True) + except BaseException as error: # noqa: BLE001 - BackgroundThreadError is a BaseException + outcome.append(type(error).__name__) + + _join(_thread(waiter, "waiter")) + _join(loop) + + assert outcome == ["BackgroundThreadError"] + assert recorder.labels() == ["invocation_start"] + + +# endregion