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