From fb8af739a92f7fd8070d928029150fe4e5c9cc61 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Wed, 16 Sep 2026 17:14:49 -0700 Subject: [PATCH 01/28] fix(insight): key export scheduling by execution ARN One plugin instance serves every execution its environment hosts, and Lambda Managed Instances makes concurrent executions in one environment routine. The export scheduler held one pending record for the whole plugin and overwrote it regardless of which execution the record belonged to. Measured before the change, 20 trials per case: 2 concurrent executions lost a terminal record in 16 of 20 trials, 10 concurrent lost 8.5 per trial, and with 5 concurrent executions and a 200 ms exporter only 1 of 5 terminal records was exported while every drain() still returned without error. Pending records are now keyed by execution ARN with a per-execution lane, so coalescing happens only within one execution and drain(execution_arn) returns only once that execution's own record has reached every exporter and a flush covering it has completed. One worker thread and one export at a time are unchanged, so an exporter never sees concurrent calls. Also fixed, each found while reviewing the change above: - A BaseException from a customer exporter (asyncio.CancelledError inherits from it, so an exporter touching asyncio can raise it without writing raise) killed the worker between consuming a record and publishing its bookkeeping. The worker slot stayed occupied by a dying thread, no replacement started, and the parked drain hung the invocation thread permanently. - A drain could request a second flush while one was already in flight, so a flush ran after drain() had returned, calling an exporter after the invocation went back to Lambda. - The per-execution lock was not reentrant while a displaced record was released under it, so a record whose finalizer re-entered a hook for the same execution self-deadlocked the invocation thread. - The closed gate was a check-then-act: customer code running under the reentrant lock could complete on_invocation_end on the same thread, after which the outer frame still scheduled its RUNNING record behind the terminal one. - A flush happened only when a record was emitted, where JS and Java flush once per sampled-in invocation end. A buffering exporter now sees the same rhythm in all three languages. - InsightExporter.flush had no docstring. It now states the cadence, the exclusivity guarantee, that a flush may cover other executions' records, what happens when an exporter omits the method, and how failures are handled. No public API changed. Record fields, emit modes, sampling, truncation and the default exporter are unchanged; the emitted record surface is byte-identical to before the change. --- .../_export_scheduler.py | 303 +++++++++-- .../plugin.py | 228 +++++--- .../types.py | 35 +- .../tests/test_export_scheduler.py | 435 +++++++++++++++- .../tests/test_plugin.py | 492 ++++++++++++++++++ 5 files changed, 1361 insertions(+), 132 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index ef22ea54..b06a2853 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -1,7 +1,17 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -"""Latest-pending asynchronous export scheduling for Workflow Insight.""" +"""Per-execution latest-pending asynchronous export scheduling for Workflow Insight. + +One plugin instance serves every execution its environment hosts, and Lambda +Managed Instances makes concurrent executions in one environment routine, so the +pending record is keyed by execution ARN: coalescing happens only within a single +execution and one execution's record can never displace another's. + +Export itself stays strictly serialized -- one worker thread, one ``export()`` at +a time -- so exporters never see concurrent calls. Parallel export is a later +phase and a contract change. +""" from __future__ import annotations @@ -16,30 +26,81 @@ _logger = logging.getLogger("aws_durable_execution_sdk_python_insight") +class _Lane: + """Per-execution export bookkeeping. One lane per execution ARN.""" + + __slots__ = ("scheduled_seq", "exported_seq", "exported_at", "waiters") + + def __init__(self) -> None: + # Newest sequence number scheduled for this execution. + self.scheduled_seq = 0 + # Newest sequence number already handed to every exporter. + self.exported_seq = 0 + # Value of the scheduler's export counter when that export finished, so + # a waiter can tell whether a completed flush covered its own record. + self.exported_at = 0 + # drain() calls currently blocked on this lane; the lane is only + # forgotten once nobody is waiting on it. + self.waiters = 0 + + class _ExportScheduler: - """Run all exporters on one lazy worker with one latest pending record.""" + """Run all exporters on one lazy worker, keeping the latest record per execution.""" def __init__(self, exporters: list[InsightExporter]) -> None: self._exporters = exporters self._condition = threading.Condition(threading.Lock()) - self._pending: dict[str, Any] | None = None + # execution ARN -> (sequence, latest record), oldest arrival first. A + # repeat schedule for an ARN replaces the value and keeps the position, + # so coalescing never lets one execution jump the queue. + self._pending: dict[str, tuple[int, dict[str, Any]]] = {} + self._lanes: dict[str, _Lane] = {} + self._seq = 0 + self._export_count = 0 + # Highest export counter value covered by a completed flush. + self._flushed_through = 0 + # Completed flushes, monotonic. Export coverage alone cannot express + # "a flush ran for this invocation end": a drain with nothing of its own + # to export -- an invocation end that emitted no record -- is trivially + # covered by an older flush, so it would return without flushing at all. + # JS and Java flush once per sampled-in invocation end whether or not a + # record was emitted, so a drain also requires a flush that COMPLETED + # AFTER it was called. Concurrent drains still share one flush: they all + # entered before it completed. + self._flushes_completed = 0 self._flush_requested = False - self._flush_event: threading.Event | None = None + # Export counter coverage of the flush the worker is running right now, or + # 0 when no flush is in flight. Published when the worker commits to a + # flush, so a waiter woken while that flush runs -- before its coverage + # reaches _flushed_through -- can tell it is already covered instead of + # requesting a second flush that would run after its drain returned. + self._flush_in_flight = 0 + # Value of the global schedule counter (_seq) when a flush was requested. + # The worker defers the flush until no record scheduled at or before that + # point is still pending. That is deliberately wider than the requester's + # own record: a drain therefore also waits for records other executions + # had pending when it was called. It excludes records scheduled after the + # request, so a steady stream of other executions cannot starve a waiting + # drain. + self._flush_barrier = 0 self._worker: threading.Thread | None = None self._disabled = False - def schedule(self, record: dict[str, Any]) -> None: - """Replace the pending snapshot and return without running exporters.""" - displaced: dict[str, Any] | None = None - failed_pending: dict[str, Any] | None = None + def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: + """Replace this execution's pending snapshot; never runs exporters inline.""" + displaced: tuple[int, dict[str, Any]] | None = None + failed_pending: dict[str, tuple[int, dict[str, Any]]] | None = None start_error: Exception | None = None with self._condition: if self._disabled: return - displaced = self._pending - self._pending = record + self._seq += 1 + lane = self._lane_locked(execution_arn) + lane.scheduled_seq = self._seq + displaced = self._pending.get(execution_arn) + self._pending[execution_arn] = (self._seq, record) failed_pending, start_error = self._ensure_worker_locked() - self._condition.notify() + self._condition.notify_all() # Releasing either record may run custom finalizers, so do it unlocked. del displaced, failed_pending if start_error is not None: @@ -49,21 +110,80 @@ def schedule(self, record: dict[str, Any]) -> None: start_error, ) - def drain(self) -> None: - """Wait until the latest pending record is exported and exporters flush.""" - failed_pending: dict[str, Any] | None = None + def drain(self, execution_arn: str) -> None: + """Wait until this execution's latest record is exported and exporters flush. + + Returns once the calling execution's own record has reached every exporter + and a flush covering it has completed; a flush triggered by another + execution never releases a waiter whose record is still pending. + + Every call waits for a flush that completed after the call started, so an + invocation end that emitted no record still flushes -- the cadence JS and + Java have. Concurrent calls can share one flush, since they all started + before it completed. + + Two paths return without exporting or flushing anything, because the + permanent ``_disabled`` latch means no record will ever be exported: the + latch was already set when this call started, or it is set while this call + is parked. Failing to start the export worker sets that latch, so a drain + that hits a worker-start failure also returns without a flush. + """ + failed_pending: dict[str, tuple[int, dict[str, Any]]] | None = None start_error: Exception | None = None with self._condition: if self._disabled: return - if not self._flush_requested: - self._flush_requested = True - self._flush_event = threading.Event() - flush_event = self._flush_event - assert flush_event is not None - failed_pending, start_error = self._ensure_worker_locked() - started = not self._disabled - self._condition.notify() + lane = self._lane_locked(execution_arn) + lane.waiters += 1 + try: + want_seq = lane.scheduled_seq + # A drain always flushes, so require a flush that covers every + # export completed before this call as well as our own. + want_flush = self._export_count + # ...and one that completed after this call, so an invocation end + # that emitted nothing still flushes exactly once instead of + # riding on a flush that finished before it started. + want_flushes = self._flushes_completed + while not self._disabled: + # Export counter value a flush has to cover to release us: + # our own record's export plus everything already exported + # when this call started. Recomputed every pass, because + # lane.exported_at only becomes ours once our record is out. + need = max(lane.exported_at, want_flush) + if ( + lane.exported_seq >= want_seq + and self._flushed_through >= need + and self._flushes_completed > want_flushes + ): + break + # A flush already in flight whose coverage reaches `need` was + # committed after our record was handed to the exporters, so + # its completion releases us. Requesting another one here -- + # which is what a waiter woken inside that flush would do, + # since the coverage is not published yet and the request it + # made has already been consumed -- runs an extra flush after + # this drain, and the invocation, returned. + # + # `_flush_in_flight` uses 0 as its "no flush is running" + # sentinel, so the naive `self._flush_in_flight >= need` + # reads as "already covered" when `need` is 0 -- precisely + # when nothing is running at all. `need` is 0 for a drain + # whose invocation emitted no record, so that form would let + # such a drain skip its request and park until some other + # execution happened to flush. Require a marker that is + # actually set AND that reaches `need`. + covered = 0 < self._flush_in_flight >= need + if not self._flush_requested and not covered: + self._flush_requested = True + self._flush_barrier = max(self._flush_barrier, self._seq) + failed_pending, start_error = self._ensure_worker_locked() + if start_error is not None: + break + self._condition.notify_all() + self._condition.wait() + finally: + lane.waiters -= 1 + self._forget_lane_locked(execution_arn, lane) del failed_pending if start_error is not None: _logger.warning( @@ -71,12 +191,33 @@ def drain(self) -> None: "asynchronous export: %s", start_error, ) - if started: - flush_event.wait() + + # -- internals ------------------------------------------------------------ + + def _lane_locked(self, execution_arn: str) -> _Lane: + lane = self._lanes.get(execution_arn) + if lane is None: + lane = _Lane() + self._lanes[execution_arn] = lane + return lane + + def _forget_lane_locked(self, execution_arn: str, lane: _Lane) -> None: + # Keep the lane while anything still depends on it; bookkeeping for a + # fully exported execution with no waiters is safe to drop, because a + # later drain then only needs a flush covering the exports so far. + if self._lanes.get(execution_arn) is not lane: + return + if lane.waiters: + return + if execution_arn in self._pending: + return + if lane.exported_seq < lane.scheduled_seq: + return + del self._lanes[execution_arn] def _ensure_worker_locked( self, - ) -> tuple[dict[str, Any] | None, Exception | None]: + ) -> tuple[dict[str, tuple[int, dict[str, Any]]] | None, Exception | None]: if self._worker is not None and self._worker.is_alive(): return None, None worker = threading.Thread( @@ -91,39 +232,111 @@ def _ensure_worker_locked( self._disabled = True self._worker = None failed_pending = self._pending - self._pending = None - failed_event = self._flush_event - self._flush_event = None + self._pending = {} + # Lanes hold plain counters, never customer objects, so they can be + # dropped under the lock. Nothing is retained once the plugin has + # given up on asynchronous export for good. + self._lanes = {} self._flush_requested = False - if failed_event is not None: - failed_event.set() + self._flush_barrier = 0 + self._flush_in_flight = 0 + # Release every waiter; the permanent disable latch means no record + # will ever be exported. + self._condition.notify_all() return failed_pending, exc return None, None + def _blocking_pending_locked(self) -> bool: + """True while a record scheduled at or before the flush barrier is pending.""" + barrier = self._flush_barrier + return any(seq <= barrier for seq, _ in self._pending.values()) + def _run(self) -> None: + # The worker slot must be empty whenever no worker is running, or + # _ensure_worker_locked() never starts a replacement and every later + # record sits pending forever. The loop's own exits clear it, but a + # BaseException from a customer exporter -- asyncio.CancelledError is one, + # so an exporter that merely touches asyncio can raise it without writing + # `raise` -- unwinds past them, and a thread that is unwinding still + # reports is_alive(), so the slot would stay occupied by a dead thread. + # Vacate it here, on every exit path, and wake anyone parked so they can + # ask for the replacement. + try: + self._run_loop() + finally: + with self._condition: + if self._worker is threading.current_thread(): + self._worker = None + self._condition.notify_all() + + def _run_loop(self) -> None: while True: + arn: str | None = None + seq = 0 record: dict[str, Any] | None = None - flush_event: threading.Event | None = None + flush_covers = 0 with self._condition: - while self._pending is None and not self._flush_requested: + while True: + if self._flush_requested and not self._blocking_pending_locked(): + self._flush_requested = False + self._flush_barrier = 0 + flush_covers = self._export_count + # Publish what this flush will cover before releasing the + # lock, so a waiter that wakes while it runs can see that + # this flush releases it and skip asking for another. + self._flush_in_flight = flush_covers + break + if self._pending: + arn, (seq, record) = next(iter(self._pending.items())) + del self._pending[arn] + break self._condition.wait() - if self._pending is not None: - record = self._pending - self._pending = None - else: - flush_event = self._flush_event - self._flush_event = None - self._flush_requested = False if record is not None: - self._export(record) + # Popping the record consumed this execution's pending slot, so + # nothing will ever export that snapshot again. The lane must + # therefore advance whatever export() did: skip it and + # lane.exported_seq never reaches a waiter's want_seq, so a drain + # parked on this execution is never released. _export() already + # contains every Exception, but a BaseException from a customer + # exporter unwinds through here. Count the attempt in a finally + # and let the exception continue out to the wrapper -- and into + # the thread's traceback -- with nothing swallowed. + try: + self._export(record) + finally: + # Release the exported record before re-locking: a custom + # finalizer may re-enter schedule(). + del record + assert arn is not None + with self._condition: + self._export_count += 1 + lane = self._lanes.get(arn) + if lane is not None: + if seq > lane.exported_seq: + lane.exported_seq = seq + lane.exported_at = self._export_count + self._forget_lane_locked(arn, lane) + self._condition.notify_all() continue - self._flush() - if flush_event is not None: - flush_event.set() + flushed = False + try: + self._flush() + flushed = True + finally: + with self._condition: + # Retire the marker whatever happened: a stale one would park + # every later waiter that trusted this flush to cover it. Only + # a flush that ran to completion publishes its coverage. + self._flush_in_flight = 0 + if flushed: + self._flushes_completed += 1 + if flush_covers > self._flushed_through: + self._flushed_through = flush_covers + self._condition.notify_all() with self._condition: - if self._pending is None and not self._flush_requested: + if not self._pending and not self._flush_requested: self._worker = None return @@ -159,4 +372,4 @@ def _worker_alive(self) -> bool: def _pending_count(self) -> int: with self._condition: - return int(self._pending is not None) + return len(self._pending) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index f19796a0..0da43e34 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -160,7 +160,14 @@ def _apply_result_override( class _ExecutionState: - __slots__ = ("start_time", "parsed_arn", "cached_input", "operations") + __slots__ = ( + "start_time", + "parsed_arn", + "cached_input", + "operations", + "closed", + "lock", + ) def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: self.start_time = start_time @@ -169,6 +176,21 @@ def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: # operation_id -> OperationInfo, adopted verbatim from the SDK's # authoritative snapshot (invocation start/end and operation-change). self.operations: dict[str, OperationInfo] = {} + # Set once the invocation this state belongs to 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). + self.closed = False + # Guards `closed`, the operations rebind and record emission for this + # execution, so a late hook can never slip a RUNNING record in after the + # terminal one. Per execution, so concurrent executions never contend. + # Reentrant on purpose: `_emit` runs the scheduler's `schedule()` inside + # this hold, and `schedule()` releases the record it displaces, which can + # run a customer finalizer that re-enters a hook for this same execution + # on this same thread. A plain lock self-deadlocks the invocation thread + # there. (Java holds no such lock: its ExecutionState carries no + # operations map, and `cachedInput` is a bare volatile field.) + self.lock = threading.RLock() class WorkflowInsightPlugin(DurableInstrumentationPlugin): @@ -226,18 +248,25 @@ def _ensure_state(self, execution_arn: str) -> _ExecutionState: self._state[execution_arn] = state return state + def _get_state(self, execution_arn: str) -> _ExecutionState | None: + # Lookup only. A hook that must never fabricate state (an + # operation-change arriving after the invocation ended, whose state has + # been discarded) uses this instead of _ensure_state. + with self._lock: + return self._state.get(execution_arn) + def _discard_state(self, execution_arn: str) -> None: with self._lock: self._state.pop(execution_arn, None) - def _adopt_operations( + def _adopt_operations_locked( self, state: _ExecutionState, operations: dict[str, OperationInfo] ) -> None: # Adopt the authoritative point-in-time snapshot. Copy so plugin state # never aliases the SDK-owned map, and rebind the attribute so a # concurrent reader holding the prior reference iterates a stable dict. - with self._lock: - state.operations = dict(operations) + # Callers hold state.lock. + state.operations = dict(operations) # -- hooks ---------------------------------------------------------------- @@ -246,44 +275,57 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: if not arn or not self._sampled_in(arn): return state = self._ensure_state(arn) - # Always adopt the service-provided execution start time when present, - # including a cold resume in a fresh environment (never the resume time, - # which would corrupt duration and the date partition). - if info.execution_start_time is not None: - state.start_time = info.execution_start_time - state.cached_input = info.execution_input - # Seed the operation map from the full snapshot on every invocation. On a - # cold resume this rebuilds prior (terminal) operations that a fresh - # plugin instance never saw via per-operation hooks. - self._adopt_operations(state, info.operations) - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + with state.lock: + if state.closed: + return + # Always adopt the service-provided execution start time when + # present, including a cold resume in a fresh environment (never the + # resume time, which would corrupt duration and the date partition). + if info.execution_start_time is not None: + state.start_time = info.execution_start_time + state.cached_input = info.execution_input + # Seed the operation map from the full snapshot on every invocation. + # On a cold resume this rebuilds prior (terminal) operations that a + # fresh plugin instance never saw via per-operation hooks. + self._adopt_operations_locked(state, info.operations) + if self._emit_mode == EmitMode.ON_CHANGE: + self._emit( + arn, + state, + status="RUNNING", + end_time=None, + output_raw=None, + error=None, + ) def on_operation_change(self, info: OperationChangeInfo) -> None: arn = info.execution_arn if not arn or not self._sampled_in(arn): return - state = self._ensure_state(arn) - # Replace state with the full operations snapshot carried by the hook. - self._adopt_operations(state, info.operations) - # on-change mode exports an updated RUNNING record on each change so - # mid-invocation progress is observable, not only at start/end. - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + # Never create state here. A change hook for a checkpoint that completed + # just before the invocation ended still arrives after on_invocation_end + # discarded the state; recreating it would fabricate start_time = now, + # emit a RUNNING record after the terminal one, and leave a state entry + # behind for an execution this environment no longer runs. + state = self._get_state(arn) + if state is None: + return + with state.lock: + if state.closed: + return + # Replace state with the full operations snapshot carried by the hook. + self._adopt_operations_locked(state, info.operations) + # on-change mode exports an updated RUNNING record on each change so + # mid-invocation progress is observable, not only at start/end. + if self._emit_mode == EmitMode.ON_CHANGE: + self._emit( + arn, + state, + status="RUNNING", + end_time=None, + output_raw=None, + error=None, + ) def on_invocation_end(self, info: InvocationEndInfo) -> None: arn = info.execution_arn @@ -294,41 +336,66 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._discard_state(arn) return state = self._ensure_state(arn) - # Refresh from the fresh end-of-invocation snapshot before emitting so - # the terminal record reflects the final operation map. - self._adopt_operations(state, info.operations) - status = _STATUS_MAP.get(info.status, "RUNNING") - is_terminal = status in ("SUCCEEDED", "FAILED") - is_failure = status == "FAILED" - - if self._emit_mode == EmitMode.ON_CHANGE: - should_emit = True - elif self._emit_mode == EmitMode.ON_FAILURE: - should_emit = is_failure - else: # on-complete - should_emit = is_terminal - - if should_emit: - # Only terminal (SUCCEEDED/FAILED) records carry an end time; a - # PENDING/RETRY invocation end maps to RUNNING (still in flight) and - # must omit endTime/durationMs. Passing end_time=None makes _emit - # drop both fields. Output and error likewise belong only to a - # terminal record. - self._emit( - arn, - state, - status=status, - end_time=datetime.datetime.now(datetime.UTC) if is_terminal else None, - output_raw=info.execution_result if is_terminal else None, - error=info.error if is_terminal else None, - ) - self._scheduler.drain() + with state.lock: + if not state.closed: + # Close the gate before emitting so a concurrent late hook for + # this execution cannot append a RUNNING record after the + # terminal one. + state.closed = True + # Refresh from the fresh end-of-invocation snapshot before + # emitting so the terminal record reflects the final operation + # map. + self._adopt_operations_locked(state, info.operations) + status = _STATUS_MAP.get(info.status, "RUNNING") + is_terminal = status in ("SUCCEEDED", "FAILED") + is_failure = status == "FAILED" + + if self._emit_mode == EmitMode.ON_CHANGE: + should_emit = True + elif self._emit_mode == EmitMode.ON_FAILURE: + should_emit = is_failure + else: # on-complete + should_emit = is_terminal + + if should_emit: + # Only terminal (SUCCEEDED/FAILED) records carry an end time; + # a PENDING/RETRY invocation end maps to RUNNING (still in + # flight) and must omit endTime/durationMs. Passing + # end_time=None makes _emit drop both fields. Output and + # error likewise belong only to a terminal record. + self._emit( + arn, + state, + status=status, + end_time=datetime.datetime.now(datetime.UTC) + if is_terminal + else None, + output_raw=info.execution_result if is_terminal else None, + error=info.error if is_terminal else None, + # This is the emit that closed the gate, so it always runs + # with `closed` already set and must never drop itself. + closing=True, + ) # Clear state after EVERY invocation end, including PENDING/RETRY. The # next invocation rebuilds it from InvocationStartInfo.operations, so a # suspended execution that never resumes in this environment (or that was - # sampled out) leaks nothing and state stays bounded. + # sampled out) leaks nothing and state stays bounded. Done before the + # drain so a late hook for this execution finds no state to update while + # the terminal record is still in flight. self._discard_state(arn) + # Drain on EVERY sampled-in invocation end, emitted record or not: JS and + # Java flush once per sampled-in invocation end regardless, and a + # buffering exporter has to see the same rhythm in all three languages + # (an on-failure/on-complete mode that emits nothing for this invocation + # may still be holding records another execution handed it). A sampled-out + # execution returns above, so it neither exports nor flushes. + # + # The drain covers this execution only: it returns once this execution's + # own record, if any, reached the exporters and a flush that completed + # after this call is done, without waiting on records scheduled after the + # call by other executions. + self._scheduler.drain(arn) # -- emission ------------------------------------------------------------- @@ -383,6 +450,7 @@ def _emit( end_time: Any, output_raw: str | None, error: Any, + closing: bool = False, ) -> None: arn = state.parsed_arn start_time = state.start_time @@ -436,7 +504,31 @@ def _emit( record["error"] = {"name": error.type, "message": error.message} record["operations"] = self._build_operations(operations) - self._scheduler.schedule(record) + # INVARIANT: no record for an execution reaches the scheduler after that + # execution's closing record -- the exporters never see a RUNNING record + # follow the terminal one for the same execution. + # + # Re-check the gate here, because each hook's own `if state.closed` is a + # check-then-act and this is the act. Everything between the two runs + # customer code while holding state.lock: the input/output transforms + # above, a result override in _build_operations, and __del__ on any object + # the record carries. state.lock is a reentrant RLock on purpose (so that + # customer code re-entering a hook on this thread does not self-deadlock), + # which means such a re-entrant call can run on_invocation_end all the way + # through -- set `closed`, emit the terminal record, discard the state and + # drain it -- and then return here. Without this re-check the outer frame + # hands its already-built RUNNING record to the scheduler afterwards, and + # one execution exports ['SUCCEEDED', 'RUNNING'] from a single hook call, + # no concurrency required. + # + # `closing` marks the emit that set the gate (on_invocation_end's own), + # which by construction always runs with `closed` set and must not drop + # itself. It is not the same test as "the record is terminal": in + # on-change mode a PENDING/RETRY invocation end legitimately emits a + # RUNNING record, and that record is the closing one. + if state.closed and not closing: + return + self._scheduler.schedule(execution_arn, record) def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPlugin: diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py index 82426609..4e58ac7b 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py @@ -61,7 +61,40 @@ def render(self, record: dict[str, Any]) -> Any: ... # pragma: no cover def export(self, record: dict[str, Any]) -> None: ... # pragma: no cover - def flush(self) -> None: ... # pragma: no cover + def flush(self) -> None: # pragma: no cover + """Push any records this exporter is buffering to their destination. + + Only an exporter that buffers needs a body here; one that writes + synchronously inside ``export()`` can leave it empty. The method itself is + *not* optional: it is part of this protocol, so an exporter without it + fails the static protocol check, and at run time the plugin's call raises + ``AttributeError``, which is caught and logged as an exporter failure on + every flush. + + When the plugin calls it: + + * Once per sampled-in invocation end, after that invocation's own record + -- if the emit mode produced one -- has been handed to every exporter. + An execution that is sampled out neither exports nor flushes. + Invocation ends that overlap in one environment may share a single + flush, so the call count is at most one per sampled-in invocation end. + * Never concurrently with ``export()`` on the same plugin instance: one + worker thread runs both, one call at a time. + + A flush may cover records belonging to other executions running in the + same environment, so it is not a per-execution barrier. + + It must return promptly. The invocation that triggered it cannot return + until it does, so a slow flush is billed to the customer's invocation. + + Failures are isolated: an ``Exception`` is logged, never retried, never + propagated into the execution, and never prevents another exporter from + flushing. A ``BaseException`` (``asyncio.CancelledError`` is one) is not + contained -- it skips the remaining exporters for that flush and ends the + export worker -- but it still never reaches the execution, and the + waiting invocation is released by a replacement worker running the flush + it asked for. + """ @dataclass(frozen=True) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 34879059..cfad4097 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -13,6 +13,11 @@ ) +ARN = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-{}/inv-1" +ARN_A = ARN.format("a") +ARN_B = ARN.format("b") + + def _record(value: str) -> dict[str, Any]: return {"status": "RUNNING", "value": value, "operations": []} @@ -62,20 +67,38 @@ def flush(self) -> None: raise RuntimeError("flush failed") -def test_latest_pending_coalesces_without_blocking_schedule() -> None: +class ExporterBaseException(BaseException): + """Stands in for the BaseExceptions a customer exporter can raise.""" + + +class BaseExceptionExporter(CaptureExporter): + """Raises a ``BaseException`` out of its first ``export()`` call.""" + + def __init__(self) -> None: + super().__init__() + self.exports = 0 + + def export(self, record: dict[str, Any]) -> None: + self.exports += 1 + if self.exports == 1: + raise ExporterBaseException("export exploded") + super().export(record) + + +def test_latest_pending_coalesces_within_one_execution() -> None: exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) - scheduler.schedule(_record("first")) + scheduler.schedule(ARN_A, _record("first")) assert exporter.started.wait(5.0) start = time.monotonic() - scheduler.schedule(_record("middle")) - scheduler.schedule(_record("latest")) + scheduler.schedule(ARN_A, _record("middle")) + scheduler.schedule(ARN_A, _record("latest")) assert time.monotonic() - start < 0.5 assert scheduler._pending_count() == 1 exporter.release.set() - scheduler.drain() + scheduler.drain(ARN_A) assert exporter.calls == [ ("export", "first"), ("export", "latest"), @@ -89,19 +112,48 @@ def test_exporter_failure_does_not_block_other_exporters() -> None: capture = CaptureExporter() scheduler = _ExportScheduler([failing, capture]) - scheduler.schedule(_record("terminal")) + scheduler.schedule(ARN_A, _record("terminal")) - scheduler.drain() + scheduler.drain(ARN_A) assert capture.calls == [("export", "terminal"), ("flush", None)] +def test_base_exception_from_export_still_releases_drain() -> None: + # _export() contains every Exception, but a BaseException from a customer + # exporter -- asyncio.CancelledError is one -- unwinds out of the worker + # instead. The record has already been taken out of _pending by then and + # nothing will re-export that snapshot, so the export has to count and the + # worker slot has to be vacated anyway; otherwise the drain, and with it the + # invocation thread, parks forever. Every wait here is bounded so a + # regression fails instead of hanging the suite. + exporter = BaseExceptionExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + + assert returned.wait(10.0), "drain() never returned after export() raised" + thread.join(5.0) + assert not thread.is_alive() + # The snapshot was consumed, not retried, and the drain still got its flush. + assert exporter.exports == 1 + assert exporter.calls == [("flush", None)] + assert _wait_until(lambda: not scheduler._worker_alive()) + + def test_drain_flushes_after_export() -> None: capture = CaptureExporter() scheduler = _ExportScheduler([capture]) - scheduler.schedule(_record("terminal")) + scheduler.schedule(ARN_A, _record("terminal")) - scheduler.drain() + scheduler.drain(ARN_A) assert capture.calls == [("export", "terminal"), ("flush", None)] @@ -112,8 +164,8 @@ def fail_start(self) -> None: # noqa: ARG001 monkeypatch.setattr(threading.Thread, "start", fail_start) scheduler = _ExportScheduler([CaptureExporter()]) - scheduler.schedule(_record("dropped")) - scheduler.drain() + scheduler.schedule(ARN_A, _record("dropped")) + scheduler.drain(ARN_A) assert scheduler._pending_count() == 0 @@ -122,33 +174,33 @@ def test_superseded_record_finalizes_after_lane_unlock() -> None: scheduler = _ExportScheduler([BlockingExporter()]) exporter = scheduler._exporters[0] assert isinstance(exporter, BlockingExporter) - scheduler.schedule(_record("inflight")) + scheduler.schedule(ARN_A, _record("inflight")) assert exporter.started.wait(5.0) finalized = threading.Event() class ReentrantValue: def __del__(self) -> None: - scheduler.schedule(_record("from-finalizer")) + scheduler.schedule(ARN_A, _record("from-finalizer")) finalized.set() pending = _record("superseded") pending["payload"] = ReentrantValue() - scheduler.schedule(pending) + scheduler.schedule(ARN_A, pending) del pending - scheduler.schedule(_record("replacement")) + scheduler.schedule(ARN_A, _record("replacement")) assert finalized.wait(5.0) exporter.release.set() - scheduler.drain() + scheduler.drain(ARN_A) def test_drain_waits_for_blocked_exporter() -> None: exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) - scheduler.schedule(_record("terminal")) + scheduler.schedule(ARN_A, _record("terminal")) assert exporter.started.wait(5.0) - drain_thread = threading.Thread(target=scheduler.drain) + drain_thread = threading.Thread(target=scheduler.drain, args=(ARN_A,)) drain_thread.start() assert _wait_until(drain_thread.is_alive) @@ -157,3 +209,350 @@ def test_drain_waits_for_blocked_exporter() -> None: assert not drain_thread.is_alive() assert exporter.calls == [("export", "terminal"), ("flush", None)] + + +# -- concurrent executions in one environment (LMI) --------------------------- + + +class EventLogExporter: + """Records ``(kind, executionArn, status)`` events from every thread.""" + + max_record_size_bytes: int | None = None + + def __init__(self) -> None: + self.events: list[tuple[str, str | None, str | None]] = [] + self.lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self.log("export", record["executionArn"], record["status"]) + + def flush(self) -> None: + self.log("flush", None, None) + + def log(self, kind: str, arn: str | None, status: str | None) -> None: + with self.lock: + self.events.append((kind, arn, status)) + + def snapshot(self) -> list[tuple[str, str | None, str | None]]: + with self.lock: + return list(self.events) + + +class GatedEventLogExporter(EventLogExporter): + """Blocks inside the first ``export()`` until released.""" + + def __init__(self) -> None: + super().__init__() + self.first_export_started = threading.Event() + self.release = threading.Event() + + def export(self, record: dict[str, Any]) -> None: + if not self.first_export_started.is_set(): + self.first_export_started.set() + self.release.wait(10.0) + super().export(record) + + +class FlushGateEventLogExporter(EventLogExporter): + """Logs flush begin/end and blocks inside the first ``flush()`` until released.""" + + def __init__(self) -> None: + super().__init__() + self.first_flush_started = threading.Event() + self.release_flush = threading.Event() + self.flush_count = 0 + + def flush(self) -> None: + with self.lock: + self.flush_count += 1 + first = self.flush_count == 1 + self.log("flush_begin", None, None) + if first: + self.first_flush_started.set() + self.release_flush.wait(10.0) + self.log("flush_end", None, None) + + def flushes(self) -> int: + with self.lock: + return self.flush_count + + +def _execution_record(arn: str, status: str) -> dict[str, Any]: + return {"executionArn": arn, "status": status, "operations": []} + + +def _scheduler_is_empty(scheduler: _ExportScheduler) -> bool: + with scheduler._condition: + return not scheduler._pending and not scheduler._lanes + + +def _drain_waiters(scheduler: _ExportScheduler, arn: str) -> int: + """How many drain() calls are currently parked on this execution's lane.""" + with scheduler._condition: + lane = scheduler._lanes.get(arn) + return 0 if lane is None else lane.waiters + + +def test_drain_stays_parked_until_a_flush_covering_its_record_completes() -> None: + # The guarantee the per-execution scheduler exists for: drain() returns only + # after the calling execution's own record reached the exporters AND a flush + # that started after that export has itself completed. Gating the exporter + # inside flush() makes that deterministic -- while the gate is held the flush + # provably cannot have completed, so a drain that returns is a violation. + exporter = FlushGateEventLogExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + exporter.log("drain", ARN_A, None) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert exporter.first_flush_started.wait(5.0) + # Its own record went to the exporters before this flush was even started. + assert ("export", ARN_A, "SUCCEEDED") in exporter.snapshot() + assert not returned.wait(0.25), "drain returned while its flush was still running" + assert thread.is_alive() + + exporter.release_flush.set() + thread.join(5.0) + assert not thread.is_alive() + + events = exporter.snapshot() + exported = events.index(("export", ARN_A, "SUCCEEDED")) + flush_begin = events.index(("flush_begin", None, None)) + flush_end = events.index(("flush_end", None, None)) + drained = events.index(("drain", ARN_A, None)) + assert exported < flush_begin < flush_end < drained + + +def test_no_redundant_flush_runs_after_drain_returned() -> None: + # A drain woken while a flush is in flight must recognise that the in-flight + # flush already covers it. Otherwise it re-requests one (its own request has + # been consumed and the coverage is not published yet) and that second flush + # calls the exporters after drain(), and with it the invocation, returned. + exporter = FlushGateEventLogExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + returned = threading.Event() + flushes_at_return: list[int] = [] + + def drain() -> None: + scheduler.drain(ARN_A) + flushes_at_return.append(exporter.flushes()) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + # The worker is inside the flush it committed to: it has consumed the drain's + # request and has not published the flush's coverage yet. + assert exporter.first_flush_started.wait(5.0) + # Wake the parked drain exactly inside that window, which is the interleaving + # a loaded environment produces by itself. + for _ in range(3): + with scheduler._condition: + scheduler._condition.notify_all() + time.sleep(0.01) + assert not returned.is_set() + + exporter.release_flush.set() + thread.join(5.0) + assert not thread.is_alive() + # The worker only retires once nothing is pending and no flush is requested, + # so this settles the question without sleeping for a late flush. + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert flushes_at_return == [1] + assert exporter.flushes() == 1, "a second flush ran after drain() returned" + + +def test_drain_with_nothing_to_export_still_flushes_exactly_once() -> None: + # A drain for an execution that scheduled no record -- the invocation end + # emitted nothing -- has `need` == 0: no export has to be covered to release + # it. `_flush_in_flight` uses 0 for "no flush is running", so a naive + # `_flush_in_flight >= need` reads as "a flush already covers me" exactly + # when nothing is running at all; the drain would skip its request and park + # until some other execution happened to flush. It has to flush once and + # return. + capture = CaptureExporter() + scheduler = _ExportScheduler([capture]) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert returned.wait(10.0), "drain() parked with nothing of its own to export" + thread.join(5.0) + assert not thread.is_alive() + assert capture.calls == [("flush", None)] + # Nothing arrives after it returned either. + assert _wait_until(lambda: not scheduler._worker_alive()) + assert capture.calls == [("flush", None)] + + +def test_drain_never_rides_on_a_flush_that_finished_before_it_started() -> None: + # Export coverage alone would let the second drain return immediately: every + # export is already covered by the first drain's flush. A drain must wait for + # a flush that completed after it was called, so that one invocation end + # means one flush. + capture = CaptureExporter() + scheduler = _ExportScheduler([capture]) + scheduler.schedule(ARN_A, _record("terminal")) + + scheduler.drain(ARN_A) + scheduler.drain(ARN_A) + + assert capture.calls == [ + ("export", "terminal"), + ("flush", None), + ("flush", None), + ] + + +def test_disabled_latch_retains_no_lanes_or_pending_records(monkeypatch) -> None: + # The _disabled latch is permanent: nothing will ever be exported again, so + # the scheduler must not hold on to records or per-execution bookkeeping for + # the remaining life of the environment. + def fail_start(self) -> None: # noqa: ARG001 + raise RuntimeError("cannot start") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler = _ExportScheduler([CaptureExporter()]) + + scheduler.schedule(ARN_A, _record("dropped")) + scheduler.drain(ARN_A) + scheduler.schedule(ARN_B, _record("also-dropped")) + scheduler.drain(ARN_B) + + with scheduler._condition: + assert scheduler._disabled + assert scheduler._pending == {} + assert scheduler._lanes == {} + assert scheduler._flush_requested is False + assert scheduler._flush_in_flight == 0 + + +def test_disabled_latch_clears_a_published_flush_in_flight_marker(monkeypatch) -> None: + # `_flush_in_flight` is the coverage of the flush the worker is running right + # now, published so a waiter woken during that flush can tell it is already + # covered and skip requesting another. 0 means "no flush is running", so the + # marker is a claim that a flush is in flight and will complete. + # + # The _disabled latch makes that claim permanently false: no worker exists and + # none will ever be started again, so the published flush can never complete. + # The latch therefore has to retire the marker along with the pending records + # and the lanes. test_disabled_latch_retains_no_lanes_or_pending_records + # asserts the same field, but reaches the latch with the marker already at 0, + # so it holds whether or not the latch clears it; this one arms the marker + # first. + scheduler = _ExportScheduler([CaptureExporter()]) + with scheduler._condition: + scheduler._flush_in_flight = 7 + + def fail_start(self) -> None: # noqa: ARG001 + raise RuntimeError("cannot start") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler.schedule(ARN_A, _record("dropped")) + + with scheduler._condition: + assert scheduler._disabled + assert scheduler._flush_in_flight == 0, ( + "the _disabled latch left a flush-in-flight marker behind for a flush " + "that can never run" + ) + + +def test_every_concurrent_execution_delivers_its_terminal_record_once() -> None: + # One plugin instance (one scheduler) serves every execution the environment + # hosts. Ten executions running at once must each land their terminal record + # exactly once: a pending record keyed per execution is never displaced by a + # different execution's record. + executions = 10 + exporter = EventLogExporter() + scheduler = _ExportScheduler([exporter]) + arns = [ARN.format(index) for index in range(executions)] + ready = threading.Barrier(executions) + + def run(arn: str) -> None: + ready.wait(10.0) + scheduler.schedule(arn, _execution_record(arn, "RUNNING")) + scheduler.schedule(arn, _execution_record(arn, "SUCCEEDED")) + scheduler.drain(arn) + + threads = [threading.Thread(target=run, args=(arn,)) for arn in arns] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30.0) + assert not any(thread.is_alive() for thread in threads) + + terminal = [ + arn + for kind, arn, status in exporter.snapshot() + if kind == "export" and status == "SUCCEEDED" and arn is not None + ] + assert sorted(terminal) == sorted(arns) # each exactly once, none lost + # Nothing per-execution is retained once every execution has drained. + assert _wait_until(lambda: _scheduler_is_empty(scheduler)) + + +def test_blocked_export_never_loses_another_executions_terminal_record() -> None: + # The worker is busy exporting when two executions queue their terminal + # records. Neither may be dropped, and each drain must be released by its own + # record reaching the exporters -- not by another execution's flush. + exporter = GatedEventLogExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "RUNNING")) + assert exporter.first_export_started.wait(5.0) + + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + scheduler.schedule(ARN_B, _execution_record(ARN_B, "SUCCEEDED")) + returned: dict[str, bool] = {} + + def drain(arn: str) -> None: + scheduler.drain(arn) + exporter.log("drain", arn, None) + returned[arn] = True + + drains = [threading.Thread(target=drain, args=(arn,)) for arn in (ARN_A, ARN_B)] + for thread in drains: + thread.start() + # Both drains really parked: each registered on its own lane (which drain() + # only does from inside its wait loop) and neither has returned while the + # exporter still holds the worker inside the very first export. + assert _wait_until( + lambda: ( + _drain_waiters(scheduler, ARN_A) == 1 + and _drain_waiters(scheduler, ARN_B) == 1 + ) + ) + assert returned == {} + assert all(thread.is_alive() for thread in drains) + + exporter.release.set() + for thread in drains: + thread.join(10.0) + assert not any(thread.is_alive() for thread in drains) + assert returned == {ARN_A: True, ARN_B: True} + + events = exporter.snapshot() + assert ("export", ARN_A, "SUCCEEDED") in events + assert ("export", ARN_B, "SUCCEEDED") in events + for arn in (ARN_A, ARN_B): + exported = events.index(("export", arn, "SUCCEEDED")) + returned_at = events.index(("drain", arn, None)) + # Each drain returned only after its own record was exported and a flush + # covering it completed. + assert exported < returned_at + assert ("flush", None, None) in events[exported:returned_at] + assert _wait_until(lambda: _scheduler_is_empty(scheduler)) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 9d7638eb..35a9d98d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -13,6 +13,9 @@ from __future__ import annotations import datetime +import itertools +import threading +import time from typing import Any from aws_durable_execution_sdk_python.lambda_service import ( @@ -187,6 +190,46 @@ def test_on_failure_success_emits_nothing(): ) _run(plugin, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) assert exporter.records == [] + # No record, but the invocation end still flushed once: a sampled-in + # invocation end flushes whether or not this emit mode produced a record + # (JS/Java cadence), because the exporter may be buffering another + # execution's records. + assert exporter.flush_count == 1 + assert _wait_until(lambda: not plugin._scheduler._worker_alive()) + + +def test_invocation_end_that_emits_no_record_still_flushes_exactly_once(): + # on-complete mode with a PENDING end: the execution suspended, so nothing is + # emitted. JS and Java flush once per sampled-in invocation end regardless of + # whether a record was emitted, and a buffering exporter has to see the same + # rhythm in every SDK, so this end must still flush -- exactly once, not + # twice, and not zero times. + exporter = CaptureExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin.on_invocation_start(_start(operations={})) + plugin.on_invocation_end( + _end(operations={}, status=InvocationStatus.PENDING, result=None) + ) + assert exporter.records == [] + assert exporter.flush_count == 1 + # The worker retires, so no later flush can arrive after the invocation + # returned. + assert _wait_until(lambda: not plugin._scheduler._worker_alive()) + assert exporter.flush_count == 1 + + +def test_sampled_out_invocation_end_neither_exports_nor_flushes(): + # The sampled-out path is the one exception to the cadence above: a sampled + # out execution exports nothing and must not flush either, so instrumenting + # a fraction of executions costs the rest nothing. + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + plugin.on_invocation_end(_end(operations=_ops(op))) + assert exporter.records == [] assert exporter.flush_count == 0 assert not plugin._scheduler._worker_alive() @@ -587,3 +630,452 @@ def test_failed_end_is_terminal_with_end_time_and_duration(): assert rec["endTime"] is not None assert rec["durationMs"] is not None assert rec["error"]["name"] == "StepError" + + +# -- concurrent executions in one environment (LMI) --------------------------- + + +class ConcurrentCaptureExporter: + """CaptureExporter for multi-threaded drives; appends under a lock.""" + + max_record_size_bytes = None + + def __init__(self) -> None: + self.records: list[dict[str, Any]] = [] + self._lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + with self._lock: + self.records.append(record) + + def flush(self) -> None: + pass + + def snapshot(self) -> list[dict[str, Any]]: + with self._lock: + return list(self.records) + + +def test_concurrent_executions_each_deliver_their_terminal_record(): + # One plugin instance serves every execution its environment hosts, and LMI + # runs several at once. Drive the real hooks concurrently: every execution's + # terminal record must arrive exactly once. + executions = 5 + exporter = ConcurrentCaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + arns = [ + f"arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-{index}/inv-1" + for index in range(executions) + ] + ready = threading.Barrier(executions) + + def run(arn: str) -> None: + op = _step("s", op_id="1") + ready.wait(10.0) + plugin.on_invocation_start(_start(arn=arn, operations={})) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=arn, updated_operations=_ops(op), operations=_ops(op) + ) + ) + plugin.on_invocation_end(_end(arn=arn, operations=_ops(op))) + + threads = [threading.Thread(target=run, args=(arn,)) for arn in arns] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30.0) + assert not any(thread.is_alive() for thread in threads) + + terminal = [ + record["executionArn"] + for record in exporter.snapshot() + if record["status"] == "SUCCEEDED" + ] + assert sorted(terminal) == sorted(arns) # each exactly once, none lost + assert plugin._state == {} + # Nothing per-execution is retained in the scheduler either. + assert _wait_until(lambda: _scheduler_is_empty(plugin)) + + +def _wait_until(predicate, timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def _scheduler_is_empty(plugin) -> bool: + scheduler = plugin._scheduler + with scheduler._condition: + return not scheduler._pending and not scheduler._lanes + + +class PinnedWorkerExporter: + """Blocks inside every ``export()`` until released, pinning the one worker.""" + + max_record_size_bytes = None + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + self.entered.set() + self.release.wait(30.0) + + def flush(self) -> None: + pass + + +def test_reentrant_finalizer_in_a_hook_does_not_deadlock(): + # _emit runs the scheduler's schedule() while holding the execution's lock, + # and schedule() releases the record it displaces inside that hold, on + # purpose: a record can carry customer objects whose finalizers run arbitrary + # code. A finalizer that re-enters a hook for the same execution therefore + # re-acquires that lock on the thread that already owns it, which a + # non-reentrant lock turns into a permanent hang of the invocation thread. + exporter = PinnedWorkerExporter() + reentered = threading.Event() + holder: dict[str, Any] = {} + + class ReentrantPayload: + """A customer object that reaches the record through a content transform.""" + + def __del__(self) -> None: + if reentered.is_set(): + return + reentered.set() + holder["plugin"].on_operation_change( + OperationChangeInfo( + execution_arn=ARN, + updated_operations=holder["ops"], + operations=holder["ops"], + ) + ) + + plugin = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=lambda _value: ReentrantPayload()), + ) + ) + op = _step("s", op_id="1") + holder["plugin"] = plugin + holder["ops"] = _ops(op) + change = OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + try: + # The first emit pins the single worker inside export()... + plugin.on_invocation_start(_start(operations={})) + assert exporter.entered.wait(5.0) + # ...so this emit stays this execution's pending record... + plugin.on_operation_change(change) + + returned = threading.Event() + + def hook() -> None: + # ...and this one displaces it, releasing the displaced record (and + # running the payload's finalizer) on this thread, inside the lock. + plugin.on_operation_change(change) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), ( + "the hook never returned: a record finalizer that re-entered a hook " + "for the same execution deadlocked the invocation thread" + ) + assert reentered.is_set() # the finalizer really did re-enter a hook + finally: + exporter.release.set() + + +# -- late hooks after the invocation ended (closed gate) ---------------------- + + +def test_change_hook_reaching_the_lock_after_invocation_end_emits_nothing(): + # A change hook for a checkpoint that completed just before the invocation + # ended can reach the execution's lock while on_invocation_end still holds it. + # It must find the gate closed and emit nothing, so no RUNNING record can + # follow the terminal one. + exporter = ConcurrentCaptureExporter() + in_terminal_emit = threading.Event() + release_terminal = threading.Event() + + def blocking_output(value: Any) -> Any: + # Runs inside _emit, which runs inside the execution's lock, and only for + # a terminal record (a RUNNING record carries no output). + in_terminal_emit.set() + release_terminal.wait(10.0) + return value + + plugin = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(output=blocking_output), + ) + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + + end_returned = threading.Event() + + def end() -> None: + plugin.on_invocation_end(_end(operations=_ops(op))) + end_returned.set() + + end_thread = threading.Thread(target=end, daemon=True) + end_thread.start() + # on_invocation_end has closed the gate and is building the terminal record, + # still holding the execution's lock. + assert in_terminal_emit.wait(5.0) + + change_returned = threading.Event() + + def change() -> None: + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + change_returned.set() + + change_thread = threading.Thread(target=change, daemon=True) + change_thread.start() + # It cannot get past the execution's lock while the end hook holds it. + assert not change_returned.wait(0.25) + + release_terminal.set() + end_thread.join(5.0) + change_thread.join(5.0) + assert end_returned.is_set() + assert change_returned.is_set() + plugin._scheduler.drain(ARN) + + statuses = [record["status"] for record in exporter.snapshot()] + assert "SUCCEEDED" in statuses + # Nothing at all after the terminal record. + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + assert plugin._state == {} + + +def test_invocation_end_waits_for_an_in_flight_change_hook_emit(): + # The mirror interleaving: a change hook is already inside its emit, holding + # the execution's lock, when the invocation ends. Closing the gate and + # emitting the terminal record has to wait for it, otherwise the change hook + # finishes afterwards and appends a RUNNING record after the terminal one. + exporter = ConcurrentCaptureExporter() + in_change_emit = threading.Event() + release_change = threading.Event() + calls = itertools.count() + lock = threading.Lock() + + def blocking_input(value: Any) -> Any: + # Runs inside _emit for every record. Block only on the change hook's + # emit, which is the second one (invocation start emits the first). + with lock: + index = next(calls) + if index == 1: + in_change_emit.set() + release_change.wait(10.0) + return value + + plugin = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=blocking_input), + ) + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + + change_returned = threading.Event() + + def change() -> None: + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + change_returned.set() + + change_thread = threading.Thread(target=change, daemon=True) + change_thread.start() + assert in_change_emit.wait(5.0) + + end_returned = threading.Event() + + def end() -> None: + plugin.on_invocation_end(_end(operations=_ops(op))) + end_returned.set() + + end_thread = threading.Thread(target=end, daemon=True) + end_thread.start() + # The terminal record cannot be emitted while the change hook holds the lock. + assert not end_returned.wait(0.25) + + release_change.set() + change_thread.join(5.0) + end_thread.join(10.0) + assert change_returned.is_set() + assert end_returned.is_set() + plugin._scheduler.drain(ARN) + + statuses = [record["status"] for record in exporter.snapshot()] + assert "SUCCEEDED" in statuses + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + assert plugin._state == {} + + +def test_operation_change_after_invocation_end_emits_nothing(): + # A checkpoint that completed just before the invocation ended still delivers + # its operation-change hook. It must not recreate state, must not fabricate a + # start time, and must not append a RUNNING record after the terminal one. + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + plugin.on_invocation_end(_end(operations=_ops(op))) + before = list(exporter.records) + assert before and before[-1]["status"] == "SUCCEEDED" + + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + plugin._scheduler.drain(ARN) + + assert exporter.records == before # nothing emitted after the terminal record + assert [record["status"] for record in exporter.records][-1] == "SUCCEEDED" + # No fabricated start time: every record still reports the execution start. + assert {record["startTime"] for record in exporter.records} == { + "2026-01-01T00:00:00Z" + } + assert plugin._state == {} # and no state entry recreated + + +def test_late_invocation_start_finds_the_closed_gate_shut(monkeypatch): + # on_invocation_end sets `closed` and emits the terminal record while holding + # the execution's lock, RELEASES the lock, and only then discards the state. + # A concurrent on_invocation_start that already resolved that state reference + # gets the lock inside that window and finds a state that is closed but not + # yet gone; the gate has to shut it out. + # + # In production that window is sub-microsecond, so it is entered here + # deterministically: the late hook runs from inside _discard_state, which is + # exactly where the window sits. + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={}, input_value="World")) + state = plugin._state[ARN] + real_discard = plugin._discard_state + late = threading.Event() + + def discard_after_a_late_start(arn: str) -> None: + if not late.is_set(): + late.set() + assert state.closed # the window: closed, emitted, lock free, state alive + plugin.on_invocation_start( + _start( + operations=_ops(_step("late", op_id="2")), + input_value="late-input", + execution_start_time=T1, + ) + ) + real_discard(arn) + + monkeypatch.setattr(plugin, "_discard_state", discard_after_a_late_start) + plugin.on_invocation_end(_end(operations=_ops(op))) + plugin._scheduler.drain(ARN) + + assert late.is_set() # the late hook really did run inside the window + statuses = [record["status"] for record in exporter.records] + assert "SUCCEEDED" in statuses + # Nothing follows the terminal record... + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + # ...and the closed state was not re-seeded on the way out. A late start that + # got past the gate adopts its own operation snapshot, input and start time, + # which is the observable effect of the gate: the emission it would also have + # produced is stopped a second time by the re-check in _emit, so these are + # what pin the hook's own check. + assert state.cached_input == "World" + assert state.start_time == T0 + assert [info.name for info in state.operations.values()] == ["s"] + assert plugin._state == {} + + +def test_reentrant_invocation_end_stops_the_outer_running_record(): + # The gate at the top of each hook is a check-then-act, and _emit is the act. + # Between them _emit runs customer code while holding the execution's lock -- + # here a content transform -- and the lock is reentrant, so that customer code + # can run on_invocation_end to completion on this same thread: `closed` set, + # terminal record scheduled, state discarded and drained. The outer frame then + # resumes with a fully built RUNNING record, which must NOT reach the + # exporters after the terminal one. One hook call, no concurrency. + exporter = ConcurrentCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if not reentered.is_set(): + reentered.set() + holder["plugin"].on_invocation_end(_end(operations=_ops(_step("s")))) + return value + + plugin = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + holder["plugin"] = plugin + + # On a bounded thread, so a regression that makes the lock non-reentrant + # again fails here instead of hanging the suite. + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(_start(operations={})) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), ( + "the hook never returned: re-entering on_invocation_end from customer " + "code inside _emit deadlocked the invocation thread" + ) + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() # the re-entrant end hook really did run + # Force everything the plugin scheduled to reach the exporters, so a record + # that slipped past the gate is observed here rather than left pending. + plugin._scheduler.drain(ARN) + + statuses = [record["status"] for record in exporter.snapshot()] + assert statuses == ["SUCCEEDED"], ( + "a non-terminal record reached the exporters after the terminal one for " + f"the same execution: {statuses}" + ) + assert _wait_until(lambda: not plugin._scheduler._worker_alive()) From 49747479c67d379f1f0cae4f19146d5206f88b7c Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 17 Sep 2026 06:44:53 -0700 Subject: [PATCH 02/28] refactor(plugin): create one plugin instance per invocation A plugin instance used to live as long as the execution environment while serving every execution that landed on it. Under Lambda Managed Instances several executions run concurrently in one environment, so every plugin had to key its own state by execution ARN and clean that map up itself. Getting that wrong loses telemetry silently, and both bundled plugins had got it wrong. The SDK now builds one plugin per invocation and drops it when the invocation ends. `plugins` takes factories, not instances: @durable_execution(plugins=[lambda info: MyPlugin(shared_exporter)]) A factory is any callable taking `InvocationStartInfo` and returning a plugin. Environment-lifetime state stays in a callable class or a closure; per-invocation state becomes ordinary instance attributes. This deletes the provider path rather than adding a parallel one. While an instance path exists a plugin cannot delete its ARN-keyed dict, which is the entire point of the change. Removed: `DurableInstrumentationPluginProvider` and `DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION`. Added: `DurableInstrumentationPluginFactory`. Entry points in the `aws_durable_execution.plugins` group now name a callable, so environment-based loading works unchanged; `load_configured_plugins` dedups by factory identity instead of by plugin type. Workflow Insight now holds no ARN-keyed structure at all. Its `RLock` and its pre-hand-off re-check both stay: hooks dispatch on the producing thread, so races within one invocation remain, and `schedule()` releases displaced records while holding the lock, which can re-enter through a customer finalizer. The OTel plugin keeps each invocation's spans in plain attributes, so the state-reset machinery is gone. `install_log_filter` now rebinds the installed filter to the current invocation's plugin. Without that, a logging handler outlives the plugin that installed it, keeps asking a discarded instance for a span context, and log correlation stops after the first invocation while the dead instance stays reachable for the life of the environment. BREAKING CHANGE: `plugins` accepts factories instead of plugin instances, and the provider class is removed. Where you passed `MyPlugin()`, pass a factory: `lambda info: MyPlugin()`. A plugin class is itself a valid factory when its constructor takes the info argument, because calling a class in Python constructs an instance. Entry points in the `aws_durable_execution.plugins` group must now name a callable rather than a provider object. --- .../README.md | 9 +- .../src/common.py | 21 +- .../src/otel_10_wait_for_callback.py | 4 +- .../src/otel_11_chained_invoke.py | 6 +- .../src/otel_12_child_context_failure.py | 4 +- .../src/otel_13_parallel_failure.py | 4 +- .../src/otel_14_map_failure.py | 4 +- .../src/otel_15_wait_interrupted.py | 4 +- .../src/otel_16_wait_for_condition_failure.py | 4 +- .../src/otel_17_wait_for_callback_failure.py | 4 +- .../src/otel_18_chained_invoke_failure.py | 6 +- .../src/otel_19_execution_failure.py | 4 +- .../src/otel_1_success.py | 4 +- .../src/otel_20_virtual_context.py | 4 +- .../src/otel_2_wait_resume.py | 4 +- .../src/otel_3_retry.py | 4 +- .../src/otel_4_terminal_failure.py | 4 +- .../src/otel_5_child_context.py | 4 +- .../src/otel_6_parallel.py | 4 +- .../src/otel_7_map.py | 4 +- .../src/otel_8_handled_failure.py | 4 +- .../src/otel_9_wait_for_condition.py | 4 +- .../src/otel_long_running_1_wait.py | 4 +- .../src/otel_long_running_2_retry.py | 4 +- .../src/otel_long_running_3_callback.py | 4 +- .../src/otel_long_running_4_chained_invoke.py | 6 +- .../tests/test_otel_examples.py | 4 +- .../plugin/plugin_attempt_hooks_retry.py | 2 +- .../plugin/plugin_attempt_info_shape.py | 2 +- .../plugin/plugin_context_info_shape.py | 2 +- .../handlers/plugin/plugin_error_isolation.py | 2 +- .../plugin_external_update_on_invoke.py | 2 +- .../plugin/plugin_faulty_and_healthy.py | 4 +- .../plugin/plugin_first_invocation_flag.py | 2 +- .../plugin/plugin_invocation_info_shape.py | 2 +- .../plugin/plugin_invocation_lifecycle.py | 2 +- .../plugin/plugin_multiple_plugins.py | 2 +- .../plugin/plugin_nested_parent_linkage.py | 2 +- .../plugin/plugin_operation_change.py | 2 +- .../plugin/plugin_operation_change_shape.py | 2 +- .../plugin/plugin_operation_info_shape.py | 2 +- .../plugin/plugin_operation_lifecycle.py | 2 +- .../plugin/plugin_parallel_branch_hooks.py | 2 +- .../handlers/plugin/plugin_replay_flags.py | 2 +- .../plugin/plugin_retry_exhaustion.py | 2 +- .../plugin_suspension_invocation_end.py | 2 +- .../plugin/plugin_terminal_failure.py | 2 +- .../plugin/plugin_terminal_payloads.py | 2 +- .../plugin/plugin_wait_operation_hooks.py | 2 +- .../plugin/plugin_wait_replay_flag.py | 2 +- .../src/otel/otel_logger_example.py | 11 +- .../src/plugin/execution_with_otel.py | 4 +- .../src/plugin/execution_with_plugin.py | 2 +- .../src/plugin/execution_with_wait_plugin.py | 2 +- .../README.md | 10 +- .../_export_scheduler.py | 186 ++++---- .../plugin.py | 422 ++++++++++-------- .../tests/e2e/wait_suspend_resume_int_test.py | 9 +- .../tests/test_export_scheduler.py | 79 +++- .../tests/test_plugin.py | 352 ++++++++------- .../README.md | 70 ++- .../pyproject.toml | 4 +- .../__init__.py | 10 + .../execution_plugin.py | 60 ++- .../invocation_plugin.py | 61 ++- .../log_filter.py | 27 +- .../plugin_factory.py | 90 ++++ .../plugin_provider.py | 23 - .../e2e/test_invocation_wait_resume_int.py | 22 +- .../tests/test_execution_plugin.py | 40 +- .../test_execution_plugin_integration.py | 9 +- .../tests/test_invocation_plugin.py | 33 +- .../test_invocation_plugin_integration.py | 9 +- .../tests/test_log_filter.py | 31 ++ .../tests/test_plugin_factory.py | 177 ++++++++ .../tests/test_plugin_provider.py | 56 --- .../tests/e2e/wait_suspend_replay_test.py | 11 +- .../README.md | 32 +- .../execution.py | 9 +- .../plugin.py | 92 +++- .../plugin_discovery.py | 131 ++---- .../tests/context_test.py | 18 +- .../plugin_invocation_operations_int_test.py | 52 ++- .../e2e/plugin_invocation_payload_int_test.py | 10 +- ...plugin_user_function_lifecycle_int_test.py | 3 +- .../tests/execution_test.py | 111 ++++- .../tests/plugin_discovery_test.py | 372 +++++++-------- .../tests/plugin_test.py | 335 ++++++++++++-- .../tests/state_test.py | 91 ++-- .../tests/test_helpers.py | 48 ++ 90 files changed, 2108 insertions(+), 1190 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py delete mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py delete mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 5ff978c8..b36a1b2c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -35,9 +35,9 @@ tests/ # contract tests for the templates and handlers The 20 invocation and 20 execution requirements reuse the same scenario handlers; the view is selected per function through the `OTEL_PLUGIN_MODE` -environment variable, which `common.otel_plugin()` reads to pick -`InvocationOtelPlugin` or `ExecutionOtelPlugin`. `template.yaml` deploys only the -view named by its `OtelSuite` parameter. +environment variable, which `common.otel_plugin_factory()` reads to pick +`InvocationOtelPluginFactory` or `ExecutionOtelPluginFactory`. `template.yaml` +deploys only the view named by its `OtelSuite` parameter. ## Scenarios @@ -183,7 +183,8 @@ write access; the runner identity needs list, read, and cleanup access. `test-requirements//.yaml`. New requirement IDs must be registered there first. 2. Add `src/otel__.py` exporting `handler`. Select the plugin with - `common.otel_plugin()` and guard the input with `common.require_scenario()`. + `common.otel_plugin_factory()` and guard the input with + `common.require_scenario()`. Use the SDK's real API; never hand-roll behavior to force an expected result. 3. Register the function in `template.yaml` (or `template-long-running.yaml`) with `Handler: .handler` and `TestDescription: [""]`, and add a diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py index 7c8b7c59..bc8cbc37 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py @@ -9,20 +9,27 @@ from collections.abc import Mapping from typing import Any -from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPluginFactory, +) from aws_durable_execution_sdk_python_otel import ( - ExecutionOtelPlugin, - InvocationOtelPlugin, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, OtelPluginConfig, ) -def otel_plugin() -> DurableInstrumentationPlugin: - """Select the telemetry view configured for this deployed function.""" +def otel_plugin_factory() -> DurableInstrumentationPluginFactory: + """Select the telemetry view configured for this deployed function. + + Returns a factory, which is what ``durable_execution(plugins=[...])`` takes: + the SDK calls it once per invocation to build that invocation's plugin. The + view is still resolved once, when the handler module is imported. + """ if os.environ.get("OTEL_PLUGIN_MODE") == "execution": - return ExecutionOtelPlugin(OtelPluginConfig()) - return InvocationOtelPlugin(OtelPluginConfig()) + return ExecutionOtelPluginFactory(OtelPluginConfig()) + return InvocationOtelPluginFactory(OtelPluginConfig()) def require_scenario(event: Mapping[str, Any], expected: str) -> None: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py index 3807d2ae..107362c9 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def submit_callback( @@ -19,7 +19,7 @@ def submit_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "wait-for-callback") return context.wait_for_callback( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py index 0a43b520..777d222c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py @@ -9,10 +9,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler( event: dict[str, Any], context: DurableContext, @@ -25,7 +25,7 @@ def handler( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( event: dict[str, Any], _context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py index 2d06b9a3..365c2923 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py @@ -12,7 +12,7 @@ durable_execution, durable_with_child_context, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_with_child_context @@ -20,7 +20,7 @@ def fail_child_context(_context: DurableContext) -> None: raise RuntimeError("Intentional child-context failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "child-context-failure") context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py index c36fadf7..416b8bec 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py @@ -13,7 +13,7 @@ durable_parallel_branch, ) from aws_durable_execution_sdk_python.config import ParallelConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_parallel_branch(name="otel-failed-parallel-branch") @@ -21,7 +21,7 @@ def fail_parallel_branch(_context: DurableContext) -> None: raise RuntimeError("Intentional parallel branch failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "parallel-failure") result = context.parallel( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py index b77b1ddd..6e1bd35c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py @@ -10,7 +10,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import MapConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def fail_map_item( @@ -22,7 +22,7 @@ def fail_map_item( raise RuntimeError("Intentional map iteration failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "map-failure") result = context.map( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py index f8c225dc..584259c0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py @@ -9,10 +9,10 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import Duration -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-interrupted") context.wait( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py index db60d3c8..23aaa72a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py @@ -14,7 +14,7 @@ WaitForConditionConfig, WaitForConditionDecision, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def fail_condition_check( @@ -31,7 +31,7 @@ def continue_condition( return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-for-condition-failure") context.wait_for_condition( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py index 8aa55881..5b0a1b9f 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def submit_failed_callback( @@ -19,7 +19,7 @@ def submit_failed_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-for-callback-failure") context.wait_for_callback( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py index c8acdd28..d7e4dd9b 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py @@ -9,10 +9,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "chained-invoke-failure") context.invoke( @@ -22,7 +22,7 @@ def handler(event: dict[str, Any], context: DurableContext) -> None: ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( _event: dict[str, Any], _context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py index 887a488c..cd14daf3 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py @@ -8,10 +8,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], _context: DurableContext) -> None: require_scenario(event, "execution-failure") raise RuntimeError("Intentional execution failure") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py index 0409588f..1fb1ef2e 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py @@ -13,7 +13,7 @@ durable_execution, durable_step, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -21,7 +21,7 @@ def complete_successfully(_step_context: StepContext) -> str: return "success" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "success") return context.step(complete_successfully(), name="otel-success") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py index 1f957bdb..2e867ae1 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py @@ -13,7 +13,7 @@ durable_with_child_context, ) from aws_durable_execution_sdk_python.config import ChildConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_with_child_context @@ -21,7 +21,7 @@ def run_virtual_context(_context: DurableContext) -> str: return "virtual-complete" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "virtual-context") return context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py index 21041a94..40e13596 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py @@ -14,7 +14,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import Duration -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -22,7 +22,7 @@ def complete_after_resume(_step_context: StepContext) -> str: return "resumed" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "wait-resume") context.wait(Duration.from_seconds(1), name="otel-wait") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py index 761f76b4..87f51ad9 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -28,7 +28,7 @@ def succeed_on_retry(step_context: StepContext) -> str: return "retried" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "retry") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py index 5850757e..cd4f81d0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -26,7 +26,7 @@ def fail_terminally(_step_context: StepContext) -> None: raise RuntimeError("Intentional terminal failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "terminal-failure") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py index b9d47534..61d049e0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py @@ -14,7 +14,7 @@ durable_step, durable_with_child_context, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -27,7 +27,7 @@ def run_child_workflow(context: DurableContext) -> str: return context.step(complete_child_step(), name="otel-child-step") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "child-context") return context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py index e2ac6adb..e7b54e39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py @@ -15,7 +15,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import ParallelConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -39,7 +39,7 @@ def run_parallel_branch_b(context: DurableContext) -> str: ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> list[str]: require_scenario(event, "parallel-hierarchy") return context.parallel( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py index b35ad6e7..e69c68ab 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py @@ -15,7 +15,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import MapConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -35,7 +35,7 @@ def process_map_item( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> list[int]: require_scenario(event, "map-hierarchy") return context.map( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py index 44fe595d..20407101 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py @@ -20,7 +20,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -33,7 +33,7 @@ def recover_after_failure(_step_context: StepContext) -> str: return "recovered" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "handled-failure") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py index 5daf3b5e..578b9b10 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py @@ -14,7 +14,7 @@ WaitForConditionConfig, WaitForConditionDecision, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def increment_condition( @@ -33,7 +33,7 @@ def stop_after_second_attempt( return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> int: require_scenario(event, "wait-for-condition") return context.wait_for_condition( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py index 3278f3c5..8b65a555 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py @@ -14,7 +14,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import Duration -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario @durable_step @@ -22,7 +22,7 @@ def complete_after_long_wait(_step_context: StepContext) -> str: return "resumed" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-wait") context.wait( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py index 58db3540..59d337b6 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario @durable_step @@ -28,7 +28,7 @@ def succeed_after_long_retry(step_context: StepContext) -> str: return "retried" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-retry") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py index a4078f10..d187b455 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario def submit_callback( @@ -19,7 +19,7 @@ def submit_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-callback") long_delay_seconds(event) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py index a490d44d..170fef6a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py @@ -10,10 +10,10 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import Duration -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler( event: dict[str, Any], context: DurableContext, @@ -26,7 +26,7 @@ def handler( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( event: dict[str, Any], context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 5908dd7d..2d086d7c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -519,5 +519,5 @@ def test_common_selects_the_plugin_from_the_deployed_view() -> None: common: str = (SRC_DIR / "common.py").read_text(encoding="utf-8") assert 'os.environ.get("OTEL_PLUGIN_MODE") == "execution"' in common - assert "ExecutionOtelPlugin(OtelPluginConfig())" in common - assert "InvocationOtelPlugin(OtelPluginConfig())" in common + assert "ExecutionOtelPluginFactory(OtelPluginConfig())" in common + assert "InvocationOtelPluginFactory(OtelPluginConfig())" in common diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py index a8cdc04b..ef684516 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py @@ -89,7 +89,7 @@ def unreliable_operation(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[AttemptPlugin()]) +@durable_execution(plugins=[lambda _info: AttemptPlugin()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py index 1a770cc5..a03e98c7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py @@ -86,7 +86,7 @@ def flaky(step_context: StepContext) -> str: return "ok" -@durable_execution(plugins=[AttemptInfoShapePlugin()]) +@durable_execution(plugins=[lambda _info: AttemptInfoShapePlugin()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py index 0debf47b..1c9f8873 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py @@ -94,7 +94,7 @@ def branch_b(_context: DurableContext) -> str: return "b-done" -@durable_execution(plugins=[ContextInfoShapePlugin()]) +@durable_execution(plugins=[lambda _info: ContextInfoShapePlugin()]) def handler(_event: Any, context: DurableContext) -> list[str]: result: BatchResult[str] = context.parallel( [ diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py index 0a116f88..ac0a44a8 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py @@ -87,7 +87,7 @@ def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[FaultyPlugin()]) +@durable_execution(plugins=[lambda _info: FaultyPlugin()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py index 7abbc401..580100d5 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py @@ -65,7 +65,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[ExternalUpdatePlugin()]) +@durable_execution(plugins=[lambda _info: ExternalUpdatePlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py index ef11896a..c45dee84 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py @@ -187,7 +187,9 @@ def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[FaultyPlugin(), HealthyPlugin()]) +@durable_execution( + plugins=[lambda _info: FaultyPlugin(), lambda _info: HealthyPlugin()] +) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py index 8b5ecefd..55b422a1 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py @@ -47,7 +47,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[FirstInvocationPlugin()]) +@durable_execution(plugins=[lambda _info: FirstInvocationPlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py index 75b53f7b..f6ec386e 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py @@ -58,7 +58,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: _emit(record, info.execution_arn) -@durable_execution(plugins=[InvocationInfoShapePlugin()]) +@durable_execution(plugins=[lambda _info: InvocationInfoShapePlugin()]) def handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return f"done-{event}" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py index ddd22673..9e40133f 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py @@ -56,7 +56,7 @@ def greet(step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[LifecyclePlugin()]) +@durable_execution(plugins=[lambda _info: LifecyclePlugin()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py index f9669a8e..c5f80166 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py @@ -65,7 +65,7 @@ def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[PluginA(), PluginB()]) +@durable_execution(plugins=[lambda _info: PluginA(), lambda _info: PluginB()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py index 82f6ad8b..ec49bd39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py @@ -66,7 +66,7 @@ def child_operation(ctx: DurableContext, name: str) -> str: return ctx.step(greet(name)) -@durable_execution(plugins=[ParentLinkagePlugin()]) +@durable_execution(plugins=[lambda _info: ParentLinkagePlugin()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.run_in_child_context(child_operation(str(event))) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py index 14c36c2a..366d3ebe 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py @@ -63,7 +63,7 @@ def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[OperationChangePlugin()]) +@durable_execution(plugins=[lambda _info: OperationChangePlugin()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py index bbb2988b..e4f86021 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py @@ -83,7 +83,7 @@ def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[OperationChangeShapePlugin()]) +@durable_execution(plugins=[lambda _info: OperationChangeShapePlugin()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py index 07e4a55f..a0294291 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py @@ -79,7 +79,7 @@ def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[OperationInfoShapePlugin()]) +@durable_execution(plugins=[lambda _info: OperationInfoShapePlugin()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py index 48b25cf4..734177c8 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py @@ -77,7 +77,7 @@ def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[OperationLifecyclePlugin()]) +@durable_execution(plugins=[lambda _info: OperationLifecyclePlugin()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py index 2475d4be..7c642d47 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py @@ -83,7 +83,7 @@ def branch1(_ctx: DurableContext) -> str: return "task-2" -@durable_execution(plugins=[ParallelBranchPlugin()]) +@durable_execution(plugins=[lambda _info: ParallelBranchPlugin()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [branch0, branch1], diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py index 12a29056..e0490f69 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py @@ -90,7 +90,7 @@ def step_b(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[ReplayFlagPlugin()]) +@durable_execution(plugins=[lambda _info: ReplayFlagPlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) retry_config = RetryStrategyConfig( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py index e97d4528..f851fce1 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py @@ -99,7 +99,7 @@ def always_fail(_step_context: StepContext) -> str: raise RuntimeError(msg) -@durable_execution(plugins=[RetryExhaustionPlugin()]) +@durable_execution(plugins=[lambda _info: RetryExhaustionPlugin()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=2, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py index 327763ba..496dca5b 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py @@ -57,7 +57,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[SuspensionPlugin()]) +@durable_execution(plugins=[lambda _info: SuspensionPlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py index 2939d09e..237fe227 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py @@ -58,7 +58,7 @@ def failing_step(_step_context: StepContext) -> str: raise RuntimeError(msg) -@durable_execution(plugins=[TerminalFailurePlugin()]) +@durable_execution(plugins=[lambda _info: TerminalFailurePlugin()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step( failing_step(), diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py index e54049c1..849e533e 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py @@ -73,7 +73,7 @@ def step_b(_step_context: StepContext) -> str: raise RuntimeError(msg) -@durable_execution(plugins=[TerminalPayloadPlugin()]) +@durable_execution(plugins=[lambda _info: TerminalPayloadPlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) result: str = context.step( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py index df7e79fa..d18e7958 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py @@ -68,7 +68,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[WaitOperationPlugin()]) +@durable_execution(plugins=[lambda _info: WaitOperationPlugin()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py index eef6d3ee..882cf92e 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py @@ -91,7 +91,7 @@ def wait_long(ctx: DurableContext) -> str: return "long-done" -@durable_execution(plugins=[WaitReplayFlagPlugin()]) +@durable_execution(plugins=[lambda _info: WaitReplayFlagPlugin()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [wait_short, wait_long], diff --git a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py index a727bf73..a94e1409 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py @@ -1,8 +1,9 @@ """Demonstrates OTel-enriched logging in a durable execution. -The InvocationOtelPlugin installs a logging filter on the root logger -(enrich_logger=True by default) when the plugin is constructed. The filter -stamps the active OpenTelemetry trace context (traceId, spanId, +InvocationOtelPluginFactory is the plugin factory the SDK registers; it builds +one InvocationOtelPlugin per invocation. Each plugin installs a logging filter +on the root logger (enrich_logger=True by default) when it is constructed. The +filter stamps the active OpenTelemetry trace context (traceId, spanId, otelTraceSampled) onto every log record that flows through the root handler. This includes logs emitted via context.logger / step_context.logger as well as direct logging.getLogger() calls and third-party library logs, so logs @@ -16,7 +17,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.context import ( @@ -44,7 +45,7 @@ def greet_in_child(child_context: DurableContext, name: str) -> str: return result -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: # Logged at the top level: enriched with the invocation span_id. context.logger.info("Workflow started") diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py index 3d001d46..b61a6420 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py @@ -2,7 +2,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.config import Duration @@ -32,7 +32,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(_event: Any, context: DurableContext) -> int: result = 0 for i in range(3): diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py index d8858baa..56e908b4 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py @@ -51,7 +51,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[MyPlugin()]) +@durable_execution(plugins=[lambda _info: MyPlugin()]) def handler(_event: Any, context: DurableContext) -> int: result: int = context.run_in_child_context( add_numbers_in_child(6, 4), diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py index fb13cc81..72491e90 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py @@ -42,7 +42,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[RecordingWaitPlugin()]) +@durable_execution(plugins=[lambda _info: RecordingWaitPlugin()]) def handler(_event: Any, context: DurableContext) -> dict[str, Any]: context.wait(Duration.from_seconds(1), name="plugin-wait") return { diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index be1027df..98b1383d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -42,6 +42,12 @@ def handler(event, context): ... ``` +`workflow_insight()` returns a plugin *factory*, which is what the SDK's +`plugins` argument takes: the SDK calls it once per invocation to build that +invocation's plugin instance. The factory holds the resolved configuration and +the exporters, so configuration is per handler while record state is per +invocation. + With no exporter configured, records are written to the function's own CloudWatch log group as single JSON lines (the `LambdaLogExporter` default), carrying the name-keyed `operationsByName` summary. The `S3Exporter` writes the @@ -56,8 +62,8 @@ Behavior is validated cross-SDK by the `insight` conformance suite (`aws-durable-execution-conformance-tests-insight`). > **Note (asynchronous export).** Export rendering, truncation, `export()`, and -> `flush()` run on one lazy background worker per plugin. Checkpoint hooks only -> replace the latest pending snapshot and wake the worker. Consecutive +> `flush()` run on one lazy background worker per registered factory. Checkpoint +> hooks only replace the latest pending snapshot and wake the worker. Consecutive > `on-change` snapshots may coalesce while an export is in flight. An invocation > that emits a record drains the latest snapshot and flushes exporters before it > returns; invocations that emit nothing do not start or flush the worker. diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index b06a2853..fdd92db7 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -3,10 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 """Per-execution latest-pending asynchronous export scheduling for Workflow Insight. -One plugin instance serves every execution its environment hosts, and Lambda -Managed Instances makes concurrent executions in one environment routine, so the -pending record is keyed by execution ARN: coalescing happens only within a single -execution and one execution's record can never displace another's. +One scheduler serves every execution its environment hosts -- it is owned by the +handler-lifetime plugin factory, because serializing export is a cross-execution +job -- and Lambda Managed Instances makes concurrent executions in one +environment routine. So the pending record and the bookkeeping that goes with it +live on the per-execution object the caller passes in, which is the caller's own +per-invocation plugin instance (:class:`_ExportState` is mixed into it): +coalescing happens only within a single execution and one execution's record can +never displace another's. The scheduler holds those objects; it has no notion of +an execution ARN and nothing to look up. Export itself stays strictly serialized -- one worker thread, one ``export()`` at a time -- so exporters never see concurrent calls. Parallel export is a later @@ -26,35 +31,61 @@ _logger = logging.getLogger("aws_durable_execution_sdk_python_insight") -class _Lane: - """Per-execution export bookkeeping. One lane per execution ARN.""" +class _ExportState: + """One execution's export bookkeeping, and its slot in the export queue. - __slots__ = ("scheduled_seq", "exported_seq", "exported_at", "waiters") + Mixed into the per-invocation plugin instance, so one object carries both an + execution's hook-facing state and its export bookkeeping. Those used to live + in three ARN-keyed structures -- the plugin's execution registry, the + scheduler's pending record and its per-execution lane -- three views of one + execution that had to agree about whether it still had work outstanding. In + Java the same shape produced a defect where two of the views disagreed. There + is one view now, and the scheduler holds the object itself. + + Every field here is guarded by ``_ExportScheduler._condition``. They belong to + the scheduler: nothing outside it reads or writes them, and it never touches + the hook-facing state on the same object. + """ def __init__(self) -> None: - # Newest sequence number scheduled for this execution. + # Newest sequence number the scheduler assigned to this execution. self.scheduled_seq = 0 + # This execution's latest record, waiting for the export worker, or None + # when nothing of its own is outstanding. A repeat emission replaces it, + # which is what per-execution coalescing means; the scheduler's queue + # holds this object exactly while this field is set. + self.pending_record: dict[str, Any] | None = None # Newest sequence number already handed to every exporter. self.exported_seq = 0 # Value of the scheduler's export counter when that export finished, so # a waiter can tell whether a completed flush covered its own record. self.exported_at = 0 - # drain() calls currently blocked on this lane; the lane is only - # forgotten once nobody is waiting on it. + # drain() calls currently parked on this execution. Nothing depends on + # it: a waiter holds this object directly, so the bookkeeping it waits + # on can no longer be reclaimed from under it. It is kept because it is + # the only way to observe that a drain really parked rather than raced + # past. self.waiters = 0 +# What a caller has to release once it is back outside the lock: the records the +# `_disabled` latch dropped, each with the execution object that was carrying it. +# Both can run customer finalizers. +_Dropped = list[tuple[_ExportState, dict[str, Any] | None]] + + class _ExportScheduler: """Run all exporters on one lazy worker, keeping the latest record per execution.""" def __init__(self, exporters: list[InsightExporter]) -> None: self._exporters = exporters self._condition = threading.Condition(threading.Lock()) - # execution ARN -> (sequence, latest record), oldest arrival first. A - # repeat schedule for an ARN replaces the value and keeps the position, - # so coalescing never lets one execution jump the queue. - self._pending: dict[str, tuple[int, dict[str, Any]]] = {} - self._lanes: dict[str, _Lane] = {} + # Executions with a record waiting, oldest arrival first -- an ordered + # set, keyed by the execution object itself. A repeat schedule for an + # execution replaces the record the object carries and keeps the object's + # position, so coalescing never lets one execution jump the queue. An + # execution is in here exactly while its `pending_record` is set. + self._pending: dict[_ExportState, None] = {} self._seq = 0 self._export_count = 0 # Highest export counter value covered by a completed flush. @@ -86,19 +117,21 @@ def __init__(self, exporters: list[InsightExporter]) -> None: self._worker: threading.Thread | None = None self._disabled = False - def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: + def schedule(self, execution: _ExportState, record: dict[str, Any]) -> None: """Replace this execution's pending snapshot; never runs exporters inline.""" - displaced: tuple[int, dict[str, Any]] | None = None - failed_pending: dict[str, tuple[int, dict[str, Any]]] | None = None + displaced: dict[str, Any] | None = None + failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: if self._disabled: return self._seq += 1 - lane = self._lane_locked(execution_arn) - lane.scheduled_seq = self._seq - displaced = self._pending.get(execution_arn) - self._pending[execution_arn] = (self._seq, record) + execution.scheduled_seq = self._seq + displaced = execution.pending_record + execution.pending_record = record + # Re-queuing an execution that is already queued is a no-op that + # keeps its arrival position. + self._pending[execution] = None failed_pending, start_error = self._ensure_worker_locked() self._condition.notify_all() # Releasing either record may run custom finalizers, so do it unlocked. @@ -110,7 +143,7 @@ def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: start_error, ) - def drain(self, execution_arn: str) -> None: + def drain(self, execution: _ExportState) -> None: """Wait until this execution's latest record is exported and exporters flush. Returns once the calling execution's own record has reached every exporter @@ -122,21 +155,25 @@ def drain(self, execution_arn: str) -> None: Java have. Concurrent calls can share one flush, since they all started before it completed. + The caller passes the execution object rather than an ARN, so there is + nothing to look up and nothing to create: an execution that never + scheduled a record simply carries zeroed bookkeeping, which is exactly + "nothing of my own is outstanding, flush and return". + Two paths return without exporting or flushing anything, because the permanent ``_disabled`` latch means no record will ever be exported: the latch was already set when this call started, or it is set while this call is parked. Failing to start the export worker sets that latch, so a drain that hits a worker-start failure also returns without a flush. """ - failed_pending: dict[str, tuple[int, dict[str, Any]]] | None = None + failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: if self._disabled: return - lane = self._lane_locked(execution_arn) - lane.waiters += 1 + execution.waiters += 1 try: - want_seq = lane.scheduled_seq + want_seq = execution.scheduled_seq # A drain always flushes, so require a flush that covers every # export completed before this call as well as our own. want_flush = self._export_count @@ -148,10 +185,11 @@ def drain(self, execution_arn: str) -> None: # Export counter value a flush has to cover to release us: # our own record's export plus everything already exported # when this call started. Recomputed every pass, because - # lane.exported_at only becomes ours once our record is out. - need = max(lane.exported_at, want_flush) + # execution.exported_at only becomes ours once our record is + # out. + need = max(execution.exported_at, want_flush) if ( - lane.exported_seq >= want_seq + execution.exported_seq >= want_seq and self._flushed_through >= need and self._flushes_completed > want_flushes ): @@ -182,8 +220,7 @@ def drain(self, execution_arn: str) -> None: self._condition.notify_all() self._condition.wait() finally: - lane.waiters -= 1 - self._forget_lane_locked(execution_arn, lane) + execution.waiters -= 1 del failed_pending if start_error is not None: _logger.warning( @@ -194,30 +231,7 @@ def drain(self, execution_arn: str) -> None: # -- internals ------------------------------------------------------------ - def _lane_locked(self, execution_arn: str) -> _Lane: - lane = self._lanes.get(execution_arn) - if lane is None: - lane = _Lane() - self._lanes[execution_arn] = lane - return lane - - def _forget_lane_locked(self, execution_arn: str, lane: _Lane) -> None: - # Keep the lane while anything still depends on it; bookkeeping for a - # fully exported execution with no waiters is safe to drop, because a - # later drain then only needs a flush covering the exports so far. - if self._lanes.get(execution_arn) is not lane: - return - if lane.waiters: - return - if execution_arn in self._pending: - return - if lane.exported_seq < lane.scheduled_seq: - return - del self._lanes[execution_arn] - - def _ensure_worker_locked( - self, - ) -> tuple[dict[str, tuple[int, dict[str, Any]]] | None, Exception | None]: + def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: if self._worker is not None and self._worker.is_alive(): return None, None worker = threading.Thread( @@ -231,12 +245,19 @@ def _ensure_worker_locked( except Exception as exc: # noqa: BLE001 - instrumentation must not escape hooks self._disabled = True self._worker = None - failed_pending = self._pending + # Nothing is retained once the plugin has given up on asynchronous + # export for good: the queue is the scheduler's only per-execution + # structure, so emptying it drops every reference it holds. Both the + # records and the execution objects go back to the CALLER to release + # outside the lock -- a record can carry customer objects whose + # finalizers run arbitrary code, and so can an execution whose hook + # state the plugin has already discarded. + failed_pending = [ + (execution, execution.pending_record) for execution in self._pending + ] + for execution, _ in failed_pending: + execution.pending_record = None self._pending = {} - # Lanes hold plain counters, never customer objects, so they can be - # dropped under the lock. Nothing is retained once the plugin has - # given up on asynchronous export for good. - self._lanes = {} self._flush_requested = False self._flush_barrier = 0 self._flush_in_flight = 0 @@ -249,7 +270,7 @@ def _ensure_worker_locked( def _blocking_pending_locked(self) -> bool: """True while a record scheduled at or before the flush barrier is pending.""" barrier = self._flush_barrier - return any(seq <= barrier for seq, _ in self._pending.values()) + return any(execution.scheduled_seq <= barrier for execution in self._pending) def _run(self) -> None: # The worker slot must be empty whenever no worker is running, or @@ -271,7 +292,7 @@ def _run(self) -> None: def _run_loop(self) -> None: while True: - arn: str | None = None + execution: _ExportState | None = None seq = 0 record: dict[str, Any] | None = None flush_covers = 0 @@ -287,36 +308,39 @@ def _run_loop(self) -> None: self._flush_in_flight = flush_covers break if self._pending: - arn, (seq, record) = next(iter(self._pending.items())) - del self._pending[arn] + execution = next(iter(self._pending)) + del self._pending[execution] + record = execution.pending_record + execution.pending_record = None + seq = execution.scheduled_seq break self._condition.wait() if record is not None: - # Popping the record consumed this execution's pending slot, so - # nothing will ever export that snapshot again. The lane must - # therefore advance whatever export() did: skip it and - # lane.exported_seq never reaches a waiter's want_seq, so a drain - # parked on this execution is never released. _export() already - # contains every Exception, but a BaseException from a customer - # exporter unwinds through here. Count the attempt in a finally - # and let the exception continue out to the wrapper -- and into - # the thread's traceback -- with nothing swallowed. + assert execution is not None + # Taking the record consumed this execution's pending slot, so + # nothing will ever export that snapshot again. The bookkeeping + # must therefore advance whatever export() did: skip it and + # execution.exported_seq never reaches a waiter's want_seq, so a + # drain parked on this execution is never released. _export() + # already contains every Exception, but a BaseException from a + # customer exporter unwinds through here. Count the attempt in a + # finally and let the exception continue out to the wrapper -- + # and into the thread's traceback -- with nothing swallowed. + # + # The record and its bookkeeping are one object, so there is no + # second lookup left to come back empty: publishing cannot miss. try: self._export(record) finally: # Release the exported record before re-locking: a custom # finalizer may re-enter schedule(). del record - assert arn is not None with self._condition: self._export_count += 1 - lane = self._lanes.get(arn) - if lane is not None: - if seq > lane.exported_seq: - lane.exported_seq = seq - lane.exported_at = self._export_count - self._forget_lane_locked(arn, lane) + if seq > execution.exported_seq: + execution.exported_seq = seq + execution.exported_at = self._export_count self._condition.notify_all() continue diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 0da43e34..35f8e024 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -9,8 +9,20 @@ record keeps the JS camelCase field names so records read identically across SDKs. +Two lifetimes: + ``workflow_insight(config)`` returns a FACTORY, which is what + ``@durable_execution(plugins=[...])`` takes. The factory lives as long as the + handler and owns everything that is not per-execution: the resolved immutable + config, the exporters, and the ``_ExportScheduler`` that serializes export + across executions. The SDK calls it once per invocation and drops the instance + it returns when the invocation scope exits, so a + :class:`WorkflowInsightPlugin` instance serves exactly one invocation of one + execution. Everything this environment holds for that execution is therefore + ordinary instance state: no ARN-keyed registry, and no hook can reach an + instance other than its own. + Operation-map sourcing: - The Python SDK invocation hooks now carry the full operation map directly: + The Python SDK invocation hooks carry the full operation map directly: ``InvocationStartInfo.operations`` (a point-in-time snapshot at invocation start), ``InvocationEndInfo.operations`` (a fresh snapshot at invocation end), and ``OperationChangeInfo.operations`` (the full map at the change). Alongside @@ -19,8 +31,8 @@ snapshots as the authoritative operation state -- it does NOT reconstruct the map by accumulating per-operation ``on_operation_end`` events. Because every invocation start re-seeds the map from the snapshot, a cold resume in a fresh - Lambda environment (a brand-new plugin instance) still reports the prior - terminal operations. + Lambda environment (a brand-new instance, as every invocation now gets) still + reports the prior terminal operations. The Python SDK has no ``pluginsConfig.childOperationsDepth`` equivalent, so ``full-tree`` records rely on the child operations being present in the @@ -46,7 +58,10 @@ OperationType, ) -from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportScheduler +from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _ExportScheduler, + _ExportState, +) from aws_durable_execution_sdk_python_insight.exporters.lambda_log_exporter import ( LambdaLogExporter, ) @@ -159,200 +174,150 @@ def _apply_result_override( return None -class _ExecutionState: - __slots__ = ( - "start_time", - "parsed_arn", - "cached_input", - "operations", - "closed", - "lock", - ) - - def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: - self.start_time = start_time - self.parsed_arn = parsed_arn - self.cached_input: Any = None +class WorkflowInsightPlugin(DurableInstrumentationPlugin, _ExportState): + """Everything this environment holds for one invocation of one execution. + + Built by the factory ``workflow_insight()`` returns, once per invocation, + from that invocation's ``InvocationStartInfo`` -- the same object its + ``on_invocation_start`` then receives. Identity comes from that info and is + never revised afterwards: the execution ARN, the sampling decision, the + execution start time and the cached input. + + Per-execution state is plain instance state, and the object doubles as its + own export queue entry (via :class:`_ExportState`). Where three ARN-keyed + structures had to agree about one execution -- this plugin's registry, the + scheduler's pending record and its per-execution lane -- there is now one + object and no ARN to resolve. In Java the same three-view shape produced a + defect where two of the views disagreed. + + Two locks, disjoint field sets, and neither is ever taken to reach the + other's fields: + + * ``_lock`` guards ``_closed``, the ``_operations`` rebind and record + emission. + * ``_ExportScheduler._condition``'s lock guards the export bookkeeping this + instance carries for the scheduler; nothing outside the scheduler reads or + writes those fields, and the scheduler never touches the fields above. + + Shared, handler-lifetime state -- resolved config, exporters, scheduler -- + lives on ``_shared`` and is read-only from here. + """ + + def __init__( + self, shared: _WorkflowInsightFactory, info: InvocationStartInfo + ) -> None: + _ExportState.__init__(self) + self._shared = shared + execution_arn = info.execution_arn or "" + self._execution_arn = execution_arn + # Deterministic per-ARN, so the decision could be recomputed on every + # hook; taken once here because the instance now has a place to keep it, + # and because an unsampled instance then never parses the ARN. + self._sampled_in = bool(execution_arn) and _should_sample( + execution_arn, shared._sampling_rate + ) + self._parsed_arn: dict[str, str] = ( + _parse_execution_arn(execution_arn) if self._sampled_in else {} + ) + # Always the service-provided execution start time when present, + # including on a cold resume in a fresh environment -- never the resume + # time, which would corrupt duration and the date partition. `now` is the + # fallback for an info that carries no start time at all. + self._start_time: Any = ( + info.execution_start_time + if info.execution_start_time is not None + else datetime.datetime.now(datetime.UTC) + ) + self._cached_input: Any = info.execution_input # operation_id -> OperationInfo, adopted verbatim from the SDK's # authoritative snapshot (invocation start/end and operation-change). - self.operations: dict[str, OperationInfo] = {} - # Set once the invocation this state belongs to 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). - self.closed = False - # Guards `closed`, the operations rebind and record emission for this - # execution, so a late hook can never slip a RUNNING record in after the - # terminal one. Per execution, so concurrent executions never contend. + 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). + self._closed = False + # Guards `_closed`, the operations rebind and record emission, so a late + # hook can never slip a RUNNING record in after the terminal one. Still + # earns its place with one instance per invocation: the SDK dispatches + # every hook synchronously on the thread that produced the event, so an + # operation-change raised off the checkpointing path runs concurrently + # with the invocation thread's on_invocation_end -- two hooks, one + # instance, genuinely racing. + # # Reentrant on purpose: `_emit` runs the scheduler's `schedule()` inside # this hold, and `schedule()` releases the record it displaces, which can # run a customer finalizer that re-enters a hook for this same execution # on this same thread. A plain lock self-deadlocks the invocation thread # there. (Java holds no such lock: its ExecutionState carries no # operations map, and `cachedInput` is a bare volatile field.) - self.lock = threading.RLock() - - -class WorkflowInsightPlugin(DurableInstrumentationPlugin): - def __init__(self, config: WorkflowInsightConfig) -> None: - self._sampling_rate = _resolve_sampling_rate(config.sampling_rate) - # config.emit_mode / operation_detail are already normalized to enum - # members (or None) by WorkflowInsightConfig.__post_init__; re-wrap to - # satisfy the static type of the union-typed config fields. - self._emit_mode: EmitMode = ( - EmitMode(config.emit_mode) - if config.emit_mode is not None - else EmitMode.ON_COMPLETE - ) - detail = ( - OperationDetail(config.operation_detail) - if config.operation_detail is not None - else OperationDetail.TOP_LEVEL - ) - self._top_level_only = detail != OperationDetail.FULL_TREE - content: ContentConfig | None = config.content - self._content = content - ops = content.operations if content and content.operations else None - self._include_errors = ( - True if ops is None or ops.include_errors is None else ops.include_errors - ) - self._overrides_by_name: dict[str, OperationOverride] = {} - if ops is not None: - for override in ops.overrides: - self._overrides_by_name[override.operation_name] = override - # Default-exporter parity with the JS plugin: an omitted OR an explicitly - # empty exporter list falls back to the Lambda log exporter, so the - # plugin is never a silent no-op. A non-empty list is used verbatim. - self._exporters: list[InsightExporter] = ( - list(config.exporters) if config.exporters else [LambdaLogExporter()] - ) - self._scheduler = _ExportScheduler(self._exporters) - self._state: dict[str, _ExecutionState] = {} - self._lock = threading.Lock() - - # -- sampling / state ----------------------------------------------------- - - def _sampled_in(self, execution_arn: str) -> bool: - # Deterministic per-ARN, so every hook for one execution agrees without - # needing to persist the decision in state. - return _should_sample(execution_arn, self._sampling_rate) - - def _ensure_state(self, execution_arn: str) -> _ExecutionState: - with self._lock: - state = self._state.get(execution_arn) - if state is None: - state = _ExecutionState( - start_time=datetime.datetime.now(datetime.UTC), - parsed_arn=_parse_execution_arn(execution_arn), - ) - self._state[execution_arn] = state - return state - - def _get_state(self, execution_arn: str) -> _ExecutionState | None: - # Lookup only. A hook that must never fabricate state (an - # operation-change arriving after the invocation ended, whose state has - # been discarded) uses this instead of _ensure_state. - with self._lock: - return self._state.get(execution_arn) + self._lock = threading.RLock() - def _discard_state(self, execution_arn: str) -> None: - with self._lock: - self._state.pop(execution_arn, None) + # -- state ---------------------------------------------------------------- - def _adopt_operations_locked( - self, state: _ExecutionState, operations: dict[str, OperationInfo] - ) -> None: + def _adopt_operations_locked(self, operations: dict[str, OperationInfo]) -> None: # Adopt the authoritative point-in-time snapshot. Copy so plugin state # never aliases the SDK-owned map, and rebind the attribute so a # concurrent reader holding the prior reference iterates a stable dict. - # Callers hold state.lock. - state.operations = dict(operations) + # Callers hold self._lock. + self._operations = dict(operations) # -- hooks ---------------------------------------------------------------- def on_invocation_start(self, info: InvocationStartInfo) -> None: - arn = info.execution_arn - if not arn or not self._sampled_in(arn): + if not self._sampled_in: return - state = self._ensure_state(arn) - with state.lock: - if state.closed: + with self._lock: + if self._closed: return - # Always adopt the service-provided execution start time when - # present, including a cold resume in a fresh environment (never the - # resume time, which would corrupt duration and the date partition). - if info.execution_start_time is not None: - state.start_time = info.execution_start_time - state.cached_input = info.execution_input - # Seed the operation map from the full snapshot on every invocation. - # On a cold resume this rebuilds prior (terminal) operations that a - # fresh plugin instance never saw via per-operation hooks. - self._adopt_operations_locked(state, info.operations) - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + # Seed the operation map from the full snapshot. On a cold resume + # this rebuilds prior (terminal) operations that a fresh instance + # never saw via per-operation hooks. + self._adopt_operations_locked(info.operations) + if self._shared._emit_mode == EmitMode.ON_CHANGE: + self._emit(status="RUNNING", end_time=None, output_raw=None, error=None) def on_operation_change(self, info: OperationChangeInfo) -> None: - arn = info.execution_arn - if not arn or not self._sampled_in(arn): - return - # Never create state here. A change hook for a checkpoint that completed - # just before the invocation ended still arrives after on_invocation_end - # discarded the state; recreating it would fabricate start_time = now, - # emit a RUNNING record after the terminal one, and leave a state entry - # behind for an execution this environment no longer runs. - state = self._get_state(arn) - if state is None: + # No ARN to resolve: this instance belongs to the invocation the change + # was raised in, so the hook's `execution_arn` is this execution's by + # construction. A change hook can no longer fabricate state for an + # execution whose invocation has ended -- there is no registry to + # fabricate it in, and the instance it reaches is its own. + if not self._sampled_in: return - with state.lock: - if state.closed: + with self._lock: + if self._closed: return # Replace state with the full operations snapshot carried by the hook. - self._adopt_operations_locked(state, info.operations) + self._adopt_operations_locked(info.operations) # on-change mode exports an updated RUNNING record on each change so # mid-invocation progress is observable, not only at start/end. - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + if self._shared._emit_mode == EmitMode.ON_CHANGE: + self._emit(status="RUNNING", end_time=None, output_raw=None, error=None) def on_invocation_end(self, info: InvocationEndInfo) -> None: - arn = info.execution_arn - if not arn: + if not self._sampled_in: + # Sampled-out executions process nothing: they neither export nor + # flush, so instrumenting a fraction of executions costs the rest + # nothing. return - if not self._sampled_in(arn): - # Sampled-out executions process no operations and retain no state. - self._discard_state(arn) - return - state = self._ensure_state(arn) - with state.lock: - if not state.closed: + emit_mode = self._shared._emit_mode + with self._lock: + if not self._closed: # Close the gate before emitting so a concurrent late hook for # this execution cannot append a RUNNING record after the # terminal one. - state.closed = True + self._closed = True # Refresh from the fresh end-of-invocation snapshot before # emitting so the terminal record reflects the final operation # map. - self._adopt_operations_locked(state, info.operations) + self._adopt_operations_locked(info.operations) status = _STATUS_MAP.get(info.status, "RUNNING") is_terminal = status in ("SUCCEEDED", "FAILED") is_failure = status == "FAILED" - if self._emit_mode == EmitMode.ON_CHANGE: + if emit_mode == EmitMode.ON_CHANGE: should_emit = True - elif self._emit_mode == EmitMode.ON_FAILURE: + elif emit_mode == EmitMode.ON_FAILURE: should_emit = is_failure else: # on-complete should_emit = is_terminal @@ -364,8 +329,6 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # end_time=None makes _emit drop both fields. Output and # error likewise belong only to a terminal record. self._emit( - arn, - state, status=status, end_time=datetime.datetime.now(datetime.UTC) if is_terminal @@ -373,17 +336,16 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: output_raw=info.execution_result if is_terminal else None, error=info.error if is_terminal else None, # This is the emit that closed the gate, so it always runs - # with `closed` already set and must never drop itself. + # with `_closed` already set and must never drop itself. closing=True, ) - # Clear state after EVERY invocation end, including PENDING/RETRY. The - # next invocation rebuilds it from InvocationStartInfo.operations, so a - # suspended execution that never resumes in this environment (or that was - # sampled out) leaks nothing and state stays bounded. Done before the - # drain so a late hook for this execution finds no state to update while - # the terminal record is still in flight. - self._discard_state(arn) + # Nothing has to be cleared after an invocation end, including a + # PENDING/RETRY one: this instance IS the state, and the SDK drops it + # when the invocation scope exits. A suspended execution that resumes + # here later gets a fresh instance, seeded from + # InvocationStartInfo.operations. + # # Drain on EVERY sampled-in invocation end, emitted record or not: JS and # Java flush once per sampled-in invocation end regardless, and a # buffering exporter has to see the same rhythm in all three languages @@ -391,26 +353,28 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # may still be holding records another execution handed it). A sampled-out # execution returns above, so it neither exports nor flushes. # - # The drain covers this execution only: it returns once this execution's + # The drain covers this execution only -- it names this instance, which + # carries its own export bookkeeping: it returns once this execution's # own record, if any, reached the exporters and a flush that completed # after this call is done, without waiting on records scheduled after the # call by other executions. - self._scheduler.drain(arn) + self._shared._scheduler.drain(self) # -- emission ------------------------------------------------------------- def _build_operations( self, operations: dict[str, OperationInfo] ) -> list[dict[str, Any]]: + shared = self._shared records: list[dict[str, Any]] = [] for op in operations.values(): if op.operation_type == OperationType.EXECUTION: continue if not op.name: continue - if self._top_level_only and op.parent_id: + if shared._top_level_only and op.parent_id: continue - override = self._overrides_by_name.get(op.name) + override = shared._overrides_by_name.get(op.name) if override is not None and override.exclude: continue @@ -432,7 +396,7 @@ def _build_operations( entry["durationMs"] = dur if op.attempt is not None: entry["attempt"] = op.attempt - if self._include_errors and op.error is not None: + if shared._include_errors and op.error is not None: entry["error"] = {"name": op.error.type, "message": op.error.message} if override is not None and override.result is not None: value = _apply_result_override(override.result, op.result) @@ -443,8 +407,6 @@ def _build_operations( def _emit( self, - execution_arn: str, - state: _ExecutionState, *, status: str, end_time: Any, @@ -452,21 +414,21 @@ def _emit( error: Any, closing: bool = False, ) -> None: - arn = state.parsed_arn - start_time = state.start_time + arn = self._parsed_arn + start_time = self._start_time duration = _duration_ms(start_time, end_time) # Snapshot the operations reference once so a concurrent adopt() rebind # cannot change the map mid-build. - operations = state.operations + operations = self._operations - content = self._content + content = self._shared._content record: dict[str, Any] = { "recordType": "WorkflowInsight", "schemaVersion": "1.0", "emittedAt": datetime.datetime.now(datetime.UTC) .isoformat() .replace("+00:00", "Z"), - "executionArn": execution_arn, + "executionArn": self._execution_arn, } if arn.get("executionName"): record["executionName"] = arn["executionName"] @@ -491,7 +453,7 @@ def _emit( except (json.JSONDecodeError, TypeError): parsed_output = output_raw input_value = _apply_data_content( - state.cached_input, content.input if content else None + self._cached_input, content.input if content else None ) output_value = _apply_data_content( parsed_output, content.output if content else None @@ -508,29 +470,93 @@ def _emit( # execution's closing record -- the exporters never see a RUNNING record # follow the terminal one for the same execution. # - # Re-check the gate here, because each hook's own `if state.closed` is a - # check-then-act and this is the act. Everything between the two runs - # customer code while holding state.lock: the input/output transforms + # Re-check the gate here, because each hook's own `if self._closed` is a + # check-then-act and this is the act. One instance per invocation removed + # one of the two windows this used to cover -- state can no longer be + # discarded and recreated underneath a hook, because there is no registry + # to recreate it in -- but not the other: everything between the two runs + # customer code while holding self._lock (the input/output transforms # above, a result override in _build_operations, and __del__ on any object - # the record carries. state.lock is a reentrant RLock on purpose (so that - # customer code re-entering a hook on this thread does not self-deadlock), - # which means such a re-entrant call can run on_invocation_end all the way - # through -- set `closed`, emit the terminal record, discard the state and - # drain it -- and then return here. Without this re-check the outer frame - # hands its already-built RUNNING record to the scheduler afterwards, and - # one execution exports ['SUCCEEDED', 'RUNNING'] from a single hook call, - # no concurrency required. + # the record carries), and self._lock is a reentrant RLock on purpose (so + # that customer code re-entering a hook on this thread does not + # self-deadlock). Such a re-entrant call can therefore run + # on_invocation_end all the way through -- set `_closed`, emit the + # terminal record and drain -- and then return here. Without this + # re-check the outer frame hands its already-built RUNNING record to the + # scheduler afterwards, and one execution exports + # ['SUCCEEDED', 'RUNNING'] from a single hook call, no concurrency + # required. # # `closing` marks the emit that set the gate (on_invocation_end's own), - # which by construction always runs with `closed` set and must not drop + # which by construction always runs with `_closed` set and must not drop # itself. It is not the same test as "the record is terminal": in # on-change mode a PENDING/RETRY invocation end legitimately emits a # RUNNING record, and that record is the closing one. - if state.closed and not closing: + if self._closed and not closing: return - self._scheduler.schedule(execution_arn, record) + self._shared._scheduler.schedule(self, record) + + +class _WorkflowInsightFactory: + """The handler-lifetime half of the plugin: what is NOT per-execution. + + Satisfies the SDK's ``DurableInstrumentationPluginFactory`` -- it is called + with an ``InvocationStartInfo`` and returns the plugin instance for that + invocation. Everything it holds is either immutable after construction (the + resolved config) or deliberately shared across executions: + + * the exporters, which are customer objects registered once, and + * the ``_ExportScheduler``, because export serialization is cross-execution: + one worker, one ``export()`` at a time, whatever the instance that + scheduled the record. + + A callable class rather than a closure so the resolved config stays + inspectable (``factory._emit_mode``, ``factory._exporters``) instead of being + buried in cell variables. + """ + + def __init__(self, config: WorkflowInsightConfig) -> None: + self._sampling_rate = _resolve_sampling_rate(config.sampling_rate) + # config.emit_mode / operation_detail are already normalized to enum + # members (or None) by WorkflowInsightConfig.__post_init__; re-wrap to + # satisfy the static type of the union-typed config fields. + self._emit_mode: EmitMode = ( + EmitMode(config.emit_mode) + if config.emit_mode is not None + else EmitMode.ON_COMPLETE + ) + detail = ( + OperationDetail(config.operation_detail) + if config.operation_detail is not None + else OperationDetail.TOP_LEVEL + ) + self._top_level_only = detail != OperationDetail.FULL_TREE + content: ContentConfig | None = config.content + self._content = content + ops = content.operations if content and content.operations else None + self._include_errors = ( + True if ops is None or ops.include_errors is None else ops.include_errors + ) + self._overrides_by_name: dict[str, OperationOverride] = {} + if ops is not None: + for override in ops.overrides: + self._overrides_by_name[override.operation_name] = override + # Default-exporter parity with the JS plugin: an omitted OR an explicitly + # empty exporter list falls back to the Lambda log exporter, so the + # plugin is never a silent no-op. A non-empty list is used verbatim. + self._exporters: list[InsightExporter] = ( + list(config.exporters) if config.exporters else [LambdaLogExporter()] + ) + self._scheduler = _ExportScheduler(self._exporters) + + def __call__(self, info: InvocationStartInfo) -> WorkflowInsightPlugin: + return WorkflowInsightPlugin(self, info) + +def workflow_insight(config: WorkflowInsightConfig) -> _WorkflowInsightFactory: + """Create a Workflow Insight plugin factory. Mirrors the JS ``workflowInsight()``. -def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPlugin: - """Create a Workflow Insight plugin. Mirrors the JS ``workflowInsight()`` factory.""" - return WorkflowInsightPlugin(config) + Pass the result straight to ``@durable_execution(plugins=[...])``: the SDK + calls it once per invocation to build that invocation's plugin instance. + """ + return _WorkflowInsightFactory(config) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py index a52bf168..e5633c21 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py @@ -66,10 +66,13 @@ def _insight_handler(event: Any, context: DurableContext) -> str: # noqa: ARG00 def test_terminal_record_includes_prior_step_and_completed_wait() -> None: capture = _CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[capture])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[capture])) # Functional form (not the decorator-factory form) so the wrapped handler's - # static type stays a plain 2-arg callable for the runner. - handler = durable_execution(_insight_handler, plugins=[plugin]) + # static type stays a plain 2-arg callable for the runner. `factory` is the + # plugin factory the SDK calls once per invocation, so the two invocations + # this test drives run on two instances -- which is what makes the assertions + # below about the resuming invocation meaningful. + handler = durable_execution(_insight_handler, plugins=[factory]) with DurableFunctionTestRunner(handler=handler, execution_timeout=15) as runner: result: DurableFunctionTestResult = runner.run(input="{}") diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index cfad4097..cbca9d4f 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -10,6 +10,7 @@ from aws_durable_execution_sdk_python_insight._export_scheduler import ( _ExportScheduler, + _ExportState, ) @@ -18,6 +19,37 @@ ARN_B = ARN.format("b") +class _ArnScheduler(_ExportScheduler): + """Supplies the ARN -> execution map that the SDK's plugin lifecycle provides. + + ``schedule()`` and ``drain()`` take the per-execution object: the scheduler + holds the object the caller already has -- in production the caller's own + per-invocation plugin instance, which carries its export bookkeeping as an + ``_ExportState`` -- instead of resolving an ARN to bookkeeping of its own, so + it does not know what an ARN is. These tests drive the scheduler directly, so + they own an ARN map of their own and are otherwise unchanged. + """ + + def __init__(self, exporters: list[Any]) -> None: + super().__init__(exporters) + self.executions: dict[str, _ExportState] = {} + self._executions_lock = threading.Lock() + + def _execution(self, execution_arn: str) -> _ExportState: + with self._executions_lock: + execution = self.executions.get(execution_arn) + if execution is None: + execution = _ExportState() + self.executions[execution_arn] = execution + return execution + + def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: # type: ignore[override] + super().schedule(self._execution(execution_arn), record) + + def drain(self, execution_arn: str) -> None: # type: ignore[override] + super().drain(self._execution(execution_arn)) + + def _record(value: str) -> dict[str, Any]: return {"status": "RUNNING", "value": value, "operations": []} @@ -87,7 +119,7 @@ def export(self, record: dict[str, Any]) -> None: def test_latest_pending_coalesces_within_one_execution() -> None: exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _record("first")) assert exporter.started.wait(5.0) @@ -110,7 +142,7 @@ def test_latest_pending_coalesces_within_one_execution() -> None: def test_exporter_failure_does_not_block_other_exporters() -> None: failing = FailingExporter() capture = CaptureExporter() - scheduler = _ExportScheduler([failing, capture]) + scheduler = _ArnScheduler([failing, capture]) scheduler.schedule(ARN_A, _record("terminal")) @@ -127,7 +159,7 @@ def test_base_exception_from_export_still_releases_drain() -> None: # invocation thread, parks forever. Every wait here is bounded so a # regression fails instead of hanging the suite. exporter = BaseExceptionExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _record("terminal")) returned = threading.Event() @@ -149,7 +181,7 @@ def drain() -> None: def test_drain_flushes_after_export() -> None: capture = CaptureExporter() - scheduler = _ExportScheduler([capture]) + scheduler = _ArnScheduler([capture]) scheduler.schedule(ARN_A, _record("terminal")) @@ -162,7 +194,7 @@ def fail_start(self) -> None: # noqa: ARG001 raise RuntimeError("cannot start") monkeypatch.setattr(threading.Thread, "start", fail_start) - scheduler = _ExportScheduler([CaptureExporter()]) + scheduler = _ArnScheduler([CaptureExporter()]) scheduler.schedule(ARN_A, _record("dropped")) scheduler.drain(ARN_A) @@ -171,7 +203,7 @@ def fail_start(self) -> None: # noqa: ARG001 def test_superseded_record_finalizes_after_lane_unlock() -> None: - scheduler = _ExportScheduler([BlockingExporter()]) + scheduler = _ArnScheduler([BlockingExporter()]) exporter = scheduler._exporters[0] assert isinstance(exporter, BlockingExporter) scheduler.schedule(ARN_A, _record("inflight")) @@ -197,7 +229,7 @@ def __del__(self) -> None: def test_drain_waits_for_blocked_exporter() -> None: exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _record("terminal")) assert exporter.started.wait(5.0) drain_thread = threading.Thread(target=scheduler.drain, args=(ARN_A,)) @@ -286,14 +318,13 @@ def _execution_record(arn: str, status: str) -> dict[str, Any]: def _scheduler_is_empty(scheduler: _ExportScheduler) -> bool: with scheduler._condition: - return not scheduler._pending and not scheduler._lanes + return not scheduler._pending -def _drain_waiters(scheduler: _ExportScheduler, arn: str) -> int: - """How many drain() calls are currently parked on this execution's lane.""" +def _drain_waiters(scheduler: _ArnScheduler, arn: str) -> int: + """How many drain() calls are currently parked on this execution.""" with scheduler._condition: - lane = scheduler._lanes.get(arn) - return 0 if lane is None else lane.waiters + return scheduler._execution(arn).waiters def test_drain_stays_parked_until_a_flush_covering_its_record_completes() -> None: @@ -303,7 +334,7 @@ def test_drain_stays_parked_until_a_flush_covering_its_record_completes() -> Non # inside flush() makes that deterministic -- while the gate is held the flush # provably cannot have completed, so a drain that returns is a violation. exporter = FlushGateEventLogExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) returned = threading.Event() @@ -338,7 +369,7 @@ def test_no_redundant_flush_runs_after_drain_returned() -> None: # been consumed and the coverage is not published yet) and that second flush # calls the exporters after drain(), and with it the invocation, returned. exporter = FlushGateEventLogExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) returned = threading.Event() flushes_at_return: list[int] = [] @@ -381,7 +412,7 @@ def test_drain_with_nothing_to_export_still_flushes_exactly_once() -> None: # until some other execution happened to flush. It has to flush once and # return. capture = CaptureExporter() - scheduler = _ExportScheduler([capture]) + scheduler = _ArnScheduler([capture]) returned = threading.Event() def drain() -> None: @@ -405,7 +436,7 @@ def test_drain_never_rides_on_a_flush_that_finished_before_it_started() -> None: # a flush that completed after it was called, so that one invocation end # means one flush. capture = CaptureExporter() - scheduler = _ExportScheduler([capture]) + scheduler = _ArnScheduler([capture]) scheduler.schedule(ARN_A, _record("terminal")) scheduler.drain(ARN_A) @@ -426,7 +457,7 @@ def fail_start(self) -> None: # noqa: ARG001 raise RuntimeError("cannot start") monkeypatch.setattr(threading.Thread, "start", fail_start) - scheduler = _ExportScheduler([CaptureExporter()]) + scheduler = _ArnScheduler([CaptureExporter()]) scheduler.schedule(ARN_A, _record("dropped")) scheduler.drain(ARN_A) @@ -436,7 +467,13 @@ def fail_start(self) -> None: # noqa: ARG001 with scheduler._condition: assert scheduler._disabled assert scheduler._pending == {} - assert scheduler._lanes == {} + # ...and no execution kept the record it was carrying. The per-execution + # bookkeeping that used to need clearing in a second map is on these + # objects now, so the queue and the records are one thing to release. + assert all( + execution.pending_record is None + for execution in scheduler.executions.values() + ) assert scheduler._flush_requested is False assert scheduler._flush_in_flight == 0 @@ -454,7 +491,7 @@ def test_disabled_latch_clears_a_published_flush_in_flight_marker(monkeypatch) - # asserts the same field, but reaches the latch with the marker already at 0, # so it holds whether or not the latch clears it; this one arms the marker # first. - scheduler = _ExportScheduler([CaptureExporter()]) + scheduler = _ArnScheduler([CaptureExporter()]) with scheduler._condition: scheduler._flush_in_flight = 7 @@ -479,7 +516,7 @@ def test_every_concurrent_execution_delivers_its_terminal_record_once() -> None: # different execution's record. executions = 10 exporter = EventLogExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) arns = [ARN.format(index) for index in range(executions)] ready = threading.Barrier(executions) @@ -511,7 +548,7 @@ def test_blocked_export_never_loses_another_executions_terminal_record() -> None # records. Neither may be dropped, and each drain must be released by its own # record reaching the exporters -- not by another execution's flush. exporter = GatedEventLogExporter() - scheduler = _ExportScheduler([exporter]) + scheduler = _ArnScheduler([exporter]) scheduler.schedule(ARN_A, _execution_record(ARN_A, "RUNNING")) assert exporter.first_export_started.wait(5.0) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 35a9d98d..baed5b7e 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -8,6 +8,13 @@ exercised end to end, nothing about SDK behavior is mocked). Operations reach the plugin the way the real SDK delivers them: as the point-in-time ``operations`` map on ``InvocationStartInfo`` / ``InvocationEndInfo`` / ``OperationChangeInfo``. + +``workflow_insight()`` returns a factory, so these tests hold two things where +they used to hold one: the handler-lifetime factory (``factory``, carrying the +resolved config, the exporters and the export scheduler) and the per-invocation +instance the SDK builds from it (``plugin``). ``_invocation()`` does what the SDK +does -- build the instance from the invocation's start info, then dispatch that +same info to its first hook. """ from __future__ import annotations @@ -41,7 +48,11 @@ WorkflowInsightConfig, workflow_insight, ) -from aws_durable_execution_sdk_python_insight.plugin import _resolve_sampling_rate +from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportState +from aws_durable_execution_sdk_python_insight.plugin import ( + WorkflowInsightPlugin, + _resolve_sampling_rate, +) ARN = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-1/inv-1" ARN_B = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-2/inv-1" @@ -138,8 +149,20 @@ def _end( ) +def _invocation(factory, info: InvocationStartInfo) -> WorkflowInsightPlugin: + """Enter one invocation the way the SDK does. + + The SDK builds one plugin instance per invocation from that invocation's + start info and dispatches the very same object to its first hook. A test that + drives hooks directly does both. + """ + plugin = factory(info) + plugin.on_invocation_start(info) + return plugin + + def _run( - plugin, + factory, *, ops, status=InvocationStatus.SUCCEEDED, @@ -150,10 +173,13 @@ def _run( """Single-invocation drive: the full operation map is present in both the start and the end snapshot (the terminal record is built from the end one).""" operations = _ops(*ops) - plugin.on_invocation_start(_start(operations=operations, input_value=input_value)) + plugin = _invocation( + factory, _start(operations=operations, input_value=input_value) + ) plugin.on_invocation_end( _end(operations=operations, status=status, result=result, error=error) ) + return plugin # -- existing record-building coverage --------------------------------------- @@ -161,8 +187,8 @@ def _run( def test_basic_success_record(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) - _run(plugin, ops=[_step("greet")]) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + _run(factory, ops=[_step("greet")]) assert len(exporter.records) == 1 rec = exporter.records[0] assert rec["recordType"] == "WorkflowInsight" @@ -185,17 +211,17 @@ def test_basic_success_record(): def test_on_failure_success_emits_nothing(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-failure") ) - _run(plugin, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) + _run(factory, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) assert exporter.records == [] # No record, but the invocation end still flushed once: a sampled-in # invocation end flushes whether or not this emit mode produced a record # (JS/Java cadence), because the exporter may be buffering another # execution's records. assert exporter.flush_count == 1 - assert _wait_until(lambda: not plugin._scheduler._worker_alive()) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) def test_invocation_end_that_emits_no_record_still_flushes_exactly_once(): @@ -205,8 +231,8 @@ def test_invocation_end_that_emits_no_record_still_flushes_exactly_once(): # rhythm in every SDK, so this end must still flush -- exactly once, not # twice, and not zero times. exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) - plugin.on_invocation_start(_start(operations={})) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end( _end(operations={}, status=InvocationStatus.PENDING, result=None) ) @@ -214,7 +240,7 @@ def test_invocation_end_that_emits_no_record_still_flushes_exactly_once(): assert exporter.flush_count == 1 # The worker retires, so no later flush can arrive after the invocation # returned. - assert _wait_until(lambda: not plugin._scheduler._worker_alive()) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) assert exporter.flush_count == 1 @@ -223,23 +249,23 @@ def test_sampled_out_invocation_end_neither_exports_nor_flushes(): # out execution exports nothing and must not flush either, so instrumenting # a fraction of executions costs the rest nothing. exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end(_end(operations=_ops(op))) assert exporter.records == [] assert exporter.flush_count == 0 - assert not plugin._scheduler._worker_alive() + assert not factory._scheduler._worker_alive() def test_sampling_zero_emits_nothing(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) assert exporter.records == [] @@ -251,22 +277,22 @@ def test_resolve_sampling_rate_nan_fails_open_to_one(): def test_nan_sampling_rate_emits_instead_of_silently_disabling(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=float("nan")) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) # A NaN rate must not disable instrumentation: the record is still emitted. assert len(exporter.records) == 1 def test_content_omit_input_output_without_drop_flags(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig(input=False, output=False) ) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) rec = exporter.records[0] assert "input" not in rec and "output" not in rec assert "droppedInput" not in rec and "droppedOutput" not in rec @@ -274,7 +300,7 @@ def test_content_omit_input_output_without_drop_flags(): def test_result_opt_in(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig( @@ -284,14 +310,14 @@ def test_result_opt_in(): ), ) ) - _run(plugin, ops=[_step("compute", result="42")], result="42") + _run(factory, ops=[_step("compute", result="42")], result="42") op = exporter.records[0]["operations"][0] assert op["result"] == 42 # checkpointed JSON string parsed def test_include_errors_false_drops_op_error_keeps_record_error(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig(operations=ContentOperations(include_errors=False)), @@ -302,7 +328,7 @@ def test_include_errors_false_drops_op_error_keeps_record_error(): message="boom", type="InsightTestError", data=None, stack_trace=None ) _run( - plugin, + factory, ops=[_step("failing-step", status=OperationStatus.FAILED, error=op_err)], status=InvocationStatus.FAILED, result=None, @@ -315,7 +341,7 @@ def test_include_errors_false_drops_op_error_keeps_record_error(): def test_top_level_only_drops_children(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) parent = _step( "parallel-work", op_id="p", @@ -323,14 +349,14 @@ def test_top_level_only_drops_children(): sub_type=OperationSubType.PARALLEL, ) child = _step("branch-a-step", parent_id="p", op_id="c") - _run(plugin, ops=[parent, child]) + _run(factory, ops=[parent, child]) names = [op["name"] for op in exporter.records[0]["operations"]] assert names == ["parallel-work"] def test_full_tree_includes_children_with_parent_id(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], operation_detail="full-tree") ) parent = _step( @@ -340,7 +366,7 @@ def test_full_tree_includes_children_with_parent_id(): sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, ) child = _step("child-step", parent_id="p", op_id="c") - _run(plugin, ops=[parent, child]) + _run(factory, ops=[parent, child]) ops = {op["name"]: op for op in exporter.records[0]["operations"]} assert set(ops) == {"parent-context", "child-step"} assert ops["child-step"]["parentId"] == "p" @@ -348,9 +374,9 @@ def test_full_tree_includes_children_with_parent_id(): def test_unnamed_operation_dropped(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) unnamed = _step(None, op_id="u") # type: ignore[arg-type] - _run(plugin, ops=[_step("named-step"), unnamed]) + _run(factory, ops=[_step("named-step"), unnamed]) names = [op["name"] for op in exporter.records[0]["operations"]] assert names == ["named-step"] @@ -359,9 +385,10 @@ def test_unnamed_operation_dropped(): def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): - # Invocation 1 (plugin A): a step completes, then a wait suspends -> PENDING. + # Invocation 1 (environment A): a step completes, then a wait suspends -> + # PENDING. exporter1 = CaptureExporter() - plugin1 = workflow_insight(WorkflowInsightConfig(exporters=[exporter1])) + factory1 = workflow_insight(WorkflowInsightConfig(exporters=[exporter1])) step = _step("greet", op_id="op-step") wait_pending = _step( "pause", @@ -371,7 +398,7 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): status=OperationStatus.PENDING, end_time=None, ) - plugin1.on_invocation_start(_start(operations={})) + plugin1 = _invocation(factory1, _start(operations={})) plugin1.on_invocation_end( _end( operations=_ops(step, wait_pending), @@ -380,12 +407,15 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): ) ) assert exporter1.records == [] # on-complete emits nothing for a suspend - assert plugin1._state == {} # and retains nothing + # And nothing is retained: the instance that served the suspending invocation + # is dropped by the SDK, and the scheduler holds no execution either. + assert _wait_until(lambda: _scheduler_is_empty(factory1)) - # Invocation 2 on a *fresh* plugin instance (new Lambda environment): the - # resume start snapshot carries the prior terminal step + resolved wait. + # Invocation 2 in a *fresh* Lambda environment -- a new factory, and so also a + # new instance: the resume start snapshot carries the prior terminal step + + # resolved wait. exporter2 = CaptureExporter() - plugin2 = workflow_insight(WorkflowInsightConfig(exporters=[exporter2])) + factory2 = workflow_insight(WorkflowInsightConfig(exporters=[exporter2])) step_done = _step("greet", op_id="op-step") wait_done = _step( "pause", @@ -395,8 +425,8 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): status=OperationStatus.SUCCEEDED, ) resume_ops = _ops(step_done, wait_done) - plugin2.on_invocation_start( - _start(operations=resume_ops, is_first=False, execution_start_time=T0) + plugin2 = _invocation( + factory2, _start(operations=resume_ops, is_first=False, execution_start_time=T0) ) plugin2.on_invocation_end( _end(operations=resume_ops, is_first=False, execution_start_time=T0) @@ -416,13 +446,13 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): def test_on_change_schedules_running_and_delivers_terminal(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op1 = _step("s1", op_id="1") op2 = _step("s2", op_id="2") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1) @@ -451,15 +481,17 @@ def test_on_change_schedules_running_and_delivers_terminal(): def test_concurrent_executions_do_not_cross_contaminate(): # A and B both suspend; B is the most-recently started (the old insertion- # order heuristic would have attributed A's resume to B). A then resumes to - # a terminal state. Its record must contain only A's data. + # a terminal state. Its record must contain only A's data. Each invocation + # gets its own instance from the one shared factory, which is what the SDK + # does for concurrent executions in one environment. exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) a_op = _step("a-step", op_id="a1") b_op = _step("b-step", op_id="b1") - plugin.on_invocation_start(_start(arn=ARN, operations={}, input_value="A")) - plugin.on_invocation_start(_start(arn=ARN_B, operations={}, input_value="B")) - plugin.on_invocation_end( + a_first = _invocation(factory, _start(arn=ARN, operations={}, input_value="A")) + b_first = _invocation(factory, _start(arn=ARN_B, operations={}, input_value="B")) + b_first.on_invocation_end( _end( arn=ARN_B, operations=_ops(b_op), @@ -467,7 +499,7 @@ def test_concurrent_executions_do_not_cross_contaminate(): result=None, ) ) - plugin.on_invocation_end( + a_first.on_invocation_end( _end( arn=ARN, operations=_ops(a_op), @@ -478,10 +510,11 @@ def test_concurrent_executions_do_not_cross_contaminate(): assert exporter.records == [] # both suspended, nothing terminal yet a_done = _step("a-step", op_id="a1") - plugin.on_invocation_start( - _start(arn=ARN, operations=_ops(a_done), is_first=False, input_value="A") + a_resume = _invocation( + factory, + _start(arn=ARN, operations=_ops(a_done), is_first=False, input_value="A"), ) - plugin.on_invocation_end( + a_resume.on_invocation_end( _end(arn=ARN, operations=_ops(a_done), is_first=False, result='"A-done"') ) @@ -492,22 +525,25 @@ def test_concurrent_executions_do_not_cross_contaminate(): assert [op["name"] for op in rec["operations"]] == ["a-step"] -# -- state lifecycle: clear after every invocation end (comment 4) ----------- +# -- nothing retained after an invocation end (comment 4) -------------------- -def test_state_cleared_after_pending_and_retry(): +def test_nothing_retained_after_pending_and_retry(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) - plugin.on_invocation_end( + suspend = _invocation(factory, _start(operations={})) + suspend.on_invocation_end( _end(operations=_ops(op), status=InvocationStatus.PENDING, result=None) ) - assert plugin._state == {} # no leak after suspend + # The instance that served the suspending invocation is dropped by the SDK, + # so the only thing that could retain anything for this execution is the + # scheduler, and it holds nothing either. + assert _wait_until(lambda: _scheduler_is_empty(factory)) - plugin.on_invocation_start(_start(operations=_ops(op), is_first=False)) - plugin.on_invocation_end( + retry = _invocation(factory, _start(operations=_ops(op), is_first=False)) + retry.on_invocation_end( _end( operations=_ops(op), status=InvocationStatus.RETRY, @@ -515,17 +551,17 @@ def test_state_cleared_after_pending_and_retry(): is_first=False, ) ) - assert plugin._state == {} # no leak after retry + assert _wait_until(lambda: _scheduler_is_empty(factory)) assert exporter.records == [] # on-complete emits nothing for non-terminal def test_sampled_out_processes_nothing_and_retains_no_state(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) @@ -533,33 +569,35 @@ def test_sampled_out_processes_nothing_and_retains_no_state(): ) plugin.on_invocation_end(_end(operations=_ops(op))) assert exporter.records == [] - assert plugin._state == {} + # A sampled-out invocation adopts no operations and schedules nothing. + assert plugin._operations == {} + assert _scheduler_is_empty(factory) # -- default exporter parity with JS (comment 6) ----------------------------- def test_default_exporter_when_config_omits_exporters(): - plugin = workflow_insight(WorkflowInsightConfig()) - assert len(plugin._exporters) == 1 - assert isinstance(plugin._exporters[0], LambdaLogExporter) + factory = workflow_insight(WorkflowInsightConfig()) + assert len(factory._exporters) == 1 + assert isinstance(factory._exporters[0], LambdaLogExporter) def test_default_exporter_when_exporters_explicitly_empty(): - plugin = workflow_insight(WorkflowInsightConfig(exporters=[])) - assert len(plugin._exporters) == 1 - assert isinstance(plugin._exporters[0], LambdaLogExporter) + factory = workflow_insight(WorkflowInsightConfig(exporters=[])) + assert len(factory._exporters) == 1 + assert isinstance(factory._exporters[0], LambdaLogExporter) def test_explicit_exporters_are_preserved(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) - assert plugin._exporters == [exporter] + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + assert factory._exporters == [exporter] def test_default_exporter_actually_emits_to_stdout(capsys): - plugin = workflow_insight(WorkflowInsightConfig()) - _run(plugin, ops=[_step("greet")]) + factory = workflow_insight(WorkflowInsightConfig()) + _run(factory, ops=[_step("greet")]) out = capsys.readouterr().out assert '"recordType":"WorkflowInsight"' in out # compact JSON via LambdaLogExporter assert '"operationsByName"' in out @@ -576,11 +614,11 @@ def _last_record_for_end_status(status, *, emit_mode="on-change", result=None): emit nothing for a non-terminal end. """ exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode=emit_mode) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end(_end(operations=_ops(op), status=status, result=result)) return exporter.records[-1] @@ -612,11 +650,11 @@ def test_succeeded_end_is_terminal_with_end_time_and_duration(): def test_failed_end_is_terminal_with_end_time_and_duration(): err = ErrorObject(message="boom", type="StepError", data=None, stack_trace=None) exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op = _step("s", op_id="1", status=OperationStatus.FAILED) - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end( _end( operations=_ops(op), @@ -660,12 +698,13 @@ def snapshot(self) -> list[dict[str, Any]]: def test_concurrent_executions_each_deliver_their_terminal_record(): - # One plugin instance serves every execution its environment hosts, and LMI - # runs several at once. Drive the real hooks concurrently: every execution's - # terminal record must arrive exactly once. + # One factory (one scheduler) serves every execution its environment hosts, + # and LMI runs several at once. Drive the real hooks concurrently, each + # execution on its own instance: every execution's terminal record must + # arrive exactly once. executions = 5 exporter = ConcurrentCaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) arns = [ @@ -677,7 +716,7 @@ def test_concurrent_executions_each_deliver_their_terminal_record(): def run(arn: str) -> None: op = _step("s", op_id="1") ready.wait(10.0) - plugin.on_invocation_start(_start(arn=arn, operations={})) + plugin = _invocation(factory, _start(arn=arn, operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=arn, updated_operations=_ops(op), operations=_ops(op) @@ -698,9 +737,9 @@ def run(arn: str) -> None: if record["status"] == "SUCCEEDED" ] assert sorted(terminal) == sorted(arns) # each exactly once, none lost - assert plugin._state == {} - # Nothing per-execution is retained in the scheduler either. - assert _wait_until(lambda: _scheduler_is_empty(plugin)) + # Nothing per-execution is retained: each instance is dropped by the SDK with + # its invocation, and the scheduler holds no execution either. + assert _wait_until(lambda: _scheduler_is_empty(factory)) def _wait_until(predicate, timeout: float = 5.0) -> bool: @@ -712,10 +751,25 @@ def _wait_until(predicate, timeout: float = 5.0) -> bool: return predicate() -def _scheduler_is_empty(plugin) -> bool: - scheduler = plugin._scheduler +def _scheduler_is_empty(factory) -> bool: + scheduler = factory._scheduler with scheduler._condition: - return not scheduler._pending and not scheduler._lanes + # One structure: the queue of executions with a record waiting. An + # execution's export bookkeeping lives on the per-invocation instance + # itself, which the SDK drops when the invocation scope exits, so an + # empty queue means the scheduler retains nothing. + return not scheduler._pending + + +def _force_drain(factory) -> None: + """Push everything this factory's instances scheduled out to the exporters. + + drain() takes the per-execution object, which in production is the plugin + instance itself. These tests present a bare ``_ExportState`` instead, exactly + as an execution that scheduled nothing of its own would: the flush it requests + is held back until every record pending when it was called has been exported. + """ + factory._scheduler.drain(_ExportState()) class PinnedWorkerExporter: @@ -764,7 +818,7 @@ def __del__(self) -> None: ) ) - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], emit_mode="on-change", @@ -772,6 +826,8 @@ def __del__(self) -> None: ) ) op = _step("s", op_id="1") + start = _start(operations={}) + plugin = factory(start) holder["plugin"] = plugin holder["ops"] = _ops(op) change = OperationChangeInfo( @@ -779,7 +835,7 @@ def __del__(self) -> None: ) try: # The first emit pins the single worker inside export()... - plugin.on_invocation_start(_start(operations={})) + plugin.on_invocation_start(start) assert exporter.entered.wait(5.0) # ...so this emit stays this execution's pending record... plugin.on_operation_change(change) @@ -822,7 +878,7 @@ def blocking_output(value: Any) -> Any: release_terminal.wait(10.0) return value - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], emit_mode="on-change", @@ -830,7 +886,7 @@ def blocking_output(value: Any) -> Any: ) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) end_returned = threading.Event() @@ -864,13 +920,12 @@ def change() -> None: change_thread.join(5.0) assert end_returned.is_set() assert change_returned.is_set() - plugin._scheduler.drain(ARN) + _force_drain(factory) statuses = [record["status"] for record in exporter.snapshot()] assert "SUCCEEDED" in statuses # Nothing at all after the terminal record. assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] - assert plugin._state == {} def test_invocation_end_waits_for_an_in_flight_change_hook_emit(): @@ -894,7 +949,7 @@ def blocking_input(value: Any) -> Any: release_change.wait(10.0) return value - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], emit_mode="on-change", @@ -902,7 +957,7 @@ def blocking_input(value: Any) -> Any: ) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) change_returned = threading.Event() @@ -934,24 +989,25 @@ def end() -> None: end_thread.join(10.0) assert change_returned.is_set() assert end_returned.is_set() - plugin._scheduler.drain(ARN) + _force_drain(factory) statuses = [record["status"] for record in exporter.snapshot()] assert "SUCCEEDED" in statuses assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] - assert plugin._state == {} def test_operation_change_after_invocation_end_emits_nothing(): # A checkpoint that completed just before the invocation ended still delivers - # its operation-change hook. It must not recreate state, must not fabricate a - # start time, and must not append a RUNNING record after the terminal one. + # its operation-change hook. It reaches the instance for the invocation that + # just ended -- there is no registry it could recreate an entry in -- and must + # find the gate closed: no fabricated start time, and no RUNNING record after + # the terminal one. exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end(_end(operations=_ops(op))) before = list(exporter.records) assert before and before[-1]["status"] == "SUCCEEDED" @@ -961,7 +1017,7 @@ def test_operation_change_after_invocation_end_emits_nothing(): execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) ) ) - plugin._scheduler.drain(ARN) + _force_drain(factory) assert exporter.records == before # nothing emitted after the terminal record assert [record["status"] for record in exporter.records][-1] == "SUCCEEDED" @@ -969,70 +1025,56 @@ def test_operation_change_after_invocation_end_emits_nothing(): assert {record["startTime"] for record in exporter.records} == { "2026-01-01T00:00:00Z" } - assert plugin._state == {} # and no state entry recreated - - -def test_late_invocation_start_finds_the_closed_gate_shut(monkeypatch): - # on_invocation_end sets `closed` and emits the terminal record while holding - # the execution's lock, RELEASES the lock, and only then discards the state. - # A concurrent on_invocation_start that already resolved that state reference - # gets the lock inside that window and finds a state that is closed but not - # yet gone; the gate has to shut it out. - # - # In production that window is sub-microsecond, so it is entered here - # deterministically: the late hook runs from inside _discard_state, which is - # exactly where the window sits. + + +def test_invocation_start_after_the_gate_closed_changes_nothing(): + # Each invocation gets its own instance and exactly one invocation-start hook, + # so a start hook arriving on a closed instance is no longer reachable through + # a state registry -- there is none, and no instance can be handed to a second + # invocation. The `closed` gate is what holds that contract from the plugin's + # side: a start hook that arrives after the invocation ended must not re-seed + # the closed instance and must emit nothing. exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={}, input_value="World")) - state = plugin._state[ARN] - real_discard = plugin._discard_state - late = threading.Event() - - def discard_after_a_late_start(arn: str) -> None: - if not late.is_set(): - late.set() - assert state.closed # the window: closed, emitted, lock free, state alive - plugin.on_invocation_start( - _start( - operations=_ops(_step("late", op_id="2")), - input_value="late-input", - execution_start_time=T1, - ) - ) - real_discard(arn) - - monkeypatch.setattr(plugin, "_discard_state", discard_after_a_late_start) + plugin = _invocation(factory, _start(operations={}, input_value="World")) plugin.on_invocation_end(_end(operations=_ops(op))) - plugin._scheduler.drain(ARN) + assert plugin._closed + + plugin.on_invocation_start( + _start( + operations=_ops(_step("late", op_id="2")), + input_value="late-input", + execution_start_time=T1, + ) + ) + _force_drain(factory) - assert late.is_set() # the late hook really did run inside the window statuses = [record["status"] for record in exporter.records] assert "SUCCEEDED" in statuses # Nothing follows the terminal record... assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] - # ...and the closed state was not re-seeded on the way out. A late start that - # got past the gate adopts its own operation snapshot, input and start time, - # which is the observable effect of the gate: the emission it would also have - # produced is stopped a second time by the re-check in _emit, so these are - # what pin the hook's own check. - assert state.cached_input == "World" - assert state.start_time == T0 - assert [info.name for info in state.operations.values()] == ["s"] - assert plugin._state == {} + # ...and the closed instance was not re-seeded. A late start that got past the + # gate would adopt its own operation snapshot; the emission it would also have + # produced is stopped a second time by the re-check in _emit, so the adopted + # state is what pins the hook's own check. Input and start time are fixed by + # the constructor from the invocation's own start info, so no later hook can + # move them at all. + assert plugin._cached_input == "World" + assert plugin._start_time == T0 + assert [info.name for info in plugin._operations.values()] == ["s"] def test_reentrant_invocation_end_stops_the_outer_running_record(): # The gate at the top of each hook is a check-then-act, and _emit is the act. # Between them _emit runs customer code while holding the execution's lock -- # here a content transform -- and the lock is reentrant, so that customer code - # can run on_invocation_end to completion on this same thread: `closed` set, - # terminal record scheduled, state discarded and drained. The outer frame then - # resumes with a fully built RUNNING record, which must NOT reach the - # exporters after the terminal one. One hook call, no concurrency. + # can run on_invocation_end to completion on this same thread: `_closed` set, + # terminal record scheduled and drained. The outer frame then resumes with a + # fully built RUNNING record, which must NOT reach the exporters after the + # terminal one. One hook call, one instance, no concurrency. exporter = ConcurrentCaptureExporter() holder: dict[str, Any] = {} reentered = threading.Event() @@ -1043,13 +1085,15 @@ def reentering_input(value: Any) -> Any: holder["plugin"].on_invocation_end(_end(operations=_ops(_step("s")))) return value - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], emit_mode="on-change", content=ContentConfig(input=reentering_input), ) ) + start = _start(operations={}) + plugin = factory(start) holder["plugin"] = plugin # On a bounded thread, so a regression that makes the lock non-reentrant @@ -1057,7 +1101,7 @@ def reentering_input(value: Any) -> Any: returned = threading.Event() def hook() -> None: - plugin.on_invocation_start(_start(operations={})) + plugin.on_invocation_start(start) returned.set() thread = threading.Thread(target=hook, daemon=True) @@ -1071,11 +1115,11 @@ def hook() -> None: assert reentered.is_set() # the re-entrant end hook really did run # Force everything the plugin scheduled to reach the exporters, so a record # that slipped past the gate is observed here rather than left pending. - plugin._scheduler.drain(ARN) + _force_drain(factory) statuses = [record["status"] for record in exporter.snapshot()] assert statuses == ["SUCCEEDED"], ( "a non-terminal record reached the exporters after the terminal one for " f"the same execution: {statuses}" ) - assert _wait_until(lambda: not plugin._scheduler._worker_alive()) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 5e78f2a9..f532ed35 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -39,9 +39,15 @@ processors, and exporter. 1. Add the [ADOT Lambda Layer](#1-adot-lambda-layer) to your function and set `AWS_LAMBDA_EXEC_WRAPPER=/opt/otel-instrument` 2. Enable [X-Ray Active Tracing](#2-aws-x-ray-active-tracing) on the function -3. Pass `InvocationOtelPlugin` to your handler's `plugins` list +3. Pass `InvocationOtelPluginFactory()` to your handler's `plugins` list 4. Add X-Ray write permissions +The SDK's `plugins` list takes plugin *factories*, not plugin instances: it calls +each factory once per invocation and the plugin it returns serves that one +invocation. `InvocationOtelPluginFactory` and `ExecutionOtelPluginFactory` are +the factories for the two bundled plugins; each takes the optional +`OtelPluginConfig` that every plugin it builds will use. + Alternatively, install this package in the function artifact or a Lambda layer and select either OTel plugin by entry-point name: @@ -50,9 +56,10 @@ DURABLE_EXECUTION_PLUGINS=otel-invocation DURABLE_EXECUTION_PLUGINS=otel-execution ``` -`otel-invocation` creates `InvocationOtelPlugin`; `otel-execution` creates -`ExecutionOtelPlugin`. The SDK discovers the selected package entry point at -cold start, so the handler does not need to import or explicitly register the +`otel-invocation` names a default-configured `InvocationOtelPluginFactory`; +`otel-execution` names a default-configured `ExecutionOtelPluginFactory`. The SDK +discovers the selected package entry point at cold start and calls it once per +invocation, so the handler does not need to import or explicitly register the plugin. ### 1. ADOT Lambda Layer @@ -161,10 +168,10 @@ lambda_.Function( ```python from aws_durable_execution_sdk_python import DurableContext from aws_durable_execution_sdk_python.execution import durable_execution -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(event: dict, context: DurableContext) -> dict: result = context.step(lambda _: fetch_data(event["id"]), name="fetch-data") @@ -199,12 +206,12 @@ See the [ADOT sampling configuration](https://aws-otel.github.io/docs/getting-st ```python from aws_durable_execution_sdk_python_otel import ( - InvocationOtelPlugin, + InvocationOtelPluginFactory, OtelPluginConfig, xray_context_extractor, ) -plugin = InvocationOtelPlugin( +plugin_factory = InvocationOtelPluginFactory( OtelPluginConfig( # Use a custom context extractor (default: xray_context_extractor). context_extractor=xray_context_extractor, @@ -218,6 +225,9 @@ plugin = InvocationOtelPlugin( ) ``` +The config is resolved once and shared by every plugin the factory builds, so +configuration is per handler while plugin state is per invocation. + ### Context Extractors Context extractors return an `ExtractedContext` object, or `None` when no @@ -232,17 +242,19 @@ context: ```python from aws_durable_execution_sdk_python_otel import ( - InvocationOtelPlugin, + InvocationOtelPluginFactory, OtelPluginConfig, w3c_client_context_extractor, xray_context_extractor, ) # Default: X-Ray trace header (recommended for most Lambda deployments). -InvocationOtelPlugin(OtelPluginConfig(context_extractor=xray_context_extractor)) +InvocationOtelPluginFactory(OtelPluginConfig(context_extractor=xray_context_extractor)) # W3C Trace Context via clientContext (placeholder for backend propagation support). -InvocationOtelPlugin(OtelPluginConfig(context_extractor=w3c_client_context_extractor)) +InvocationOtelPluginFactory( + OtelPluginConfig(context_extractor=w3c_client_context_extractor) +) ``` Custom extractors should return `ExtractedContext`, not an OpenTelemetry @@ -339,12 +351,15 @@ After deploying your function with the plugin configured: ## API Reference -### `InvocationOtelPlugin` +### `InvocationOtelPluginFactory` -Invocation-rooted view. Implements `DurableInstrumentationPlugin` from `aws_durable_execution_sdk_python`. +Factory for the invocation-rooted plugin, and what belongs in the SDK's `plugins` +list. Satisfies `DurableInstrumentationPluginFactory` from +`aws_durable_execution_sdk_python`: calling it with an `InvocationStartInfo` +returns the `InvocationOtelPlugin` for that invocation. ```python -InvocationOtelPlugin( +InvocationOtelPluginFactory( OtelPluginConfig( tracer_provider=None, context_extractor=None, @@ -356,13 +371,36 @@ InvocationOtelPlugin( ``` Pass `tracer_provider=...` when the application owns the OpenTelemetry SDK -provider. When omitted, the globally configured provider is used. +provider. When omitted, the globally configured provider is used, resolved per +invocation so a provider installed after the handler module is imported is still +picked up. + +`INVOCATION_OTEL_PLUGIN_FACTORY` is the default-configured instance the +`otel-invocation` entry point names. + +### `ExecutionOtelPluginFactory` + +Factory for the execution-rooted plugin, with the same construction and config as +`InvocationOtelPluginFactory`. `EXECUTION_OTEL_PLUGIN_FACTORY` is the +default-configured instance the `otel-execution` entry point names. + +### `InvocationOtelPlugin` + +Invocation-rooted view. Implements `DurableInstrumentationPlugin` from +`aws_durable_execution_sdk_python`. One instance serves exactly one invocation: +the SDK builds it from the factory before the first hook fires and drops it when +the invocation ends, so it holds its span registry and context tokens in ordinary +instance state. Construct it directly only when driving the hooks yourself; in a +handler, register the factory instead. ### `ExecutionOtelPlugin` Execution-rooted view. Uses the same execution ancestor and sampling behavior as `InvocationOtelPlugin`, but parents operation spans under Workflow and links -them to Invocation. +them to Invocation. Also one instance per invocation; the identities that must +agree across invocations (the trace ID, the Workflow span ID, and each +operation's span ID) are derived deterministically from the execution ARN rather +than carried in memory. ### `DeterministicIdGenerator` diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index 36e04b31..02c695c4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -26,8 +26,8 @@ dependencies = [ ] [project.entry-points."aws_durable_execution.plugins"] -otel-invocation = "aws_durable_execution_sdk_python_otel.plugin_provider:INVOCATION_OTEL_PLUGIN_PROVIDER" -otel-execution = "aws_durable_execution_sdk_python_otel.plugin_provider:EXECUTION_OTEL_PLUGIN_PROVIDER" +otel-invocation = "aws_durable_execution_sdk_python_otel.plugin_factory:INVOCATION_OTEL_PLUGIN_FACTORY" +otel-execution = "aws_durable_execution_sdk_python_otel.plugin_factory:EXECUTION_OTEL_PLUGIN_FACTORY" [project.optional-dependencies] # Lambda telemetry layers provide a version-aligned OpenTelemetry distribution. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py index a8285dcb..774ae771 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py @@ -27,6 +27,12 @@ from aws_durable_execution_sdk_python_otel.invocation_plugin import ( InvocationOtelPlugin, ) +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + EXECUTION_OTEL_PLUGIN_FACTORY, + INVOCATION_OTEL_PLUGIN_FACTORY, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) from aws_durable_execution_sdk_python_otel.provider import ( ProviderResult, create_tracer_provider, @@ -35,12 +41,16 @@ __all__ = [ "__version__", + "EXECUTION_OTEL_PLUGIN_FACTORY", + "INVOCATION_OTEL_PLUGIN_FACTORY", "ContextExtractor", "DeterministicIdGenerator", "ExecutionOtelPlugin", + "ExecutionOtelPluginFactory", "ExtractedContext", "OtelPluginConfig", "InvocationOtelPlugin", + "InvocationOtelPluginFactory", "OtelContextLogFilter", "Sampling", "ProviderResult", 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 1a46cafc..c522278f 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 @@ -117,6 +117,15 @@ def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: class ExecutionOtelPlugin(DurableInstrumentationPlugin): """OTel plugin that renders a durable execution as one Workflow-rooted trace. + Lifetime: one instance per invocation. The SDK builds it from + :class:`~aws_durable_execution_sdk_python_otel.plugin_factory.ExecutionOtelPluginFactory` + before the first hook fires and drops it when the invocation scope exits, so + every field below is per-invocation state that no other invocation can + observe. The execution-scoped identities the plugin needs across invocations + -- the canonical trace ID, the Workflow span ID and each operation's span ID + -- are derived deterministically from the execution ARN, so a fresh instance + rejoins the same trace without carrying anything over. + Args: config: Shared plugin configuration. When omitted, defaults are used (globally configured provider, X-Ray extractor, "Workflow" root @@ -139,7 +148,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() - # Per-invocation state. + # Per-invocation state. The SDK builds one plugin instance per + # invocation through ExecutionOtelPluginFactory and drops it when the + # invocation scope exits, so these are ordinary instance fields that + # never have to be cleared for reuse. self._execution_arn = "" self._execution_trace_id: int | None = None self._execution_start_time: datetime.datetime | None = None @@ -165,6 +177,8 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._tracing_enabled = False if self._config.enrich_logger: + # Install (or, on a warm environment, rebind) the root-logger filter + # so every log record is stamped with this invocation's span context. install_log_filter(self) def _bind_sdk_tracer(self) -> bool: @@ -267,11 +281,10 @@ def _detach_context(self, key: str) -> None: otel_context.detach(token) # type: ignore[arg-type] def _detach_remaining_contexts(self) -> None: - """Release scopes still open, newest first, so nothing outlives the plugin. + """Release scopes still open, newest first, so nothing outlives the invocation. Reached when a lifecycle end hook never fires -- for example a user - function that suspends, or a warm invocation that starts before the - previous one was cleaned up. + function that suspends. """ with self._lock: keys = list(reversed(self._context_tokens)) @@ -413,7 +426,6 @@ def _with_sampling(self, parent_context: Context) -> Context: # ------------------------------------------------------------------ def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) - self._reset_state() if info.execution_start_time is None: logger.warning( "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -477,8 +489,8 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: # Make the Workflow span the active span so auto-instrumented spans # created during the invocation become its children. The token is - # released in _reset_state at invocation end, restoring the context that - # was active before the invocation started. + # released in _release_invocation_scope at invocation end, restoring the + # context that was active before the invocation started. if self._workflow_span is not None: self._attach_context( _INVOCATION_CONTEXT_KEY, @@ -592,7 +604,7 @@ def _start_invocation_span(self, info: InvocationStartInfo) -> None: def on_invocation_end(self, info: InvocationEndInfo) -> None: logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: - self._reset_state() + self._release_invocation_scope() return # End the invocation span regardless of terminal status. Record the @@ -628,7 +640,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._reset_state() + self._release_invocation_scope() if hasattr(self._provider, "force_flush"): try: @@ -636,20 +648,24 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") - def _reset_state(self) -> None: + def _release_invocation_scope(self) -> None: + """Release what this invocation attached, and stop instrumenting. + + Not a state reset. The instance serves exactly one invocation and is + dropped afterwards, so its fields never have to be cleared for a warm + environment's next invocation. Two things still have to happen at the + invocation boundary: + + * The OpenTelemetry context stack belongs to the thread, not to the + plugin, so any scope this plugin attached and did not release must be + detached here or it would stay current on a warm environment's thread + after the invocation returns. + * ``_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. + """ self._detach_remaining_contexts() - self._execution_arn = "" - self._execution_trace_id = None - self._extracted_context = None - self._execution_trace_context = None - self._sampling_intent = None - self._execution_start_time = None - self._workflow_span = None - self._invocation_span = None - with self._lock: - self._operation_spans = {} - self._checkpointed_context_ids = set() - self._ended_operation_ids = set() self._tracing_enabled = False # ------------------------------------------------------------------ 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 012f17f2..96ecda36 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 @@ -96,6 +96,12 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): use newly generated span IDs. Operation attributes and links to the Workflow span provide execution-scoped correlation across invocations. + Lifetime: one instance per invocation. The SDK builds it from + :class:`~aws_durable_execution_sdk_python_otel.plugin_factory.InvocationOtelPluginFactory` + before the first hook fires and drops it when the invocation scope exits, so + every field below is per-invocation state that no other invocation can + observe. + Args: config: Shared plugin configuration (the same OtelPluginConfig accepted by ExecutionOtelPlugin). When omitted, defaults are used (X-Ray @@ -120,7 +126,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: When ``enrich_logger`` is enabled (default), the plugin installs a logging filter that stamps the active OTel trace context onto every - emitted log record. + emitted log record. The filter is installed on the root logger's + handlers, which outlive this instance, so installing rebinds an + already-present filter to this invocation's plugin rather than stacking a + second one. """ self._config = config or OtelPluginConfig() self._context_extractor: ContextExtractor = ( @@ -137,7 +146,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() - # per invocation status: + # Per-invocation state. The SDK builds one plugin instance per + # invocation through InvocationOtelPluginFactory and drops it when the + # invocation scope exits, so these are ordinary instance fields that + # never have to be cleared for reuse. self._execution_arn = "" self._execution_trace_id: int | None = None self._execution_start_time: datetime.datetime | None = None @@ -166,7 +178,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # Install the root-logger filter so every log record is stamped with # the active span context. The Lambda runtime attaches its root # handler before the handler module is imported (and thus before the - # plugin is constructed), so the handlers are available here. + # plugin is constructed), so the handlers are available here. On a + # warm environment the handler already carries the filter a previous + # invocation's plugin installed, and this rebinds it to this one. install_log_filter(self) def _bind_sdk_tracer(self) -> bool: @@ -266,11 +280,10 @@ def _detach_context(self, key: str) -> None: context.detach(token) # type: ignore[arg-type] def _detach_remaining_contexts(self) -> None: - """Release scopes still open, newest first, so nothing outlives the plugin. + """Release scopes still open, newest first, so nothing outlives the invocation. Reached when a lifecycle end hook never fires -- for example a user - function that suspends, or a warm invocation that starts before the - previous one was cleaned up. + function that suspends. """ with self._operation_spans_lock: keys = list(reversed(self._context_tokens)) @@ -513,7 +526,6 @@ def _end_span( def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) - self._reset_state() if info.execution_start_time is None: logger.warning( "InvocationOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -652,7 +664,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: """Called at the end of each invocation. Ends the invocation span and flushes.""" logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: - self._reset_state() + self._release_invocation_scope() return # Spans are registered parent-first, so close pending spans in reverse @@ -700,27 +712,30 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._reset_state() + self._release_invocation_scope() # Flush before Lambda freeze if hasattr(self._provider, "force_flush"): self._provider.force_flush() - def _reset_state(self) -> None: - """Clear per-invocation state for warm Lambda environment reuse.""" + def _release_invocation_scope(self) -> None: + """Release what this invocation attached, and stop instrumenting. + + Not a state reset. The instance serves exactly one invocation and is + dropped afterwards, so its fields never have to be cleared for a warm + environment's next invocation. Two things still have to happen at the + invocation boundary: + + * The OpenTelemetry context stack belongs to the thread, not to the + plugin, so any scope this plugin attached and did not release must be + detached here or it would stay current on a warm environment's thread + after the invocation returns. + * ``_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. + """ self._detach_remaining_contexts() - self._execution_arn = "" - self._execution_trace_id = None - self._extracted_context = None - self._execution_trace_context = None - self._sampling_intent = None - self._execution_start_time = None - self._workflow_span = None - self._span_time_floor_ns = None - with self._operation_spans_lock: - self._operation_spans = {} - self._context_operation_replays = {} - self._incomplete_attempt_span_keys = set() self._tracing_enabled = False def on_operation_start(self, info: OperationStartInfo) -> None: 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 7e11ebc9..e1e8d34e 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 @@ -51,6 +51,9 @@ class OtelContextLogFilter(logging.Filter): The filter never caches identifiers and always returns ``True`` so it never drops a record. + The plugin it reads is rebindable, because a plugin instance serves exactly + one invocation while the handler it is attached to outlives the invocation. + Args: plugin: The OTel plugin instance that resolves the current span context. """ @@ -59,6 +62,19 @@ def __init__(self, plugin: _SpanContextProvider) -> None: super().__init__() self._plugin = plugin + def bind(self, plugin: _SpanContextProvider) -> None: + """Point the filter at the plugin serving the current invocation. + + A logging handler lives as long as the Lambda environment, but a plugin + instance lives for one invocation. Without rebinding, the filter + installed by the first invocation's plugin would keep asking that + already-discarded instance for a span context -- it reports none once its + invocation ended, so log correlation would stop after the first + invocation, and the dead instance would be kept reachable for the life of + the environment. + """ + self._plugin = plugin + def filter(self, record: logging.LogRecord) -> bool: """Stamp the active span context onto the record, then allow it through.""" span_context = self._plugin.get_current_span_context() @@ -83,9 +99,10 @@ def install_log_filter( filters run for every record reaching the handler. This is safe to call on every invocation: if a handler already has an - OtelContextLogFilter, it is left as-is, so warm Lambda reuse will not stack - duplicate filters. A single shared filter instance is reused across all - handlers. + OtelContextLogFilter, that filter is rebound to ``plugin`` and left in place, + so a warm Lambda environment neither stacks duplicate filters nor keeps + reading a previous invocation's plugin. A single shared filter instance is + reused across all handlers. Args: plugin: The OTel plugin that resolves the current span context. @@ -105,7 +122,9 @@ def install_log_filter( None, ) if existing is not None: - # Reuse the already-installed filter so a single instance is shared. + # Reuse the already-installed filter so a single instance is shared, + # and point it at this invocation's plugin. + existing.bind(plugin) context_filter = existing continue if context_filter is None: diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py new file mode 100644 index 00000000..ddda8dcd --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py @@ -0,0 +1,90 @@ +"""Plugin factories for the bundled durable-execution OTel plugins. + +The SDK's plugin contract is a factory called once per invocation: +``DurableInstrumentationPluginFactory = Callable[[InvocationStartInfo], +DurableInstrumentationPlugin]``. The instance a factory returns serves exactly +that one invocation and is dropped when the invocation scope exits, so a plugin +keeps its per-invocation state in ordinary instance attributes. + +Both factories are callable classes rather than closures so the configuration +they were built with stays inspectable (``factory.config``) and so the entry +points below name an object with a readable type. + +Everything else the plugins need is resolved per invocation inside the plugin +itself: the tracer provider (which for the global-provider case may only be +installed after the handler module is imported), the tracer, and the +deterministic id generator and sampler installed on it. Those installs are +idempotent and scoped to the plugin's own tracer, so building a plugin per +invocation neither stacks wrappers nor disturbs other instrumentation scopes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + ExecutionOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.invocation_plugin import ( + InvocationOtelPlugin, +) + + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.plugin import InvocationStartInfo + + from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ) + + +class InvocationOtelPluginFactory: + """Builds one :class:`InvocationOtelPlugin` per invocation. + + Register the factory itself, not a plugin, with + ``@durable_execution(plugins=[...])``:: + + @durable_execution(plugins=[InvocationOtelPluginFactory()]) + def handler(event, context): ... + + Args: + config: Shared plugin configuration handed to every plugin this factory + builds. When omitted, each plugin uses defaults (globally configured + tracer provider, X-Ray extractor, "Workflow" span name, log + enrichment on). + """ + + def __init__(self, config: OtelPluginConfig | None = None) -> None: + self.config = config + + def __call__(self, info: InvocationStartInfo) -> InvocationOtelPlugin: + """Return this invocation's plugin. + + ``info`` is accepted because the SDK passes it, and is unused: the + plugin reads the same object again in ``on_invocation_start``, which is + where all of its invocation identity is derived. + """ + return InvocationOtelPlugin(self.config) + + +class ExecutionOtelPluginFactory: + """Builds one :class:`ExecutionOtelPlugin` per invocation. + + Args: + config: Shared plugin configuration handed to every plugin this factory + builds. When omitted, each plugin uses defaults. + """ + + def __init__(self, config: OtelPluginConfig | None = None) -> None: + self.config = config + + def __call__(self, info: InvocationStartInfo) -> ExecutionOtelPlugin: + """Return this invocation's plugin. ``info`` is unused; see above.""" + return ExecutionOtelPlugin(self.config) + + +INVOCATION_OTEL_PLUGIN_FACTORY = InvocationOtelPluginFactory() +"""Default-configured factory named by the ``otel-invocation`` entry point.""" + +EXECUTION_OTEL_PLUGIN_FACTORY = ExecutionOtelPluginFactory() +"""Default-configured factory named by the ``otel-execution`` entry point.""" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py deleted file mode 100644 index c4be734c..00000000 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py +++ /dev/null @@ -1,23 +0,0 @@ -from aws_durable_execution_sdk_python.plugin import ( - DurableInstrumentationPluginProvider, -) - -from aws_durable_execution_sdk_python_otel.execution_plugin import ( - ExecutionOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.invocation_plugin import ( - InvocationOtelPlugin, -) - - -INVOCATION_OTEL_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=InvocationOtelPlugin, - factory=InvocationOtelPlugin, - plugin_api_version=1, -) - -EXECUTION_OTEL_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=ExecutionOtelPlugin, - factory=ExecutionOtelPlugin, - plugin_api_version=1, -) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 63563493..fafdfa7f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -28,9 +28,11 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, ) -from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin -from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -138,18 +140,20 @@ def checkpoint( @pytest.mark.parametrize( - "plugin_type", - [InvocationOtelPlugin, ExecutionOtelPlugin], + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], ) def test_otel_wait_resume_spans_share_default_xray_execution_trace( monkeypatch: pytest.MonkeyPatch, - plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], + factory_type: type[InvocationOtelPluginFactory] | type[ExecutionOtelPluginFactory], ) -> None: monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) - plugin = plugin_type( + # The SDK takes a factory and calls it once per invocation, so the two + # invocations below are served by two plugin instances sharing this provider. + factory = factory_type( OtelPluginConfig( tracer_provider=provider, enrich_logger=False, @@ -164,7 +168,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(1), name="otel-wait") return context.step(complete_after_resume(), name="otel-after-resume") - handler = durable_execution(handler_impl, plugins=[plugin]) + handler = durable_execution(handler_impl, plugins=[factory]) initial_operations = [_execution_operation()] first_checkpoint, first_operations = _checkpoint_store(initial_operations) @@ -226,7 +230,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: after_resume = next(span for span in spans if span.name == "otel-after-resume") assert len(invocations) >= 2 - if plugin_type is InvocationOtelPlugin: + if factory_type is InvocationOtelPluginFactory: assert len(waits) >= 2 # one segment per invocation else: assert len(waits) == 1 # one span per operation @@ -238,7 +242,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: } assert after_resume.parent is not None - if plugin_type is InvocationOtelPlugin: + if factory_type is InvocationOtelPluginFactory: assert after_resume.parent.span_id in { span.context.span_id for span in invocations } diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 5c3cc3ec..8c94d257 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -73,9 +73,16 @@ def _assert_otel_context_balanced(): def _create_plugin( context_extractor=lambda _: None, + exporter: InMemorySpanExporter | None = None, ) -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: - """Create an ExecutionOtelPlugin wired to an in-memory exporter.""" - exporter = InMemorySpanExporter() + """Create an ExecutionOtelPlugin wired to an in-memory exporter. + + One plugin instance serves exactly one invocation, so a test that spans + invocations creates a plugin per invocation and passes the same ``exporter`` + to each -- the way the SDK's factory hands successive invocations distinct + instances that publish to one provider. + """ + exporter = exporter if exporter is not None else InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) plugin = ExecutionOtelPlugin( @@ -721,7 +728,13 @@ def test_suspended_operation_held_as_non_recording_placeholder(): def test_suspend_then_resume_operation_exports_one_deterministic_span(): - """An operation spanning invocations exports one deterministic span.""" + """An operation spanning invocations exports one deterministic span. + + Each invocation gets its own plugin instance, so nothing about the operation + is carried in memory from one invocation to the next: the single exported + span and its ID come from the deterministic derivation off the execution ARN, + and the replay guard is the hook's own ``is_replayed`` flag. + """ plugin, exporter = _create_plugin() operation_id = "wait-across-invocations" @@ -745,6 +758,7 @@ def test_suspend_then_resume_operation_exports_one_deterministic_span(): assert not [s for s in exporter.get_finished_spans() if s.name == "long-wait"] # Invocation N+1: the still-open operation is replayed, then completes. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_start( OperationStartInfo( @@ -785,6 +799,7 @@ def test_suspend_then_resume_operation_exports_one_deterministic_span(): # ReplayChildren/virtual child completion callbacks are replay-only and # must not re-export the terminal deterministic span in a later invocation. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_end( OperationEndInfo( @@ -836,7 +851,8 @@ def test_suspended_child_context_exports_one_span_on_replay(): # Nothing exported for the suspended context. assert not [s for s in exporter.get_finished_spans() if s.name == context_id] - # Invocation 2: the context replays and completes. + # Invocation 2: the context replays and completes, in a fresh plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_start( OperationStartInfo( @@ -903,10 +919,12 @@ def test_checkpointless_context_end_uses_a_non_negative_duration(): def test_virtual_context_replay_uses_unique_linked_segments(): - plugin, exporter = _create_plugin() + exporter = InMemorySpanExporter() context_id = "flat-branch" for _ in range(2): + # Each invocation is served by its own plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) # Virtual contexts have no durable START hook. plugin.on_user_function_start(_context_start_info(context_id)) @@ -1401,15 +1419,21 @@ def test_invocation_end_releases_scope_of_suspended_user_function(): assert plugin._context_tokens == {} -def test_warm_invocation_reuse_restores_ambient_span_each_time(): - """Verify repeated invocations leave the ambient Lambda span current.""" - plugin, _ = _create_plugin() +def test_successive_invocations_restore_ambient_span_each_time(): + """Verify each invocation's own plugin leaves the ambient Lambda span current. + + The SDK builds a plugin per invocation, so this drives three invocations + through three instances against one warm environment. What is asserted is + that no instance leaves a context attached behind it -- state carried in the + instance is irrelevant, because the instance is gone. + """ ambient_provider = TracerProvider() ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") token = otel_context.attach(trace.set_span_in_context(ambient)) try: warm_context = otel_context.get_current() for index in range(3): + plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = f"step-{index}" plugin.on_user_function_start(_step_start_info(operation_id)) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index df3e4b65..6e8fd15c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -50,6 +50,9 @@ ) from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + ExecutionOtelPluginFactory, +) START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -303,7 +306,9 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "_TRACER_PROVIDER", None) current_provider: list[ApiTracerProvider] = [ProxyTracerProvider()] monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) - plugin = ExecutionOtelPlugin( + # A factory, because the two invocations below need two plugin instances and + # the second one must resolve the provider installed after the first ran. + factory = ExecutionOtelPluginFactory( OtelPluginConfig( context_extractor=lambda _: None, enrich_logger=False, @@ -311,6 +316,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() + plugin = factory(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -319,6 +325,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () + plugin = factory(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index f108ef72..8fbf4445 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -77,17 +77,26 @@ def _assert_otel_context_balanced(): ) -def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: +def _create_plugin( + exporter: InMemorySpanExporter | None = None, +) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: """Create a plugin wired to an in-memory span exporter.""" - return _create_plugin_with_sampler() + return _create_plugin_with_sampler(exporter=exporter) def _create_plugin_with_sampler( sampler: Sampler | None = None, context_extractor=lambda _: None, + exporter: InMemorySpanExporter | None = None, ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: - """Create a plugin wired to an in-memory span exporter.""" - exporter = InMemorySpanExporter() + """Create a plugin wired to an in-memory span exporter. + + One plugin instance serves exactly one invocation, so a test that spans + invocations creates a plugin per invocation and passes the same ``exporter`` + to each -- the way the SDK's factory hands successive invocations distinct + instances that publish to one provider. + """ + exporter = exporter if exporter is not None else InMemorySpanExporter() trace_provider = TracerProvider(sampler=sampler) trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) plugin = InvocationOtelPlugin( @@ -1759,12 +1768,14 @@ def test_checkpointed_context_first_span_uses_deterministic_id(): def test_virtual_context_replay_uses_unique_linked_segments(): - plugin, exporter = _create_plugin() + exporter = InMemorySpanExporter() operation_id = "flat-branch" span_name = f"step-{operation_id}" workflow_span_id = derive_workflow_span_id(EXECUTION_ARN) for _ in range(2): + # Each invocation is served by its own plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) # Virtual contexts have no durable START hook. plugin.on_user_function_start( @@ -1937,15 +1948,21 @@ def test_ambient_span_is_current_again_after_full_lifecycle(): ambient.end() -def test_warm_invocation_reuse_does_not_accumulate_scopes(): - """Verify repeated invocations on one plugin instance stay balanced.""" - plugin, _ = _create_plugin() +def test_successive_invocations_do_not_accumulate_scopes(): + """Verify each invocation's own plugin leaves the warm environment balanced. + + The SDK builds a plugin per invocation, so this drives three invocations + through three instances. What is asserted is that no instance leaves a + context attached behind it; state carried in the instance is irrelevant, + because the instance is gone. + """ ambient_provider = TracerProvider() ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") token = otel_context.attach(trace.set_span_in_context(ambient)) try: warm_context = otel_context.get_current() for index in range(3): + plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = f"step-{index}" plugin.on_user_function_start(_user_function_start_info(operation_id)) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py index a50a06ac..12b0b6cb 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -53,6 +53,9 @@ ) from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + InvocationOtelPluginFactory, +) START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -253,7 +256,9 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "_TRACER_PROVIDER", None) current_provider: list[ApiTracerProvider] = [ProxyTracerProvider()] monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) - plugin = InvocationOtelPlugin( + # A factory, because the two invocations below need two plugin instances and + # the second one must resolve the provider installed after the first ran. + factory = InvocationOtelPluginFactory( OtelPluginConfig( context_extractor=lambda _: None, enrich_logger=False, @@ -261,6 +266,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() + plugin = factory(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -269,6 +275,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () + plugin = factory(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index 0cce516e..65debda1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -177,6 +177,37 @@ def test_install_log_filter_is_idempotent(): target.removeHandler(handler) +def test_install_log_filter_rebinds_to_the_current_invocations_plugin(): + """A later invocation's plugin takes over the already-installed filter. + + The handler outlives the invocation but a plugin instance does not, so + without rebinding the filter would keep reading the first invocation's + discarded plugin and stop correlating logs after that invocation. + """ + first_plugin, _ = _create_plugin() + second_plugin, _ = _create_plugin() + target = logging.getLogger("test.rebind") + handler = logging.NullHandler() + target.addHandler(handler) + try: + installed = install_log_filter(first_plugin, target_logger=target) + rebound = install_log_filter(second_plugin, target_logger=target) + + assert rebound is installed + assert installed is not None + assert installed._plugin is second_plugin + + second_plugin.on_invocation_start(_invocation_start_info()) + record = _make_record() + installed.filter(record) + + expected = second_plugin.get_current_span_context() + assert expected is not None + assert record.spanId == format(expected.span_id, "016x") + finally: + target.removeHandler(handler) + + def test_install_log_filter_reuses_single_instance_across_handlers(): """A single filter instance is shared across all handlers.""" plugin, _ = _create_plugin() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py new file mode 100644 index 00000000..5999fc60 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py @@ -0,0 +1,177 @@ +"""Tests for the bundled OTel plugin factories and the entry points naming them. + +The SDK plugin contract is a factory called once per invocation, so what these +tests have to establish is that the objects the package exposes -- and the ones +its entry points name -- are callables that build a plugin, and that each call +builds a NEW plugin. The old provider-shaped assertions (a declared +``plugin_type`` and an API version) have no counterpart: the contract carries +neither. +""" + +from __future__ import annotations + +import importlib +import logging +import tomllib +from datetime import UTC, datetime +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, +) + +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + ExecutionOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.invocation_plugin import ( + InvocationOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + EXECUTION_OTEL_PLUGIN_FACTORY, + INVOCATION_OTEL_PLUGIN_FACTORY, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +PLUGIN_ENTRY_POINT_GROUP = "aws_durable_execution.plugins" +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" + + +@pytest.fixture(autouse=True) +def _remove_installed_log_filters(): + """Detach any log filter a default-configured plugin installed. + + ``OtelPluginConfig.enrich_logger`` defaults to True, so building a plugin + from a default-configured factory attaches a filter to the root logger's + handlers, which outlive the test. + """ + yield + for handler in logging.getLogger().handlers: + for installed in [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ]: + handler.removeFilter(installed) + + +def _invocation_start_info() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC), + is_first_invocation=True, + ) + + +def _declared_entry_points() -> dict[str, str]: + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + project = tomllib.load(pyproject)["project"] + return project["entry-points"][PLUGIN_ENTRY_POINT_GROUP] + + +def _resolve(spec: str) -> object: + """Resolve a ``module:attribute`` entry-point value the way the SDK does. + + Resolved from the declared value rather than from installed distribution + metadata so the assertion holds wherever the tests run, installed or not. + """ + module_name, _, attribute = spec.partition(":") + return getattr(importlib.import_module(module_name), attribute) + + +@pytest.mark.parametrize( + ("factory_type", "plugin_type"), + [ + (InvocationOtelPluginFactory, InvocationOtelPlugin), + (ExecutionOtelPluginFactory, ExecutionOtelPlugin), + ], +) +def test_factory_builds_its_plugin_type( + factory_type: type, plugin_type: type[DurableInstrumentationPlugin] +) -> None: + factory = factory_type(OtelPluginConfig(enrich_logger=False)) + + assert isinstance(factory(_invocation_start_info()), plugin_type) + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_builds_a_fresh_plugin_per_invocation(factory_type: type) -> None: + """Two calls must never hand back the same instance. + + This is the whole point of the factory contract: the returned plugin holds + one invocation's state in ordinary instance fields, so a shared instance + would leak span registries and context tokens from one invocation into the + next. + """ + factory = factory_type(OtelPluginConfig(enrich_logger=False)) + + first = factory(_invocation_start_info()) + second = factory(_invocation_start_info()) + + assert first is not second + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_passes_its_config_to_every_plugin(factory_type: type) -> None: + config = OtelPluginConfig(workflow_span_name="Custom", enrich_logger=False) + factory = factory_type(config) + + assert factory.config is config + assert factory(_invocation_start_info())._config is config + assert factory(_invocation_start_info())._config is config + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_without_config_builds_a_default_configured_plugin( + factory_type: type, +) -> None: + factory = factory_type() + + assert factory.config is None + assert factory(_invocation_start_info())._config == OtelPluginConfig() + + +def test_module_level_factories_are_default_configured() -> None: + assert INVOCATION_OTEL_PLUGIN_FACTORY.config is None + assert EXECUTION_OTEL_PLUGIN_FACTORY.config is None + + +def test_declared_entry_points_name_the_bundled_factories() -> None: + entry_points = _declared_entry_points() + + assert set(entry_points) == {"otel-invocation", "otel-execution"} + assert _resolve(entry_points["otel-invocation"]) is INVOCATION_OTEL_PLUGIN_FACTORY + assert _resolve(entry_points["otel-execution"]) is EXECUTION_OTEL_PLUGIN_FACTORY + + +def test_declared_entry_points_resolve_to_callables_that_build_plugins() -> None: + """The entry points must satisfy the SDK's factory contract, not a provider. + + ``plugin_discovery._load_factory`` accepts anything callable, so a target + that resolved to a plugin class -- or to a plugin instance -- would load + without complaint and only fail at invocation time. + """ + expected = { + "otel-invocation": InvocationOtelPlugin, + "otel-execution": ExecutionOtelPlugin, + } + + for name, spec in _declared_entry_points().items(): + factory = _resolve(spec) + assert callable(factory) + assert not isinstance(factory, type) + assert isinstance(factory(_invocation_start_info()), expected[name]) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py deleted file mode 100644 index c47098b2..00000000 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py +++ /dev/null @@ -1,56 +0,0 @@ -from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, -) -from aws_durable_execution_sdk_python.plugin_discovery import ( - PLUGIN_ENVIRONMENT_VARIABLE, - load_configured_plugins, -) - -from aws_durable_execution_sdk_python_otel.execution_plugin import ( - ExecutionOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.invocation_plugin import ( - InvocationOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.plugin_provider import ( - EXECUTION_OTEL_PLUGIN_PROVIDER, - INVOCATION_OTEL_PLUGIN_PROVIDER, -) - - -def test_invocation_otel_plugin_provider_uses_current_plugin_api() -> None: - assert INVOCATION_OTEL_PLUGIN_PROVIDER.plugin_type is InvocationOtelPlugin - assert ( - INVOCATION_OTEL_PLUGIN_PROVIDER.plugin_api_version - == DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION - ) - - -def test_invocation_otel_plugin_provider_creates_invocation_plugin() -> None: - assert isinstance(INVOCATION_OTEL_PLUGIN_PROVIDER.factory(), InvocationOtelPlugin) - - -def test_execution_otel_plugin_provider_uses_current_plugin_api() -> None: - assert EXECUTION_OTEL_PLUGIN_PROVIDER.plugin_type is ExecutionOtelPlugin - assert ( - EXECUTION_OTEL_PLUGIN_PROVIDER.plugin_api_version - == DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION - ) - - -def test_execution_otel_plugin_provider_creates_execution_plugin() -> None: - assert isinstance(EXECUTION_OTEL_PLUGIN_PROVIDER.factory(), ExecutionOtelPlugin) - - -def test_installed_otel_entry_points_load_both_plugin_types() -> None: - plugins = load_configured_plugins( - None, - environment={ - PLUGIN_ENVIRONMENT_VARIABLE: "otel-invocation,otel-execution", - }, - ) - - assert [type(plugin) for plugin in plugins] == [ - InvocationOtelPlugin, - ExecutionOtelPlugin, - ] diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py index 41789171..be00a7f7 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py @@ -37,7 +37,12 @@ class RecordingWaitPlugin(DurableInstrumentationPlugin): - """Records end notifications for the wait and counts invocations.""" + """Records end notifications for the wait and counts invocations. + + State is class-level on purpose: the SDK builds a fresh instance per + invocation, and this test spans two invocations, so what it asserts on has + to outlive any single instance. + """ invocation_count: ClassVar[int] = 0 wait_end_infos: ClassVar[list[OperationEndInfo]] = [] @@ -61,7 +66,9 @@ def _wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return "done" -wait_handler = durable_execution(_wait_handler, plugins=[RecordingWaitPlugin()]) +wait_handler = durable_execution( + _wait_handler, plugins=[lambda _info: RecordingWaitPlugin()] +) def test_wait_completed_during_suspend_is_delivered_as_new() -> None: diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index bf7776f4..a241e23b 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -40,15 +40,21 @@ DURABLE_EXECUTION_PLUGINS=otel-invocation,example_audit The SDK resolves those names from the `aws_durable_execution.plugins` Python entry-point group when the decorated handler is initialized. An unset or blank variable preserves the existing behavior. The decorator's `plugins` argument -remains supported; explicit plugins run first and take precedence when a -dynamic provider creates the same concrete plugin type. +remains supported; explicit factories run first, and a factory passed to the +decorator is not registered a second time through the environment. -Provider packages expose a versioned factory: +A plugin is registered as a *factory*, not as an instance. A factory is any +callable taking the invocation's `InvocationStartInfo` and returning a +`DurableInstrumentationPlugin`; the SDK calls it once per invocation, so the +instance it returns serves that one invocation only and can hold per-execution +state in ordinary attributes. + +Provider packages expose such a factory: ```python from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + InvocationStartInfo, ) @@ -56,26 +62,20 @@ class AuditPlugin(DurableInstrumentationPlugin): pass -AUDIT_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=AuditPlugin, - factory=AuditPlugin, - plugin_api_version=1, -) +def audit_plugin_factory(info: InvocationStartInfo) -> AuditPlugin: + return AuditPlugin() ``` -Register the provider in the package's `pyproject.toml`: +Register the factory in the package's `pyproject.toml`: ```toml [project.entry-points."aws_durable_execution.plugins"] -example_audit = "example_audit:AUDIT_PLUGIN_PROVIDER" +example_audit = "example_audit:audit_plugin_factory" ``` -Set `plugin_api_version` to the literal API version the provider implements. -Update it only after verifying the provider against that API version. - Provider names must be unique across installed distributions. Missing, -ambiguous, incompatible, or invalid providers raise `PluginLoadError` during -handler initialization with the provider and distribution details. +ambiguous, or non-callable providers raise `PluginLoadError` during handler +initialization with the provider and distribution details. ## 🚀 Quick Start diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 8ca5ef90..6a21fd67 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -27,7 +27,7 @@ OperationUpdate, ) from aws_durable_execution_sdk_python.plugin import ( - DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, PluginExecutor, ) from aws_durable_execution_sdk_python.plugin_discovery import ( @@ -169,7 +169,7 @@ def durable_execution( func: Callable[[Any, DurableContext], Any] | None = None, *, boto3_client: Boto3LambdaClient | None = None, - plugins: list[DurableInstrumentationPlugin] | None = None, + plugins: list[DurableInstrumentationPluginFactory] | None = None, ) -> Callable[[Any, LambdaContext], Any]: """ Decorator to create a durable execution handler. @@ -177,7 +177,10 @@ def durable_execution( Args: func: The user function to decorate boto3_client: Optional boto3 Lambda client to use - plugins: Optional list of instrumentation plugins to use + plugins: Optional list of instrumentation plugin factories. Each factory + is called once per invocation with that invocation's + ``InvocationStartInfo``, and the instance it returns serves only that + invocation. """ # Decorator called with parameters if func is None: 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 e9549d30..ed46429e 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 @@ -28,8 +28,6 @@ logger = logging.getLogger(__name__) -DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION = 1 - class InvocationStatus(Enum): """Invocation outcomes exposed to instrumentation plugins.""" @@ -451,25 +449,47 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: pass -@dataclass(frozen=True) -class DurableInstrumentationPluginProvider: - """Versioned factory exposed through the plugin entry-point group.""" +DurableInstrumentationPluginFactory = Callable[ + [InvocationStartInfo], DurableInstrumentationPlugin +] +"""Builds one plugin instance for one invocation. + +Called once per invocation with that invocation's :class:`InvocationStartInfo` -- +the same object the instance's ``on_invocation_start`` then receives -- before any +hook fires. The instance serves only that invocation and is dropped when it +returns, so a plugin can hold per-execution state in ordinary instance +attributes without keying it by execution ARN. + +A plain ``Callable`` alias rather than a ``Protocol``: the shape has exactly one +call signature and no other members, so a Protocol would only add a name. The +alias is also the more permissive of the two, because ``Callable`` parameters are +positional-only -- a factory may name its parameter whatever reads best +(``lambda info: ...``, ``def build(invocation): ...``), where a ``__call__`` +Protocol would pin that name. Anything callable satisfies it: a lambda, a +module-level function, a ``functools.partial``, or a class whose ``__init__`` +takes the info. +""" + - plugin_type: type[DurableInstrumentationPlugin] - factory: Callable[[], DurableInstrumentationPlugin] - plugin_api_version: int +def _factory_name(factory: object) -> str: + """Best available name for a factory, for log messages.""" + return getattr(factory, "__qualname__", None) or type(factory).__name__ class PluginExecutor: - def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): - self._plugins = plugins or [] + def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): + # Factories live for the life of the handler; the instances they build do + # not. _plugins is populated in on_invocation_start and emptied when the + # invocation scope exits, so one instance never spans two invocations. + self._plugin_factories = list(plugins or []) + self._plugins: list[DurableInstrumentationPlugin] = [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None @contextlib.contextmanager def run(self): - if self._plugins: + if self._plugin_factories: self._executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="plugin-executor", @@ -482,6 +502,37 @@ def run(self): # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) + # Drop this invocation's plugin instances. After the pool has + # drained, so no queued dispatch still holds one: nothing reachable + # from this handler-lifetime executor outlives the invocation. + self._plugins = [] + + def _create_plugins(self, info: InvocationStartInfo) -> None: + """Build this invocation's plugin instances from its start info. + + Called once per invocation, before the first hook is dispatched. A + factory that raises or returns ``None`` is contained exactly as a failing + hook is -- logged and skipped -- so a broken plugin cannot disrupt the + execution. The remaining factories still produce their instances. + """ + plugins: list[DurableInstrumentationPlugin] = [] + for factory in self._plugin_factories: + try: + plugin = factory(info) + except Exception: + # log and ignore the exception + logger.exception( + "Plugin factory %s exception ignored", _factory_name(factory) + ) + continue + if plugin is None: + logger.error( + "Plugin factory %s returned None; plugin ignored", + _factory_name(factory), + ) + continue + plugins.append(plugin) + self._plugins = plugins @staticmethod def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: @@ -535,11 +586,13 @@ def _snapshot_operation_infos( plugin that stashes the info and reads it later still sees the state as of its own hook. - Skipped entirely when no plugins are registered -- ``durable_execution()`` + Skipped entirely when no plugins are configured -- ``durable_execution()`` passes a provider unconditionally, so without this gate a plugin-free - execution would pay for a view nothing can read. + execution would pay for a view nothing can read. The gate reads the + factory list, not the instances: this runs while the start info is being + built, before any instance exists. """ - if not self._plugins or operations_provider is None: + if not self._plugin_factories or operations_provider is None: return {} try: return _to_operation_info_map(operations_provider()) @@ -572,7 +625,9 @@ def on_invocation_start( ``UpdatedOperationIds`` -- those updated while suspended. """ aws_request_id = lambda_context.aws_request_id if lambda_context else None - self._operations_provider = operations_provider if self._plugins else None + self._operations_provider = ( + operations_provider if self._plugin_factories else None + ) operations = self._snapshot_operation_infos(operations_provider) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, @@ -587,6 +642,9 @@ def on_invocation_start( if operation_id in operations }, ) + # Build this invocation's plugin instances from the very info their first + # hook receives, and before that hook is dispatched. + self._create_plugins(self._invocation_status) self.execute_plugins(self._invocation_status, sync=True) def _snapshot_execution_input(self, execution_input: Any) -> Any: @@ -602,12 +660,12 @@ def _snapshot_execution_input(self, execution_input: Any) -> Any: The copy is eager rather than deferred: the handler starts running immediately after this hook, so a lazily-taken snapshot could already have observed the handler's mutations. It is skipped when no plugins are - registered, so non-plugin executions pay nothing. + configured, so non-plugin executions pay nothing. The snapshot is shared by all plugins for this invocation; plugins should still treat it as read-only with respect to each other. """ - if not self._plugins or execution_input is None: + if not self._plugin_factories or execution_input is None: return execution_input try: return copy.deepcopy(execution_input) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index ffb07260..57d1d969 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -4,13 +4,11 @@ import os from collections.abc import Mapping, Sequence from importlib import metadata +from typing import cast -from aws_durable_execution_sdk_python.__about__ import __version__ from aws_durable_execution_sdk_python.exceptions import PluginLoadError from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + DurableInstrumentationPluginFactory, ) @@ -56,98 +54,62 @@ def _qualified_type_name(value: object) -> str: return f"{value_type.__module__}.{value_type.__qualname__}" -def _qualified_class_name(value_type: type[object]) -> str: - return f"{value_type.__module__}.{value_type.__qualname__}" - - -def _load_provider( +def _load_factory( plugin_name: str, entry_point: metadata.EntryPoint -) -> DurableInstrumentationPluginProvider: +) -> DurableInstrumentationPluginFactory: + """Resolve an entry point to a plugin factory. + + Only two things can still be checked here. The entry point has to import, + and what it resolves to has to be callable. Nothing more is knowable without + calling the factory, and calling it at load time is precisely what this + design avoids: the instance belongs to an invocation, and there is no + invocation yet. A factory that then misbehaves at invocation time is + contained by :meth:`PluginExecutor._create_plugins`. + """ try: - provider = entry_point.load() + factory = entry_point.load() except Exception as error: raise PluginLoadError( - f"Failed to load durable instrumentation plugin provider " + f"Failed to load durable instrumentation plugin factory " f"'{plugin_name}' from '{entry_point.value}' " f"({_distribution_name(entry_point)}): {error}" ) from error - if not isinstance(provider, DurableInstrumentationPluginProvider): + if not callable(factory): raise PluginLoadError( f"Durable instrumentation plugin entry point '{plugin_name}' must " - "resolve to DurableInstrumentationPluginProvider, but resolved to " - f"{_qualified_type_name(provider)}." - ) - - if provider.plugin_api_version != DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION: - raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' declares " - f"plugin API version {provider.plugin_api_version}, but " - f"aws-durable-execution-sdk-python {__version__} supports plugin API " - f"version {DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION}. Install " - "compatible SDK and plugin package versions." - ) - - declared_plugin_type: object = provider.plugin_type - if not isinstance(declared_plugin_type, type) or not issubclass( - declared_plugin_type, DurableInstrumentationPlugin - ): - declared_type_name = ( - _qualified_class_name(declared_plugin_type) - if isinstance(declared_plugin_type, type) - else _qualified_type_name(declared_plugin_type) - ) - raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' declares " - f"invalid plugin type {declared_type_name}; " - "expected a DurableInstrumentationPlugin subclass." + "resolve to a callable plugin factory, but resolved to " + f"{_qualified_type_name(factory)}." ) - return provider - - -def _create_plugin( - plugin_name: str, - entry_point: metadata.EntryPoint, - provider: DurableInstrumentationPluginProvider, -) -> DurableInstrumentationPlugin: - try: - plugin = provider.factory() - except Exception as error: - raise PluginLoadError( - f"Failed to create durable instrumentation plugin '{plugin_name}' " - f"from '{entry_point.value}' ({_distribution_name(entry_point)}): " - f"{error}" - ) from error - - if type(plugin) is not provider.plugin_type: - raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' returned " - f"{_qualified_type_name(plugin)}; expected " - f"{_qualified_class_name(provider.plugin_type)}." - ) - - return plugin + return cast(DurableInstrumentationPluginFactory, factory) def load_configured_plugins( - explicit_plugins: Sequence[DurableInstrumentationPlugin] | None, + explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, *, environment: Mapping[str, str] | None = None, -) -> list[DurableInstrumentationPlugin]: - """Combine explicit plugins with providers selected through the environment. - - Explicit plugins retain their order. Dynamically selected plugins follow in - configured order. When discovery creates a plugin whose concrete type is - already registered, the first registration wins, so explicit registration - takes precedence. +) -> list[DurableInstrumentationPluginFactory]: + """Combine explicit plugin factories with those selected through the environment. + + Explicit factories retain their order. Dynamically selected factories follow + in configured order. Every returned factory is called once per invocation. + + A factory already registered explicitly is not registered a second time + through the environment. The check is by factory identity, which is what is + knowable here: the old shape declared a ``plugin_type`` and could dedup on + it, but a factory is opaque until called, and calling it at load time is what + this design avoids. Identity still covers the case the plugin packages + document -- the same provider callable both passed to the decorator and named + in ``DURABLE_EXECUTION_PLUGINS``. Two *different* factories that happen to + build the same plugin type will now both be registered. """ - resolved_plugins = list(explicit_plugins or []) + resolved_factories = list(explicit_plugins or []) resolved_environment = os.environ if environment is None else environment plugin_names = _parse_configured_plugin_names(resolved_environment) if not plugin_names: - return resolved_plugins + return resolved_factories try: discovered_entry_points = list( @@ -163,10 +125,6 @@ def load_configured_plugins( for entry_point in discovered_entry_points: entry_points_by_name.setdefault(entry_point.name, []).append(entry_point) - registered_types: dict[type[DurableInstrumentationPlugin], str] = { - type(plugin): "the decorator's plugins argument" for plugin in resolved_plugins - } - for plugin_name in plugin_names: matching_entry_points = entry_points_by_name.get(plugin_name, []) if not matching_entry_points: @@ -190,20 +148,15 @@ def load_configured_plugins( "duplicate provider package." ) - entry_point = matching_entry_points[0] - provider = _load_provider(plugin_name, entry_point) - if existing_registration := registered_types.get(provider.plugin_type): + factory = _load_factory(plugin_name, matching_entry_points[0]) + if any(factory is registered for registered in resolved_factories): logger.warning( - "Skipping dynamically configured plugin '%s' because %s is " - "already registered by %s.", + "Skipping dynamically configured plugin '%s' because the same " + "plugin factory is already registered.", plugin_name, - _qualified_class_name(provider.plugin_type), - existing_registration, ) continue - plugin = _create_plugin(plugin_name, entry_point, provider) - resolved_plugins.append(plugin) - registered_types[provider.plugin_type] = f"dynamic provider '{plugin_name}'" + resolved_factories.append(factory) - return resolved_plugins + return resolved_factories 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 1b75325b..110ef919 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -63,7 +63,11 @@ WaitForConditionDecision, ) from tests.serdes_test import CustomDictSerDes -from tests.test_helpers import operation_id_sequence +from tests.test_helpers import ( + operation_id_sequence, + plugin_factory, + plugin_invocation, +) def create_test_context( @@ -3014,9 +3018,9 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(f"end:{info.operation_id}") - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) step_body_calls: list[bool] = [] - with plugin_executor.run(): + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 @@ -3062,8 +3066,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 @@ -3100,8 +3104,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py index 3cad3412..3cd9c806 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -26,8 +26,11 @@ OperationStatus, OperationType, ) -from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin -from tests.test_helpers import operation_id_sequence +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, +) +from tests.test_helpers import operation_id_sequence, plugin_factory class _MapRecordingPlugin(DurableInstrumentationPlugin): @@ -114,7 +117,7 @@ def test_operation_maps_on_a_completing_invocation(): """The start map holds the prior state; the end map sees the step added.""" plugin = _MapRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return context.step(lambda _ctx: "stepped", name="greet") @@ -150,17 +153,29 @@ def test_operation_maps_across_suspend_and_replay(): Invocation 1 suspends on a wait. Invocation 2 replays with the wait already SUCCEEDED and its id in ``UpdatedOperationIds``, which is exactly what ``updated_operations`` is derived from. + + One handler serves both invocations, as in production. Previously this test + had to declare a second handler with its own plugin instance, because the + single shared instance would have interleaved both invocations' records into + one list. Now the factory builds an instance per invocation, so each + instance's ``starts[0]``/``ends[0]`` unambiguously describes its own + invocation -- and the test asserts that separation directly. """ wait_id = next(operation_id_sequence()) - # --- Invocation 1: the wait starts and the execution suspends. - first = _MapRecordingPlugin() + built: list[_MapRecordingPlugin] = [] - @durable_execution(plugins=[first]) - def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + def build_plugin(info: InvocationStartInfo) -> _MapRecordingPlugin: + plugin = _MapRecordingPlugin() + built.append(plugin) + return plugin + + @durable_execution(plugins=[build_plugin]) + def wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 context.wait(Duration.from_seconds(60)) return "done" + # --- Invocation 1: the wait starts and the execution suspends. with patch( "aws_durable_execution_sdk_python.execution.LambdaClient" ) as mock_client_class: @@ -168,9 +183,11 @@ def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG mock_client.checkpoint = _tracking_checkpoint() mock_client_class.initialize_client.return_value = mock_client - first_result = suspending_handler(_event(), _lambda_context()) + first_result = wait_handler(_event(), _lambda_context()) assert first_result["Status"] == InvocationStatus.PENDING.value + assert len(built) == 1 + first = built[0] start_operations, start_updated = first.starts[0] assert start_operations == ["execution-1"] assert start_updated == [] @@ -179,13 +196,6 @@ def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG assert wait_id in end_operations # --- Invocation 2: replay with the wait completed externally. - replay = _MapRecordingPlugin() - - @durable_execution(plugins=[replay]) - def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 - context.wait(Duration.from_seconds(60)) - return "done" - completed_wait = { "Id": wait_id, "Type": OperationType.WAIT.value, @@ -200,13 +210,21 @@ def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG00 mock_client.checkpoint = _tracking_checkpoint() mock_client_class.initialize_client.return_value = mock_client - replay_result = replayed_handler( + replay_result = wait_handler( _event(extra_operations=[completed_wait], updated_operation_ids=[wait_id]), _lambda_context(), ) assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + # The second invocation got its own instance, and the first one recorded + # nothing further after it returned. + assert len(built) == 2 + replay = built[1] + assert replay is not first + assert len(first.starts) == 1 + assert len(first.ends) == 1 + start_operations, start_updated = replay.starts[0] # The replay start map carries the prior state, including the wait. assert sorted(["execution-1", wait_id]) == start_operations @@ -222,7 +240,7 @@ def test_updated_operations_ignores_ids_absent_from_the_map(): """An id the execution state does not carry must not appear in the subset.""" plugin = _MapRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return "ok" diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py index 6b934423..e0c86b17 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py @@ -27,7 +27,7 @@ OperationType, ) from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin -from tests.test_helpers import operation_id_sequence +from tests.test_helpers import operation_id_sequence, plugin_factory class _PayloadRecordingPlugin(DurableInstrumentationPlugin): @@ -116,7 +116,7 @@ def test_plugin_sees_execution_input_and_result_end_to_end(): """A completing invocation surfaces the input on both hooks and the result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> dict: # noqa: ARG001 return {"greeting": f"Hello, {event['name']}!"} @@ -146,7 +146,7 @@ def test_plugin_payload_surfaces_on_suspending_invocation(): """A suspending invocation carries the input but no execution result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(60)) return f"done-{event['name']}" @@ -175,7 +175,7 @@ def test_plugin_payload_surfaces_on_replay_invocation(): """A replay past a completed wait carries the input and the terminal result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(60)) return f"done-{event['name']}" @@ -232,7 +232,7 @@ def on_invocation_end(self, info) -> None: plugin = _MutatingPlugin() handler_saw: dict[str, Any] = {} - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 handler_saw.update( {"top": dict(event), "nested_items": list(event["nested"]["items"])} diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py index 7b30828d..71eaeca2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py @@ -28,6 +28,7 @@ UserFunctionOutcome, UserFunctionStartInfo, ) +from tests.test_helpers import plugin_factory @dataclass(frozen=True) @@ -156,7 +157,7 @@ def child_function(context: DurableContext) -> str: def user_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return context.run_in_child_context(child_function, name="charge") - handler = durable_execution(user_handler, plugins=[plugin]) + handler = durable_execution(user_handler, plugins=[plugin_factory(plugin)]) first_checkpoint, first_operations = _tracking_checkpoint() with patch( diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 85ea72a9..7d8b6898 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -60,6 +60,7 @@ WaitDetails, ) from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from tests.test_helpers import plugin_factory LARGE_RESULT = "large_success" * 1024 * 1024 @@ -3482,24 +3483,24 @@ def on_operation_attempt_end(self, info): def test_durable_execution_loads_plugins_when_handler_is_initialized(): - """Configured plugins are resolved once while the decorator initializes.""" - explicit_plugin = _RecordingPlugin() - resolved_plugin = _RecordingPlugin() + """Configured factories are resolved once while the decorator initializes.""" + explicit_factory = plugin_factory(_RecordingPlugin()) + resolved_factory = plugin_factory(_RecordingPlugin()) with ( warnings.catch_warnings(), patch( "aws_durable_execution_sdk_python.execution.load_configured_plugins", - return_value=[explicit_plugin, resolved_plugin], + return_value=[explicit_factory, resolved_factory], ) as load_plugins, ): warnings.simplefilter("error", FutureWarning) - @durable_execution(plugins=[explicit_plugin]) + @durable_execution(plugins=[explicit_factory]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} - load_plugins.assert_called_once_with([explicit_plugin]) + load_plugins.assert_called_once_with([explicit_factory]) assert callable(test_handler) @@ -3514,7 +3515,7 @@ def test_durable_execution_with_plugins_success(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3545,7 +3546,7 @@ def test_durable_execution_forwards_execution_input_to_plugins(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"echoed": event["name"]} @@ -3573,7 +3574,7 @@ def test_durable_execution_surfaces_empty_input_as_empty_mapping(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> str: return "ok" @@ -3608,7 +3609,7 @@ def on_invocation_start(self, info): def on_invocation_end(self, info): observed["end_input"] = dict(info.execution_input) - @durable_execution(plugins=[_MutatingPlugin()]) + @durable_execution(plugins=[plugin_factory(_MutatingPlugin())]) def test_handler(event: Any, context: DurableContext) -> dict: observed["handler_saw"] = dict(event) # Direction B: handler mutates its event after the start hook fired. @@ -3642,7 +3643,7 @@ class _NestedMutatingPlugin(DurableInstrumentationPlugin): def on_invocation_start(self, info): info.execution_input["outer"]["inner"].append("from_plugin") - @durable_execution(plugins=[_NestedMutatingPlugin()]) + @durable_execution(plugins=[plugin_factory(_NestedMutatingPlugin())]) def test_handler(event: Any, context: DurableContext) -> dict: observed["handler_saw"] = deepcopy(event) return {"ok": True} @@ -3668,7 +3669,7 @@ def test_durable_execution_with_plugins_failure(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: msg = "user error" raise ValueError(msg) @@ -3694,7 +3695,7 @@ def test_durable_execution_with_plugins_pending(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: raise SuspendExecution("test") @@ -3717,7 +3718,7 @@ def test_durable_execution_with_plugins_retryable_error(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: msg = "Retriable error" raise InvocationError(msg) @@ -3744,7 +3745,7 @@ def test_durable_execution_with_multiple_plugins(): plugin1 = _RecordingPlugin() plugin2 = _RecordingPlugin() - @durable_execution(plugins=[plugin1, plugin2]) + @durable_execution(plugins=[plugin_factory(plugin1), plugin_factory(plugin2)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3772,7 +3773,9 @@ def test_durable_execution_with_failing_plugin_does_not_break_execution(): failing_plugin = _FailingPlugin() recording_plugin = _RecordingPlugin() - @durable_execution(plugins=[failing_plugin, recording_plugin]) + @durable_execution( + plugins=[plugin_factory(failing_plugin), plugin_factory(recording_plugin)] + ) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3788,6 +3791,80 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert "invocation_end:SUCCEEDED" in recording_plugin.calls +def test_durable_execution_builds_a_plugin_per_invocation(): + """One handler, two invocations, two instances -- and no crosstalk. + + This is the LMI-driven property: a plugin may hold per-execution state in + ordinary instance attributes because the instance never outlives the + invocation it was built for. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + built: list[_RecordingPlugin] = [] + factory_arns: list[str | None] = [] + + def build_plugin(info) -> _RecordingPlugin: + factory_arns.append(info.execution_arn) + plugin = _RecordingPlugin() + built.append(plugin) + return plugin + + @durable_execution(plugins=[build_plugin]) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"result": "success"} + + for _ in range(2): + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + + assert len(built) == 2 + assert built[0] is not built[1] + # The factory saw the execution it was being built for. + assert factory_arns == [ + "arn:test:execution/exec1", + "arn:test:execution/exec1", + ] + # Each instance recorded exactly one invocation's worth of hooks. + for plugin in built: + assert plugin.calls.count("invocation_start") == 1 + assert plugin.calls.count("invocation_end:SUCCEEDED") == 1 + + +def test_durable_execution_with_failing_plugin_factory_does_not_break_execution(): + """A factory that raises is contained exactly as a failing hook is.""" + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + recording_plugin = _RecordingPlugin() + + def exploding_factory(info) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + @durable_execution(plugins=[exploding_factory, plugin_factory(recording_plugin)]) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"result": "success"} + + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + # The other plugin is unaffected. + assert "invocation_start" in recording_plugin.calls + assert "invocation_end:SUCCEEDED" in recording_plugin.calls + + def test_durable_execution_with_no_plugins(): """Test that passing no plugins (None) works correctly.""" mock_client = Mock(spec=DurableServiceClient) @@ -3843,7 +3920,7 @@ def test_durable_execution_decorator_with_plugins_and_boto3_client(): # When using DurableExecutionInvocationInputWithClient, boto3_client is ignored # but we verify the decorator accepts both parameters - @durable_execution(boto3_client=None, plugins=[plugin]) + @durable_execution(boto3_client=None, plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index 51200ba3..a5153132 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -2,17 +2,14 @@ import logging import os -from collections.abc import Callable -from typing import cast from unittest.mock import Mock, patch import pytest from aws_durable_execution_sdk_python.exceptions import PluginLoadError from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + InvocationStartInfo, ) from aws_durable_execution_sdk_python.plugin_discovery import ( PLUGIN_ENTRY_POINT_GROUP, @@ -21,6 +18,13 @@ ) +INVOCATION_START_INFO = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:exec", + is_first_invocation=True, +) + + class _PluginA(DurableInstrumentationPlugin): pass @@ -29,6 +33,14 @@ class _PluginB(DurableInstrumentationPlugin): pass +def _plugin_a_factory(info: InvocationStartInfo) -> _PluginA: + return _PluginA() + + +def _plugin_b_factory(info: InvocationStartInfo) -> _PluginB: + return _PluginB() + + class _FakeDistribution: def __init__(self, name: str) -> None: self.metadata = {"Name": name} @@ -59,32 +71,10 @@ def load(self) -> object: return self._loaded_value -def _provider( - factory: Callable[[], object], - *, - plugin_type: type[DurableInstrumentationPlugin] = _PluginA, - plugin_api_version: int = DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, -) -> DurableInstrumentationPluginProvider: - return DurableInstrumentationPluginProvider( - plugin_type=plugin_type, - factory=cast(Callable[[], DurableInstrumentationPlugin], factory), - plugin_api_version=plugin_api_version, - ) - - -def test_plugin_provider_requires_authored_api_version() -> None: - with pytest.raises(TypeError, match="plugin_api_version"): - DurableInstrumentationPluginProvider( - plugin_type=_PluginA, - factory=_PluginA, - ) # type: ignore[call-arg] - - @pytest.mark.parametrize("configured_value", [None, "", " "]) -def test_unconfigured_discovery_preserves_explicit_plugins( +def test_unconfigured_discovery_preserves_explicit_factories( configured_value: str | None, ) -> None: - explicit_plugin = _PluginA() environment = ( {} if configured_value is None @@ -95,16 +85,16 @@ def test_unconfigured_discovery_preserves_explicit_plugins( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points" ) as entry_points: result = load_configured_plugins( - [explicit_plugin], + [_plugin_a_factory], environment=environment, ) - assert result == [explicit_plugin] + assert result == [_plugin_a_factory] entry_points.assert_not_called() def test_discovery_uses_process_environment_by_default() -> None: - entry_point = _FakeEntryPoint("a", _provider(_PluginA)) + entry_point = _FakeEntryPoint("a", _plugin_a_factory) with ( patch.dict( @@ -119,25 +109,36 @@ def test_discovery_uses_process_environment_by_default() -> None: ): result = load_configured_plugins(None) - assert len(result) == 1 - assert isinstance(result[0], _PluginA) + assert result == [_plugin_a_factory] entry_points.assert_called_once_with(group=PLUGIN_ENTRY_POINT_GROUP) -def test_discovery_preserves_configured_order() -> None: - factory_calls: list[str] = [] +def test_discovery_returns_factories_without_calling_them() -> None: + """Discovery resolves factories only; instances belong to an invocation. - def create_a() -> _PluginA: - factory_calls.append("a") - return _PluginA() + Nothing is constructed at load time, so no plugin instance exists outside + the invocation that will use it. + """ + factory = Mock(return_value=_PluginA()) + entry_point = _FakeEntryPoint("a", factory) + + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + None, + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert result == [factory] + factory.assert_not_called() - def create_b() -> _PluginB: - factory_calls.append("b") - return _PluginB() +def test_discovery_preserves_configured_order() -> None: entry_points = [ - _FakeEntryPoint("b", _provider(create_b, plugin_type=_PluginB)), - _FakeEntryPoint("a", _provider(create_a)), + _FakeEntryPoint("b", _plugin_b_factory), + _FakeEntryPoint("a", _plugin_a_factory), ] with patch( @@ -149,8 +150,83 @@ def create_b() -> _PluginB: environment={PLUGIN_ENVIRONMENT_VARIABLE: " a, b "}, ) - assert [type(plugin) for plugin in result] == [_PluginA, _PluginB] - assert factory_calls == ["a", "b"] + assert result == [_plugin_a_factory, _plugin_b_factory] + assert [type(factory(INVOCATION_START_INFO)) for factory in result] == [ + _PluginA, + _PluginB, + ] + + +def test_explicit_factories_precede_discovered_factories() -> None: + entry_point = _FakeEntryPoint("b", _plugin_b_factory) + + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "b"}, + ) + + assert result == [_plugin_a_factory, _plugin_b_factory] + + +def test_explicit_registration_wins_over_the_same_discovered_factory( + caplog: pytest.LogCaptureFixture, +) -> None: + """The same callable passed explicitly and named in the env registers once. + + This is the narrowed form of the old type-based precedence rule. Dedup by + declared plugin type is gone with the provider object; identity still covers + the documented double-registration case. + """ + entry_point = _FakeEntryPoint("a", _plugin_a_factory) + + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ), + caplog.at_level( + logging.WARNING, + logger="aws_durable_execution_sdk_python.plugin_discovery", + ), + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert result == [_plugin_a_factory] + assert "already registered" in caplog.text + + +def test_distinct_factories_for_one_plugin_type_are_both_registered() -> None: + """Type-level dedup is gone: two distinct factories both register. + + Recorded deliberately. The provider object declared a ``plugin_type`` that + discovery could compare without constructing anything; a factory is opaque + until called, and calling it at load time would build an instance outside any + invocation. Callers that both pass a factory and name a different one in the + environment now get both plugins. + """ + + def another_plugin_a_factory(info: InvocationStartInfo) -> _PluginA: + return _PluginA() + + entry_point = _FakeEntryPoint("a", another_plugin_a_factory) + + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert result == [_plugin_a_factory, another_plugin_a_factory] @pytest.mark.parametrize("configured_value", ["a,,b", ",a", "a,"]) @@ -177,7 +253,7 @@ def test_discovery_rejects_duplicate_configured_names() -> None: def test_discovery_reports_missing_provider_and_available_names() -> None: - entry_point = _FakeEntryPoint("available", _provider(_PluginA)) + entry_point = _FakeEntryPoint("available", _plugin_a_factory) with ( patch( @@ -216,12 +292,12 @@ def test_discovery_rejects_ambiguous_provider_name() -> None: entry_points = [ _FakeEntryPoint( "duplicate", - _provider(_PluginA), + _plugin_a_factory, distribution_name="package-a", ), _FakeEntryPoint( "duplicate", - _provider(_PluginB), + _plugin_b_factory, distribution_name="package-b", ), ] @@ -256,10 +332,10 @@ def test_discovery_wraps_entry_point_enumeration_failure() -> None: ) -def test_discovery_wraps_provider_load_failure() -> None: +def test_discovery_wraps_factory_load_failure() -> None: entry_point = _FakeEntryPoint( "a", - _provider(_PluginA), + _plugin_a_factory, load_error=ImportError("missing dependency"), ) @@ -275,113 +351,19 @@ def test_discovery_wraps_provider_load_failure() -> None: environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "Failed to load durable instrumentation plugin provider 'a'" in str( + assert "Failed to load durable instrumentation plugin factory 'a'" in str( error.value ) assert "test-plugin-package" in str(error.value) assert isinstance(error.value.__cause__, ImportError) -def test_discovery_rejects_invalid_provider_type() -> None: - entry_point = _FakeEntryPoint("a", _PluginA) - - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - pytest.raises( - PluginLoadError, - match="must resolve to DurableInstrumentationPluginProvider", - ), - ): - load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) - - -def test_discovery_rejects_incompatible_plugin_api_version() -> None: - entry_point = _FakeEntryPoint( - "a", - _provider(_PluginA, plugin_api_version=99), - ) - - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - pytest.raises(PluginLoadError) as error, - ): - load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) - - assert "declares plugin API version 99" in str(error.value) - assert ( - f"supports plugin API version {DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION}" - in str(error.value) - ) - - -def test_discovery_rejects_invalid_declared_plugin_type() -> None: - provider = DurableInstrumentationPluginProvider( - plugin_type=cast(type[DurableInstrumentationPlugin], object), - factory=_PluginA, - plugin_api_version=DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - ) - entry_point = _FakeEntryPoint("a", provider) - - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - pytest.raises( - PluginLoadError, - match="declares invalid plugin type builtins.object", - ), - ): - load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) - - -def test_discovery_rejects_non_class_declared_plugin_type() -> None: - provider = DurableInstrumentationPluginProvider( - plugin_type=cast(type[DurableInstrumentationPlugin], _PluginA()), - factory=_PluginA, - plugin_api_version=DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - ) - entry_point = _FakeEntryPoint("a", provider) - - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - pytest.raises( - PluginLoadError, - match="declares invalid plugin type .*_PluginA", - ), - ): - load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) - - -def test_discovery_wraps_plugin_factory_failure() -> None: - def fail_factory() -> _PluginA: - raise RuntimeError("factory failed") - +def test_discovery_names_unknown_distribution_in_load_failure() -> None: entry_point = _FakeEntryPoint( "a", - _provider(fail_factory), + _plugin_a_factory, distribution_name=None, + load_error=ImportError("missing dependency"), ) with ( @@ -389,91 +371,71 @@ def fail_factory() -> _PluginA: "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", return_value=[entry_point], ), - pytest.raises(PluginLoadError) as error, + pytest.raises(PluginLoadError, match="unknown distribution"), ): load_configured_plugins( None, environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "Failed to create durable instrumentation plugin 'a'" in str(error.value) - assert "unknown distribution" in str(error.value) - assert isinstance(error.value.__cause__, RuntimeError) +@pytest.mark.parametrize( + ("resolved_value", "expected_type_name"), + [ + (_PluginA(), "_PluginA"), + (object(), "builtins.object"), + ("not-a-factory", "builtins.str"), + (None, "builtins.NoneType"), + ], +) +def test_discovery_rejects_non_callable_entry_point( + resolved_value: object, + expected_type_name: str, +) -> None: + """A plugin *instance* at the entry point is now the common mistake. -def test_discovery_rejects_invalid_plugin_type() -> None: - entry_point = _FakeEntryPoint("a", _provider(lambda: object())) + The old shape resolved to a provider object, so this replaces the + provider-type check with the only check that still means something: the + resolved value has to be callable. The message names what it actually was. + """ + entry_point = _FakeEntryPoint("a", resolved_value) with ( patch( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", return_value=[entry_point], ), - pytest.raises( - PluginLoadError, - match="expected .*_PluginA", - ), + pytest.raises(PluginLoadError) as error, ): load_configured_plugins( None, environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) + assert "must resolve to a callable plugin factory, but resolved to" in str( + error.value + ) + assert expected_type_name in str(error.value) -def test_explicit_plugin_registration_takes_precedence( - caplog: pytest.LogCaptureFixture, -) -> None: - explicit_plugin = _PluginA() - factory = Mock(return_value=_PluginA()) - entry_point = _FakeEntryPoint("a", _provider(factory)) - - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - caplog.at_level( - logging.WARNING, - logger="aws_durable_execution_sdk_python.plugin_discovery", - ), - ): - result = load_configured_plugins( - [explicit_plugin], - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) - assert result == [explicit_plugin] - factory.assert_not_called() - assert "already registered by the decorator's plugins argument" in caplog.text +def test_discovery_accepts_a_plugin_class_as_factory() -> None: + """A class taking the info is callable, so it is a factory in its own right.""" + class _InfoAwarePlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info -def test_first_dynamic_registration_wins_for_duplicate_plugin_type( - caplog: pytest.LogCaptureFixture, -) -> None: - first_factory = Mock(return_value=_PluginA()) - second_factory = Mock(return_value=_PluginA()) - entry_points = [ - _FakeEntryPoint("first", _provider(first_factory)), - _FakeEntryPoint("second", _provider(second_factory)), - ] + entry_point = _FakeEntryPoint("a", _InfoAwarePlugin) - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=entry_points, - ), - caplog.at_level( - logging.WARNING, - logger="aws_durable_execution_sdk_python.plugin_discovery", - ), + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], ): result = load_configured_plugins( None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "first,second"}, + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert len(result) == 1 - assert isinstance(result[0], _PluginA) - first_factory.assert_called_once_with() - second_factory.assert_not_called() - assert "already registered by dynamic provider 'first'" in caplog.text + plugin = result[0](INVOCATION_START_INFO) + assert isinstance(plugin, _InfoAwarePlugin) + assert plugin.info is INVOCATION_START_INFO diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 69cdfc50..77cb54dd 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,7 +1,9 @@ +import contextlib import datetime import logging import pickle import unittest +from collections.abc import Iterator from copy import deepcopy from dataclasses import asdict, fields from unittest.mock import MagicMock, patch @@ -34,6 +36,7 @@ UserFunctionOutcome, UserFunctionStartInfo, ) +from tests.test_helpers import plugin_factory # region Dataclass Tests @@ -44,6 +47,35 @@ LAMBDA_CTX = MagicMock() LAMBDA_CTX.aws_request_id = "req-1" + +@contextlib.contextmanager +def _invocation( + executor: PluginExecutor, + *recorders: "_TrackingPlugin", +) -> Iterator[None]: + """Open the executor's per-invocation scope around a single-hook test. + + Plugin instances are built by ``on_invocation_start`` and dropped when the + scope exits, so a hook dispatched outside an invocation reaches no plugin at + all. A test that exercises one hook therefore has to establish the invocation + that hook belongs to. + + The invocation-start hook this fires is scaffolding for the hook under test, + so it is cleared from the given recorders before the body runs. The + assertions that follow are about the hook under test only. + """ + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + ) + for recorder in recorders: + recorder.calls.clear() + yield + + OPERATION_START_INFO = OperationStartInfo( operation_id="op-2", operation_type=OperationType.CALLBACK, @@ -528,17 +560,213 @@ def test_subclass_override(self): class TestPluginExecutorInit(unittest.TestCase): def test_init_with_none(self): executor = PluginExecutor(plugins=None) + self.assertEqual(executor._plugin_factories, []) self.assertEqual(executor._plugins, []) def test_init_with_empty_list(self): executor = PluginExecutor(plugins=[]) + self.assertEqual(executor._plugin_factories, []) self.assertEqual(executor._plugins, []) - def test_init_with_plugins(self): + def test_init_records_factories_without_building_plugins(self): + """Construction stores factories only; instances belong to an invocation.""" p1 = _NoOpPlugin() p2 = _TrackingPlugin() - executor = PluginExecutor(plugins=[p1, p2]) - self.assertEqual(len(executor._plugins), 2) + executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)]) + self.assertEqual(len(executor._plugin_factories), 2) + self.assertEqual(executor._plugins, []) + + +class TestPluginLifetime(unittest.TestCase): + """The per-invocation plugin lifetime.""" + + @staticmethod + def _run_invocation(executor: PluginExecutor, request_id: str) -> None: + lambda_context = MagicMock() + lambda_context.aws_request_id = request_id + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=lambda_context, + execution_start_time=START_TS, + is_first_invocation=False, + ) + + def test_each_invocation_gets_its_own_instance(self): + """Two invocations of one handler never share a plugin instance.""" + built: list[_TrackingPlugin] = [] + + def build(info: InvocationStartInfo) -> _TrackingPlugin: + plugin = _TrackingPlugin() + built.append(plugin) + return plugin + + executor = PluginExecutor(plugins=[build]) + + self._run_invocation(executor, "req-1") + self._run_invocation(executor, "req-2") + + self.assertEqual(len(built), 2) + self.assertIsNot(built[0], built[1]) + # Each instance saw only its own invocation. + self.assertEqual(built[0].calls, ["invocation_start:req-1"]) + self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + + def test_instances_are_dropped_when_the_invocation_returns(self): + """Nothing on the handler-lifetime executor still references the instance.""" + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + ) + self.assertEqual(executor._plugins, [plugin]) + + self.assertEqual(executor._plugins, []) + + def test_factory_receives_the_info_the_first_hook_receives(self): + """The factory argument is the identical object, not a copy.""" + factory_infos: list[InvocationStartInfo] = [] + hook_infos: list[InvocationStartInfo] = [] + + class _RecordingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + hook_infos.append(info) + + def build(info: InvocationStartInfo) -> _RecordingPlugin: + factory_infos.append(info) + return _RecordingPlugin() + + executor = PluginExecutor(plugins=[build]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + execution_input={"name": "World"}, + ) + + self.assertEqual(len(factory_infos), 1) + self.assertEqual(len(hook_infos), 1) + self.assertIs(factory_infos[0], hook_infos[0]) + self.assertEqual(factory_infos[0].execution_arn, "arn:exec") + self.assertEqual(factory_infos[0].execution_input, {"name": "World"}) + + def test_factories_run_before_the_first_hook_is_dispatched(self): + """Every instance exists before any of them receives a hook.""" + events: list[str] = [] + + class _OrderedPlugin(DurableInstrumentationPlugin): + def __init__(self, label: str) -> None: + self.label = label + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + events.append(f"hook:{self.label}") + + def build(label: str): + def factory(info: InvocationStartInfo) -> _OrderedPlugin: + events.append(f"build:{label}") + return _OrderedPlugin(label) + + return factory + + executor = PluginExecutor(plugins=[build("a"), build("b")]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertEqual(events, ["build:a", "build:b", "hook:a", "hook:b"]) + + def test_failing_factory_is_contained(self): + """A raising factory is logged and skipped, like a raising hook.""" + surviving = _TrackingPlugin() + + def exploding(info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + executor = PluginExecutor( + plugins=[exploding, plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("factory boom", "\n".join(logs.output)) + # The other factory's plugin still receives its hooks. + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_factory_returning_none_is_contained(self): + """A factory that returns nothing is logged and skipped.""" + surviving = _TrackingPlugin() + + def returns_none(info: InvocationStartInfo): + return None + + executor = PluginExecutor( + plugins=[returns_none, plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + self.assertIn("returned None", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_every_failing_factory_leaves_the_executor_usable(self): + """All factories failing is not distinguishable from having no plugins.""" + + def exploding(info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + executor = PluginExecutor(plugins=[exploding]) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ): + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, []) + # Later hooks are no-ops rather than errors. + executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=ServiceInvocationStatus.SUCCEEDED, + result=None, + error=None, + ), + ) class TestPluginExecutor(unittest.TestCase): @@ -552,7 +780,7 @@ def test_no_thread_pool_when_plugins_is_empty_list(self): self.assertIsNone(executor._executor) def test_thread_pool_created_when_plugins_provided(self): - executor = PluginExecutor(plugins=[_NoOpPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_NoOpPlugin())]) with executor.run(): self.assertIsNotNone(executor._executor) @@ -629,40 +857,40 @@ class TestPluginExecutorExecutePlugins(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_dispatch_invocation_start_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(INVOCATION_START_INFO, sync=True) self.assertIn("invocation_start:req-1", self.plugin.calls) def test_dispatch_invocation_end_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(INVOCATION_END_INFO, sync=True) self.assertIn("invocation_end:req-1", self.plugin.calls) def test_dispatch_operation_end_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(OPERATION_END_INFO, sync=False) self.assertIn("operation_end:op-1", self.plugin.calls) def test_dispatch_operation_start_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(OPERATION_START_INFO, sync=False) self.assertIn("operation_start:op-2", self.plugin.calls) def test_dispatch_operation_change_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(OPERATION_CHANGE_INFO, sync=False) self.assertIn("operation_change:op-1", self.plugin.calls) def test_dispatch_user_function_start_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(USER_FUNCTION_START_INFO, sync=True) self.assertIn("user_function_start:op-1", self.plugin.calls) def test_dispatch_user_function_end_info(self): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins(USER_FUNCTION_END_INFO, sync=True) self.assertIn("user_function_end:op-1", self.plugin.calls) @@ -671,19 +899,21 @@ def test_dispatch_unknown_type_logs_exception(self): with self.assertLogs( "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.execute_plugins("not a valid info type", sync=True) def test_plugin_exception_is_swallowed(self): """If a plugin raises, the exception is logged and execution continues.""" failing_plugin = _FailingPlugin() tracking_plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[failing_plugin, tracking_plugin]) + executor = PluginExecutor( + plugins=[plugin_factory(failing_plugin), plugin_factory(tracking_plugin)] + ) with self.assertLogs( "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): - with executor.run(): + with _invocation(executor, tracking_plugin): executor.execute_plugins(OPERATION_START_INFO, sync=True) # The second plugin should still have been called @@ -692,9 +922,9 @@ def test_plugin_exception_is_swallowed(self): def test_multiple_plugins_all_called(self): p1 = _TrackingPlugin() p2 = _TrackingPlugin() - executor = PluginExecutor(plugins=[p1, p2]) + executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)]) - with executor.run(): + with _invocation(executor, p1, p2): executor.execute_plugins(OPERATION_START_INFO, sync=True) self.assertIn("operation_start:op-2", p1.calls) @@ -706,7 +936,7 @@ class TestPluginExecutorOnInvocationStart(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) self.ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) def _make_operation(self, start_time=None): @@ -780,7 +1010,7 @@ class TestPluginExecutorOnInvocationEnd(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) self.ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) def _make_operation(self, start_ts=None, end_ts=None): @@ -857,7 +1087,7 @@ def on_invocation_start(_self, info): # noqa: N805 def on_invocation_end(_self, info): # noqa: N805 self.captured.append(info) - self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + self.executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) @staticmethod def _operation(operation_id, status=OperationStatus.SUCCEEDED): @@ -1030,7 +1260,7 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_invocation_start(_self, info): # noqa: N805 seen.append(info) - executor = PluginExecutor(plugins=[_CapturingPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) with executor.run(): executor.on_invocation_start( execution_arn="arn:exec", @@ -1209,7 +1439,7 @@ def on_invocation_start(_self, info): # noqa: N805 def on_invocation_end(_self, info): # noqa: N805 self.captured.append(info) - self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + self.executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) def _fire_hooks(self): with self.executor.run(): @@ -1281,7 +1511,7 @@ class TestPluginExecutorOnOperationAction(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_start_action_fires_operation_start(self): captured: list[OperationStartInfo] = [] @@ -1292,7 +1522,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1301,7 +1531,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: update.name = "my-step" update.parent_id = "parent-1" - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action(update) self.assertIn("operation_start:op-1", self.plugin.calls) @@ -1318,7 +1548,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1334,7 +1564,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: start_timestamp=START_TS, ) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action(update, operation) self.assertEqual(captured[0].start_time, START_TS) @@ -1348,7 +1578,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1368,7 +1598,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: status=OperationStatus.READY, ) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action( update, operation=current_operation, @@ -1411,28 +1641,28 @@ def test_terminal_operation_does_not_fire_callbacks(self): for status in terminal_statuses: with self.subTest(status=status): plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.STEP, status=status, ) - with executor.run(): + with _invocation(executor, plugin): executor.on_operation_replay(operation) self.assertEqual(plugin.calls, []) def test_non_terminal_operation_fires_operation_start(self): plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.WAIT, status=OperationStatus.STARTED, ) - with executor.run(): + with _invocation(executor, plugin): executor.on_operation_replay(operation) self.assertEqual(plugin.calls, ["operation_start:op-1"]) @@ -1450,7 +1680,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: captured.append(info) plugin = _CapturingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) identifier = OperationIdentifier( operation_id="context-1", sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, @@ -1459,7 +1689,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) before = datetime.datetime.now(datetime.UTC) - with executor.run(): + with _invocation(executor, plugin): executor.on_child_context_end( identifier, OperationStatus.FAILED, @@ -1487,13 +1717,13 @@ def on_operation_end(self, info: OperationEndInfo) -> None: class TestPluginExecutorOnUserFunction(unittest.TestCase): def test_user_function_info_uses_plugin_operation_type(self): - executor = PluginExecutor(plugins=[_TrackingPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_TrackingPlugin())]) identifier = OperationIdentifier( operation_id="step-1", sub_type=OperationSubType.STEP, ) - with executor.run(): + with _invocation(executor): info = executor.on_user_function_start(identifier) self.assertIs(info.operation_type, OperationType.STEP) @@ -1504,7 +1734,7 @@ class TestPluginExecutorOnOperationUpdate(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def _make_operation( self, @@ -1532,7 +1762,7 @@ def _make_operation( def test_terminal_status_without_step_details_fires_operation_only(self): op = self._make_operation(status=OperationStatus.FAILED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1540,7 +1770,7 @@ def test_terminal_status_without_step_details_fires_operation_only(self): def test_non_terminal_status_without_step_details_fires_nothing(self): op = self._make_operation(status=OperationStatus.STARTED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertEqual(self.plugin.calls, []) @@ -1548,7 +1778,7 @@ def test_non_terminal_status_without_step_details_fires_nothing(self): def test_ready_status_fires_nothing(self): op = self._make_operation(status=OperationStatus.READY, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertEqual(self.plugin.calls, []) @@ -1556,7 +1786,7 @@ def test_ready_status_fires_nothing(self): def test_timed_out_is_terminal(self): op = self._make_operation(status=OperationStatus.TIMED_OUT, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1564,7 +1794,7 @@ def test_timed_out_is_terminal(self): def test_cancelled_is_terminal(self): op = self._make_operation(status=OperationStatus.CANCELLED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1572,7 +1802,7 @@ def test_cancelled_is_terminal(self): def test_stopped_is_terminal(self): op = self._make_operation(status=OperationStatus.STOPPED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1583,7 +1813,7 @@ class TestPluginExecutorOnOperationChange(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_operation_change_uses_invocation_and_operation_maps(self): updated_operation = Operation( @@ -1612,7 +1842,7 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) with self.executor.run(): self.executor.on_invocation_start( @@ -1644,11 +1874,19 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: self.assertEqual(updated_info.end_time, END_TS) self.assertFalse(updated_info.is_replayed) - def test_operation_change_without_invocation_start_is_noop(self): + def test_hooks_outside_an_invocation_reach_no_plugin(self): + """Nothing is dispatched before an invocation establishes the instances. + + This used to assert the ``_invocation_status is None`` guard inside + ``on_operation_update``. That guard is now redundant with the plugin + lifetime itself: instances are built by ``on_invocation_start`` and + dropped when the invocation scope exits, so outside an invocation there + is no instance to dispatch to in the first place. + """ operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.STEP, - status=OperationStatus.STARTED, + status=OperationStatus.SUCCEEDED, ) with self.executor.run(): @@ -1658,6 +1896,7 @@ def test_operation_change_without_invocation_start_is_noop(self): previous_operations={}, ) + self.assertEqual(self.executor._plugins, []) self.assertEqual(self.plugin.calls, []) diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 92b07f58..b2c408fb 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -57,6 +57,7 @@ QueuedOperation, ) from aws_durable_execution_sdk_python.threading import CompletionEvent +from tests.test_helpers import plugin_factory, plugin_invocation def test_checkpointed_result_create_from_operation_step(): @@ -4531,7 +4532,7 @@ def test_execution_state_accepts_plugin_executor_parameter(): """Test that ExecutionState can be created with a plugin_executor parameter.""" mock_client = Mock(spec=LambdaClient) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) state = ExecutionState( durable_execution_arn="test_arn", @@ -4564,8 +4565,8 @@ def test_plugin_executor_on_operation_action_called_on_checkpoint(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4598,8 +4599,8 @@ def test_plugin_executor_on_operation_action_called_on_checkpoint(): def test_async_operation_start_precedes_user_function_start(): """Async START notifies plugins before the user function begins.""" plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4619,7 +4620,12 @@ def test_async_operation_start_precedes_user_function_start(): ) plugin_executor.on_user_function_start(operation_identifier, attempt=1) - assert plugin.calls[:2] == [ + # Drop the per-invocation instance's own invocation-start hook, which now + # always precedes the operation hooks under test. + operation_calls = [ + call for call in plugin.calls if not call.startswith("invocation_") + ] + assert operation_calls[:2] == [ f"operation_start:{operation_id}", f"user_function_start:{operation_id}", ] @@ -4649,8 +4655,8 @@ def test_existing_operation_start_is_reported_as_replayed(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4697,8 +4703,8 @@ def test_plugin_executor_on_operation_update_called_for_terminal_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4747,8 +4753,8 @@ def test_plugin_executor_not_called_for_non_terminal_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4805,8 +4811,8 @@ def test_plugin_executor_called_for_multiple_updates_in_batch(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): config = CheckpointBatcherConfig( max_batch_time_seconds=0.2, max_batch_operations=10, @@ -4877,7 +4883,7 @@ def test_plugin_executor_on_operation_change_called_for_status_changes(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) with plugin_executor.run(): plugin_executor.on_invocation_start( execution_arn="test_arn", @@ -4926,8 +4932,8 @@ def test_operation_start_plugin_hook_fires_before_checkpoint_failure(): mock_client.checkpoint.side_effect = RuntimeError("API error") plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4984,8 +4990,8 @@ def on_operation_end(self, info): raise RuntimeError("plugin exploded") exploding_plugin = _ExplodingPlugin() - plugin_executor = PluginExecutor(plugins=[exploding_plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(exploding_plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5027,8 +5033,8 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_user_function_end(self, info: UserFunctionEndInfo) -> None: captured.append(info) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5082,8 +5088,8 @@ def test_plugin_executor_not_called_for_pending_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5149,8 +5155,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5184,8 +5190,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5207,8 +5213,8 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_operation_start(self, info): captured.append(info.operation_id) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5328,7 +5334,7 @@ def _wrapping_state(plugin: _RecordingPlugin) -> ExecutionState: initial_checkpoint_token="token123", # noqa: S106 operations={}, service_client=Mock(spec=LambdaClient), - plugin_executor=PluginExecutor(plugins=[plugin]), + plugin_executor=PluginExecutor(plugins=[plugin_factory(plugin)]), ) @@ -5362,10 +5368,13 @@ def test_wrap_user_function_reports_incomplete_when_no_outcome(raised): def user_function(): raise raised - with state._plugin_executor.run(), pytest.raises(type(raised)): + with plugin_invocation(state._plugin_executor), pytest.raises(type(raised)): _wrapped(state, user_function)() assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5380,10 +5389,13 @@ def test_wrap_user_function_does_not_report_incomplete_on_success(): plugin = _RecordingPlugin() state = _wrapping_state(plugin) - with state._plugin_executor.run(): + with plugin_invocation(state._plugin_executor): assert _wrapped(state, lambda: "done")() == "done" assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5400,10 +5412,16 @@ def test_wrap_user_function_does_not_report_incomplete_on_failure(): def user_function(): raise ValueError("boom") - with state._plugin_executor.run(), pytest.raises(ValueError, match="boom"): + with ( + plugin_invocation(state._plugin_executor), + pytest.raises(ValueError, match="boom"), + ): _wrapped(state, user_function)() assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5427,7 +5445,7 @@ def on_user_function_end(self, info) -> None: initial_checkpoint_token="token123", # noqa: S106 operations={}, service_client=Mock(spec=LambdaClient), - plugin_executor=PluginExecutor(plugins=[plugin]), + plugin_executor=PluginExecutor(plugins=[plugin_factory(plugin)]), ) def user_function(): @@ -5440,7 +5458,10 @@ def run_on_worker() -> None: with contextlib.suppress(SuspendExecution): _wrapped(state, user_function)() - with state._plugin_executor.run(), ThreadPoolExecutor(max_workers=1) as worker: + with ( + plugin_invocation(state._plugin_executor), + ThreadPoolExecutor(max_workers=1) as worker, + ): worker.submit(run_on_worker).result() assert hook_threads == worker_threads diff --git a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py index 77611a34..de677280 100644 --- a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py +++ b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py @@ -1,9 +1,16 @@ """Test helpers for generating expected step IDs.""" +import contextlib +from collections.abc import Iterator from unittest.mock import Mock from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext from aws_durable_execution_sdk_python.execution import ExecutionState +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, + PluginExecutor, +) def operation_id_sequence(parent_id: str | None = None): @@ -18,3 +25,44 @@ def operation_id_sequence(parent_id: str | None = None): while True: yield context._create_step_id() # noqa: SLF001 + + +def plugin_factory( + plugin: DurableInstrumentationPlugin, +) -> DurableInstrumentationPluginFactory: + """Wrap a plugin instance a test already holds a reference to as a factory. + + Production factories build a fresh instance per invocation. A test that has + to read what the plugin recorded needs the instance it passed in, so it + supplies a factory that returns that one. Only valid for a single + invocation, which is all these tests run. + """ + return lambda info: plugin + + +@contextlib.contextmanager +def plugin_invocation( + plugin_executor: PluginExecutor, + *, + execution_arn: str = "test_arn", + is_first_invocation: bool = True, +) -> Iterator[None]: + """Open a plugin executor's per-invocation scope. + + Plugin instances are built by ``on_invocation_start`` and dropped when the + ``run()`` scope exits, so a hook dispatched outside an invocation reaches no + plugin at all. Tests that exercise a single hook still have to establish the + invocation the hook belongs to; this does that and nothing else. + + The invocation-start hook this fires is scaffolding. Tests asserting on an + exact list of recorded calls should record only the hooks they care about, or + clear their recorder after entering the scope. + """ + with plugin_executor.run(): + plugin_executor.on_invocation_start( + execution_arn=execution_arn, + is_first_invocation=is_first_invocation, + execution_start_time=None, + lambda_context=None, + ) + yield From 67eabd08b11125a0b7a226400e486a70ad6d739b Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 17 Sep 2026 12:15:43 -0700 Subject: [PATCH 03/28] fix(plugin): give each invocation its own plugin session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings on the per-invocation plugin contract. The first one meant the contract change did not actually deliver what it promised under concurrency. `durable_execution` built one `PluginExecutor` at decoration time and every invocation shared it, while the executor stored that invocation's plugin instances, invocation metadata and operations provider on itself. So the SDK created one plugin per invocation and then parked it in a shared mutable slot. Measured with two invocations held inside their user function at once, invocation A's operation hook was delivered to B's plugin and A's plugin never received its own operation or end hook; in a second scenario A's scope exit cleared the shared state while B was live and B lost `on_invocation_end` entirely. A new handler-lifetime `PluginHost` now holds nothing but the resolved factory list and hands out a fresh `PluginExecutor` per invocation, which is the shape JS and Java already had. The executor is reachable only from the frames of the invocation that owns it, and a second `run()` on one executor raises, so reuse fails loudly instead of as silent crosstalk. Its thread pool is per invocation too: sharing one would let an ending invocation shut down a pool a concurrent invocation was still submitting to. The OTel log filter had the same class of bug. It held one mutable plugin reference, so with two invocations open the one that started last won and every other invocation's records carried its trace. Reading the active span from the OTel context, as the Java plugin does, is not sufficient here: the invocation span is never attached to the context, and the SDK runs the handler body on a pool thread that does not inherit the context of the thread the plugin bound on. The filter is now stateless and resolves per record, preferring the invocation that claimed the emitting thread through a `ContextVar` and falling back to the single open invocation when exactly one is open. With several open and an unclaimed thread it leaves the record unstamped, because an unattributed record is a smaller defect than one attributed to another customer's execution. Installation is serialized so concurrent first- time callers cannot stack filters. Third, the Insight scheduler used `0` for both "no flush in flight" and "a flush in flight covering zero exports", which are different facts. An invocation that emitted no record could not tell that its flush was already running, so it requested a second one that then ran after both invocations had returned — customer exporter code past the invocation boundary, which the flush contract forbids. Absence is now `None` and coverage an integer. --- .../_export_scheduler.py | 43 ++- .../tests/test_export_scheduler.py | 95 +++++- .../README.md | 5 +- .../execution_plugin.py | 37 ++- .../invocation_plugin.py | 34 +- .../log_filter.py | 189 ++++++++--- .../tests/test_invocation_plugin.py | 2 +- .../tests/test_log_filter.py | 296 +++++++++++++++--- .../execution.py | 13 +- .../plugin.py | 104 +++++- .../aws_durable_execution_sdk_python/state.py | 3 + .../tests/execution_test.py | 121 ++++++- .../tests/plugin_test.py | 52 ++- 13 files changed, 842 insertions(+), 152 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index fdd92db7..06a52b67 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -101,11 +101,17 @@ def __init__(self, exporters: list[InsightExporter]) -> None: self._flushes_completed = 0 self._flush_requested = False # Export counter coverage of the flush the worker is running right now, or - # 0 when no flush is in flight. Published when the worker commits to a - # flush, so a waiter woken while that flush runs -- before its coverage - # reaches _flushed_through -- can tell it is already covered instead of - # requesting a second flush that would run after its drain returned. - self._flush_in_flight = 0 + # None when no flush is in flight. Presence and coverage are separate + # facts: a flush that covers zero exports is an ordinary flush -- it is + # what an invocation that emitted no record asks for -- and a single + # integer cannot say both "no flush is running" and "a flush covering + # nothing is running". Encoding the first as 0 made those two states + # identical, so a waiter needing zero coverage could not tell that its + # flush was already running and requested a second one that then ran + # after its invocation had returned. Published when the worker commits to + # a flush, so a waiter woken while that flush runs -- before its coverage + # reaches _flushed_through -- can tell it is already covered. + self._flush_in_flight: int | None = None # Value of the global schedule counter (_seq) when a flush was requested. # The worker defers the flush until no record scheduled at or before that # point is still pending. That is deliberately wider than the requester's @@ -202,15 +208,15 @@ def drain(self, execution: _ExportState) -> None: # made has already been consumed -- runs an extra flush after # this drain, and the invocation, returned. # - # `_flush_in_flight` uses 0 as its "no flush is running" - # sentinel, so the naive `self._flush_in_flight >= need` - # reads as "already covered" when `need` is 0 -- precisely - # when nothing is running at all. `need` is 0 for a drain - # whose invocation emitted no record, so that form would let - # such a drain skip its request and park until some other - # execution happened to flush. Require a marker that is - # actually set AND that reaches `need`. - covered = 0 < self._flush_in_flight >= need + # `_flush_in_flight` is None exactly while no flush is + # running, so a flush that covers zero exports is still a + # flush in flight. That case is the common one, not a corner: + # `need` is 0 for a drain whose invocation emitted no record, + # and the flush it asks for covers 0 exports when nothing has + # ever been exported. Two such drains at once both see the + # other's flush and neither asks for a second. + in_flight = self._flush_in_flight + covered = in_flight is not None and in_flight >= need if not self._flush_requested and not covered: self._flush_requested = True self._flush_barrier = max(self._flush_barrier, self._seq) @@ -260,7 +266,7 @@ def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: self._pending = {} self._flush_requested = False self._flush_barrier = 0 - self._flush_in_flight = 0 + self._flush_in_flight = None # Release every waiter; the permanent disable latch means no record # will ever be exported. self._condition.notify_all() @@ -304,7 +310,10 @@ def _run_loop(self) -> None: flush_covers = self._export_count # Publish what this flush will cover before releasing the # lock, so a waiter that wakes while it runs can see that - # this flush releases it and skip asking for another. + # this flush releases it and skip asking for another. A + # coverage of 0 is published like any other: it means this + # flush covers every export so far, of which there are + # none, and it is still a flush in flight. self._flush_in_flight = flush_covers break if self._pending: @@ -353,7 +362,7 @@ def _run_loop(self) -> None: # Retire the marker whatever happened: a stale one would park # every later waiter that trusted this flush to cover it. Only # a flush that ran to completion publishes its coverage. - self._flush_in_flight = 0 + self._flush_in_flight = None if flushed: self._flushes_completed += 1 if flush_covers > self._flushed_through: diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index cbca9d4f..477a18b1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -430,6 +430,92 @@ def drain() -> None: assert capture.calls == [("flush", None)] +class GatedFlushExporter(CaptureExporter): + """Holds each ``flush()`` open until the test releases it, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self._lock = threading.Lock() + self.flush_count = 0 + self.started: dict[int, threading.Event] = {} + self.release: dict[int, threading.Event] = {} + for index in (1, 2): + self.started[index] = threading.Event() + self.release[index] = threading.Event() + + def flush(self) -> None: + with self._lock: + self.flush_count += 1 + index = self.flush_count + if index in self.started: + self.started[index].set() + self.release[index].wait(10.0) + super().flush() + + +def test_concurrent_drains_with_nothing_to_export_share_one_flush() -> None: + # Two invocations that emitted no record drain at the same time. Neither has + # an export to be covered, so both need a flush that covers zero exports. + # While `_flush_in_flight` used 0 for "no flush is running", the second drain + # could not tell that the flush it needs was already running, and asked for + # another. The completion of the first flush then released both drains, and + # the second flush ran after both invocations had already returned -- customer + # exporter code running past the invocation boundary, which is what the flush + # contract forbids. One flush must serve both, and whichever drain a flush + # belongs to must stay parked until that flush completes. + exporter = GatedFlushExporter() + scheduler = _ArnScheduler([exporter]) + returned: list[str] = [] + returned_lock = threading.Lock() + + def drain(name: str, execution_arn: str) -> None: + scheduler.drain(execution_arn) + with returned_lock: + returned.append(name) + + threads = [ + threading.Thread(target=drain, args=("first", ARN_A), daemon=True), + threading.Thread(target=drain, args=("second", ARN_B), daemon=True), + ] + try: + threads[0].start() + assert exporter.started[1].wait(10.0), "the first drain never flushed" + threads[1].start() + + first = scheduler.executions[ARN_A] + + def both_parked() -> bool: + second = scheduler.executions.get(ARN_B) + if second is None: + return False + with scheduler._condition: + return first.waiters == 1 and second.waiters == 1 + + assert _wait_until(both_parked), "a drain raced past the flush it needs" + with returned_lock: + assert returned == [], "a drain returned before its flush completed" + + exporter.release[1].set() + for thread in threads: + thread.join(10.0) + assert not any(thread.is_alive() for thread in threads) + with returned_lock: + assert sorted(returned) == ["first", "second"] + + # The redundant request, if one was made, was recorded before either + # drain returned, so the worker starts that flush without further + # prompting. Nothing arriving here is what proves no second flush was + # requested. + assert not exporter.started[2].wait(0.75), ( + "a second flush ran after both invocations had returned" + ) + assert exporter.flush_count == 1 + finally: + exporter.release[1].set() + exporter.release[2].set() + _wait_until(lambda: not scheduler._worker_alive()) + + def test_drain_never_rides_on_a_flush_that_finished_before_it_started() -> None: # Export coverage alone would let the second drain return immediately: every # export is already covered by the first drain's flush. A drain must wait for @@ -475,14 +561,15 @@ def fail_start(self) -> None: # noqa: ARG001 for execution in scheduler.executions.values() ) assert scheduler._flush_requested is False - assert scheduler._flush_in_flight == 0 + assert scheduler._flush_in_flight is None def test_disabled_latch_clears_a_published_flush_in_flight_marker(monkeypatch) -> None: # `_flush_in_flight` is the coverage of the flush the worker is running right # now, published so a waiter woken during that flush can tell it is already - # covered and skip requesting another. 0 means "no flush is running", so the - # marker is a claim that a flush is in flight and will complete. + # covered and skip requesting another. None means "no flush is running", so + # any integer -- 0 included -- is a claim that a flush is in flight and will + # complete. # # The _disabled latch makes that claim permanently false: no worker exists and # none will ever be started again, so the published flush can never complete. @@ -503,7 +590,7 @@ def fail_start(self) -> None: # noqa: ARG001 with scheduler._condition: assert scheduler._disabled - assert scheduler._flush_in_flight == 0, ( + assert scheduler._flush_in_flight is None, ( "the _disabled latch left a flush-in-flight marker behind for a flush " "that can never run" ) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index f532ed35..4cbe4850 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -430,7 +430,10 @@ Structured trace context and sampling decision returned by context extractors. The logging filter (and its installer) used to stamp trace context onto log records. Installed automatically when `enrich_logger=True`; exported for manual -setups. +setups, where `install_log_filter(target_logger)` attaches it to a logger of your +choice. The filter carries no invocation identity of its own: it resolves the +invocation a record belongs to at emit time, so one filter serves every +invocation the environment runs, including concurrent ones. ## Requirements 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 c522278f..7af7e452 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 @@ -90,7 +90,11 @@ canonical_trace_id, ) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig -from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter +from aws_durable_execution_sdk_python_otel.log_filter import ( + bind_invocation, + install_log_filter, + unbind_invocation, +) from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -177,9 +181,12 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._tracing_enabled = False if self._config.enrich_logger: - # Install (or, on a warm environment, rebind) the root-logger filter - # so every log record is stamped with this invocation's span context. - install_log_filter(self) + # Install the root-logger filter so every log record is stamped with + # the active span context. On a warm environment the handler already + # carries the filter a previous invocation installed and it is reused + # as is: the filter holds no invocation identity, and this plugin + # claims the invocation in on_invocation_start instead. + install_log_filter() def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" @@ -426,6 +433,12 @@ def _with_sampling(self, parent_context: Context) -> Context: # ------------------------------------------------------------------ def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) + # Claim log correlation for this invocation before anything can fail + # below: the claim is what keeps a concurrent invocation's records off + # this invocation's trace, and it is registered even when tracing turns + # out to be disabled so that the filter can still tell how many + # invocations are open. + bind_invocation(self) if info.execution_start_time is None: logger.warning( "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -660,12 +673,17 @@ def _release_invocation_scope(self) -> None: plugin, so any scope this plugin attached and did not release must be detached here or it would stay current on a warm environment's thread after the invocation returns. + * The log filter's record of open invocations is process-global, so this + invocation must be removed from it. Until it is, a record emitted on an + unclaimed thread could still be 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. """ self._detach_remaining_contexts() + unbind_invocation(self) self._tracing_enabled = False # ------------------------------------------------------------------ @@ -675,6 +693,14 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return + # Runs on the thread that drives the durable operation, which is the + # thread running the handler body and not the thread the + # invocation-start hook claimed. Claim it too, so records emitted from + # top-level handler code are correlated to this invocation even while + # another invocation is open in the same process. Claimed after the + # tracing-enabled gate, so a hook arriving after the invocation ended + # cannot re-register a finished invocation. + bind_invocation(self) if info.operation_type is OperationType.CONTEXT: with self._lock: self._checkpointed_context_ids.add(info.operation_id) @@ -807,6 +833,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: logger.debug("Durable user function started: %s", info) if not self._tracing_enabled: return + # Runs on the thread executing user code -- a parallel branch runs on its + # own thread -- so claim that thread for this invocation as well. + bind_invocation(self) if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP): raise RuntimeError( "on_user_function_start only supports CONTEXT and STEP operations" 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 96ecda36..b0043c6f 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 @@ -58,7 +58,11 @@ ExecutionTraceContext, canonical_trace_id, ) -from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter +from aws_durable_execution_sdk_python_otel.log_filter import ( + bind_invocation, + install_log_filter, + unbind_invocation, +) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -180,8 +184,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # handler before the handler module is imported (and thus before the # plugin is constructed), so the handlers are available here. On a # warm environment the handler already carries the filter a previous - # invocation's plugin installed, and this rebinds it to this one. - install_log_filter(self) + # invocation installed and is reused as is: the filter holds no + # invocation identity, and this plugin claims the invocation in + # on_invocation_start instead. + install_log_filter() def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" @@ -526,6 +532,12 @@ def _end_span( def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) + # Claim log correlation for this invocation before anything can fail + # below: the claim is what keeps a concurrent invocation's records off + # this invocation's trace, and it is registered even when tracing turns + # out to be disabled so that the filter can still tell how many + # invocations are open. + bind_invocation(self) if info.execution_start_time is None: logger.warning( "InvocationOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -730,12 +742,17 @@ def _release_invocation_scope(self) -> None: plugin, so any scope this plugin attached and did not release must be detached here or it would stay current on a warm environment's thread after the invocation returns. + * The log filter's record of open invocations is process-global, so this + invocation must be removed from it. Until it is, a record emitted on an + unclaimed thread could still be 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. """ self._detach_remaining_contexts() + unbind_invocation(self) self._tracing_enabled = False def on_operation_start(self, info: OperationStartInfo) -> None: @@ -743,6 +760,14 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return + # Runs on the thread that drives the durable operation, which is the + # thread running the handler body and not the thread the + # invocation-start hook claimed. Claim it too, so records emitted from + # top-level handler code are correlated to this invocation even while + # another invocation is open in the same process. Claimed after the + # tracing-enabled gate, so a hook arriving after the invocation ended + # cannot re-register a finished invocation. + bind_invocation(self) if info.operation_type is OperationType.CONTEXT: # The user-function hook owns the span, but this durable START hook # distinguishes checkpoint-backed contexts from virtual branches. @@ -815,6 +840,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: logger.debug("Durable user function started: %s", info) if not self._tracing_enabled: return + # Runs on the thread executing user code -- a parallel branch runs on its + # own thread -- so claim that thread for this invocation as well. + bind_invocation(self) # Context and Step operations are tracked using on_user_function_start if info.operation_type not in [OperationType.CONTEXT, OperationType.STEP]: raise RuntimeError( 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 e1e8d34e..65cf71fe 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 @@ -13,11 +13,51 @@ These attributes are only set when a valid span context is active. Records emitted outside an active invocation (e.g. during Lambda teardown) pass through unmodified, so any log formatter or schema must treat the fields as optional. + +Resolving *which* invocation a record belongs to +----------------------------------------------- + +A logging handler is process-global and outlives every invocation, while a +plugin instance serves exactly one invocation. Concurrent executions in one +environment (Lambda Managed Instances) therefore have several plugin instances +alive at once, all reachable from one installed filter. The filter must not hold +a single mutable reference to "the" plugin: whichever invocation started last +would win, and every other invocation's records would be stamped with its trace. + +So the binding is per invocation, not per filter: + + - ``bind_invocation`` marks an invocation open and claims the calling + thread/task for it, through a :class:`contextvars.ContextVar`. A record + emitted on a claimed thread resolves to the invocation that claimed it, + which is per-thread and per-task and so cannot be overwritten by a + concurrent invocation. + - ``unbind_invocation`` marks the invocation closed. + - A record emitted on a thread no invocation has claimed resolves to the one + open invocation, if exactly one is open. That covers the SDK's user-code + worker threads: the invocation-start hook runs on the Lambda handler + thread, the handler body runs on a pool thread the plugin has not been + given control on yet, and Python does not propagate context into new + threads. + - If several invocations are open and the thread is unclaimed, the record is + left unstamped. An unattributed record is a smaller defect than one + attributed to another customer execution. + +Reading the active span straight from the OTel context (as the Java plugin's +static ``MdcSpanEnricher`` does) is not sufficient here: the invocation span is +never attached to the OTel context, and the context of the handler thread -- the +only place the plugin attaches anything at invocation scope -- is not visible on +the worker thread that runs the handler body. Top-level records would silently +lose correlation. The plugin's ``get_current_span_context()`` still reads the +OTel context first, so records emitted inside a step or child context resolve to +the active operation span exactly as before. """ from __future__ import annotations +import contextvars import logging +import threading +import weakref from typing import TYPE_CHECKING, Protocol from opentelemetry.trace import TraceFlags @@ -38,46 +78,92 @@ class _SpanContextProvider(Protocol): def get_current_span_context(self) -> SpanContext | None: ... -class OtelContextLogFilter(logging.Filter): - """Logging filter that injects the active OTel span context onto records. +# Guards the open-invocation registry. Held for the length of a set membership +# test or a single mutation, never while a plugin is called. +_registry_lock = threading.Lock() - The filter is a pure reader of the plugin's current span context. It - resolves the span at emit time, on the thread that emits the record, via - ``plugin.get_current_span_context()``. That method returns the active - operation span inside steps and child contexts (attached to the worker - thread's OTel context) and falls back to the invocation span for top-level - handler code. +# Invocations that have started and not yet ended. Weak so that a plugin whose +# end hook never ran (a process torn down mid-invocation) cannot keep itself, +# and the spans it holds, alive for the life of the environment. +_open_invocations: weakref.WeakSet[_SpanContextProvider] = weakref.WeakSet() - The filter never caches identifiers and always returns ``True`` so it never - drops a record. +# The invocation owning the current thread/task. Set by bind_invocation on every +# thread the owning plugin is given control on. +_current_invocation: contextvars.ContextVar[_SpanContextProvider | None] = ( + contextvars.ContextVar("durable_execution_otel_invocation", default=None) +) + +# Serializes installation so two invocations starting at once cannot both find +# a handler filterless and both add a filter to it. +_install_lock = threading.Lock() - The plugin it reads is rebindable, because a plugin instance serves exactly - one invocation while the handler it is attached to outlives the invocation. + +def bind_invocation(provider: _SpanContextProvider) -> None: + """Mark ``provider``'s invocation open and claim this thread/task for it. + + Called by a plugin when it takes control on a thread: at invocation start on + the Lambda handler thread, and again from the hooks that run on the threads + executing user code. Idempotent, so a plugin can call it from every such + hook without tracking which threads it has already claimed. Args: - plugin: The OTel plugin instance that resolves the current span context. + provider: The plugin serving the invocation that owns this thread. """ + with _registry_lock: + _open_invocations.add(provider) + _current_invocation.set(provider) + - def __init__(self, plugin: _SpanContextProvider) -> None: - super().__init__() - self._plugin = plugin +def unbind_invocation(provider: _SpanContextProvider) -> None: + """Mark ``provider``'s invocation closed and release its claim on this thread. - def bind(self, plugin: _SpanContextProvider) -> None: - """Point the filter at the plugin serving the current invocation. + Claims made on other threads are not released here -- a context can only be + reset from the thread that set it -- so :func:`_resolve_provider` also checks + that a claim names a still-open invocation. That check is what keeps a + pooled thread outliving its invocation from correlating a later record to a + finished one. + + Args: + provider: The plugin whose invocation has ended. + """ + with _registry_lock: + _open_invocations.discard(provider) + if _current_invocation.get() is provider: + _current_invocation.set(None) - A logging handler lives as long as the Lambda environment, but a plugin - instance lives for one invocation. Without rebinding, the filter - installed by the first invocation's plugin would keep asking that - already-discarded instance for a span context -- it reports none once its - invocation ended, so log correlation would stop after the first - invocation, and the dead instance would be kept reachable for the life of - the environment. - """ - self._plugin = plugin + +def _resolve_provider() -> _SpanContextProvider | None: + """Return the invocation to correlate a record emitted right here against.""" + claimed = _current_invocation.get() + with _registry_lock: + if claimed is not None and claimed in _open_invocations: + return claimed + if len(_open_invocations) == 1: + return next(iter(_open_invocations)) + return None + + +class OtelContextLogFilter(logging.Filter): + """Logging filter that injects the active OTel span context onto records. + + The filter holds no state: it resolves the invocation and the span at emit + time, on the thread that emits the record, so one installed filter serves + any number of concurrent invocations. Resolution is described in the module + docstring; the span itself comes from that invocation's + ``get_current_span_context()``, which returns the active operation span + inside steps and child contexts and falls back to the invocation span for + top-level handler code. + + The filter never caches identifiers and always returns ``True`` so it never + drops a record. + """ def filter(self, record: logging.LogRecord) -> bool: """Stamp the active span context onto the record, then allow it through.""" - span_context = self._plugin.get_current_span_context() + provider = _resolve_provider() + if provider is None: + return True + span_context = provider.get_current_span_context() if span_context and span_context.is_valid: record.traceId = format(span_context.trace_id, "032x") record.spanId = format(span_context.span_id, "016x") @@ -88,7 +174,6 @@ def filter(self, record: logging.LogRecord) -> bool: def install_log_filter( - plugin: _SpanContextProvider, target_logger: logging.Logger | None = None, ) -> OtelContextLogFilter | None: """Attach an OtelContextLogFilter to a logger's handlers, idempotently. @@ -98,14 +183,14 @@ def install_log_filter( records propagated from child loggers are also enriched, since handler filters run for every record reaching the handler. - This is safe to call on every invocation: if a handler already has an - OtelContextLogFilter, that filter is rebound to ``plugin`` and left in place, - so a warm Lambda environment neither stacks duplicate filters nor keeps - reading a previous invocation's plugin. A single shared filter instance is - reused across all handlers. + This is safe to call on every invocation, and from several at once: the + check for an already-installed filter and the install that follows it happen + under one lock, so concurrent first-time callers cannot stack duplicate + filters on a handler. Installation carries no invocation identity -- see + :func:`bind_invocation` for that -- so a warm environment reuses the + filter installed by the first invocation as is. Args: - plugin: The OTel plugin that resolves the current span context. target_logger: Logger whose handlers receive the filter. Defaults to the root logger, which in AWS Lambda is where runtime log handlers live. @@ -115,20 +200,20 @@ def install_log_filter( """ logger = target_logger if target_logger is not None else logging.getLogger() - context_filter: OtelContextLogFilter | None = None - for handler in logger.handlers: - existing = next( - (f for f in handler.filters if isinstance(f, OtelContextLogFilter)), - None, - ) - if existing is not None: - # Reuse the already-installed filter so a single instance is shared, - # and point it at this invocation's plugin. - existing.bind(plugin) - context_filter = existing - continue - if context_filter is None: - context_filter = OtelContextLogFilter(plugin) - handler.addFilter(context_filter) - - return context_filter + with _install_lock: + context_filter: OtelContextLogFilter | None = None + for handler in logger.handlers: + existing = next( + (f for f in handler.filters if isinstance(f, OtelContextLogFilter)), + None, + ) + if existing is not None: + # Reuse the already-installed filter so a single instance is + # shared by every handler. + context_filter = existing + continue + if context_filter is None: + context_filter = OtelContextLogFilter() + handler.addFilter(context_filter) + + return context_filter diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 8fbf4445..1c778cda 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -381,7 +381,7 @@ def test_log_filter_uses_invocation_trace_when_ambient_trace_is_rejected(): exc_info=None, ) - OtelContextLogFilter(plugin).filter(record) + OtelContextLogFilter().filter(record) invocation_span = plugin._get_span(None) assert invocation_span is not None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index 65debda1..df2b6af0 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -3,13 +3,17 @@ from __future__ import annotations import logging +import threading from datetime import UTC, datetime +import pytest from aws_durable_execution_sdk_python.lambda_service import ( OperationStatus, ) from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, InvocationStartInfo, + InvocationStatus, OperationType, UserFunctionStartInfo, ) @@ -17,6 +21,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from aws_durable_execution_sdk_python_otel import log_filter as log_filter_module from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, @@ -29,6 +34,25 @@ EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" +@pytest.fixture(autouse=True) +def _isolated_invocation_registry(): + """Run each test against an empty open-invocation registry. + + The registry is process-global by design -- one installed filter serves every + invocation in the environment -- so a test that starts an invocation without + ending it would otherwise change what later tests resolve. + """ + saved = list(log_filter_module._open_invocations) + log_filter_module._open_invocations.clear() + token = log_filter_module._current_invocation.set(None) + try: + yield + finally: + log_filter_module._current_invocation.reset(token) + log_filter_module._open_invocations.clear() + log_filter_module._open_invocations.update(saved) + + def _create_plugin( enrich_logger: bool = True, ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -46,16 +70,26 @@ def _create_plugin( return plugin, exporter -def _invocation_start_info() -> InvocationStartInfo: +def _invocation_start_info(suffix: str = "") -> InvocationStartInfo: """Create standard invocation start info for tests.""" return InvocationStartInfo( - request_id="request-1", - execution_arn=EXECUTION_ARN, + request_id=f"request-1{suffix}", + execution_arn=f"{EXECUTION_ARN}{suffix}", execution_start_time=START_TIME, is_first_invocation=True, ) +def _invocation_end_info(suffix: str = "") -> InvocationEndInfo: + """Create standard invocation end info for tests.""" + return InvocationEndInfo( + request_id=f"request-1{suffix}", + execution_arn=f"{EXECUTION_ARN}{suffix}", + is_first_invocation=True, + status=InvocationStatus.SUCCEEDED, + ) + + def _user_function_start_info(operation_id: str) -> UserFunctionStartInfo: """Create standard user function start info for tests.""" return UserFunctionStartInfo( @@ -85,6 +119,21 @@ def _make_record() -> logging.LogRecord: ) +def _stamped(record: logging.LogRecord) -> tuple[str | None, str | None]: + """Return the trace and span identifiers a filter stamped on a record.""" + return getattr(record, "traceId", None), getattr(record, "spanId", None) + + +def _own_identifiers(plugin: InvocationOtelPlugin) -> tuple[str, str]: + """Return the trace and span identifiers of the plugin's current span.""" + span_context = plugin.get_current_span_context() + assert span_context is not None + return ( + format(span_context.trace_id, "032x"), + format(span_context.span_id, "016x"), + ) + + def _remove_otel_filters(handler: logging.Handler) -> None: """Remove any OtelContextLogFilter from a handler (test cleanup).""" for log_filter in [ @@ -95,16 +144,15 @@ def _remove_otel_filters(handler: logging.Handler) -> None: def test_filter_always_returns_true(): """The filter never drops a record, even with no active span.""" - plugin, _ = _create_plugin() - log_filter = OtelContextLogFilter(plugin) + log_filter = OtelContextLogFilter() assert log_filter.filter(_make_record()) is True def test_filter_does_not_set_fields_without_active_span(): - """With no invocation active, the filter leaves the record unmodified.""" - plugin, _ = _create_plugin() - log_filter = OtelContextLogFilter(plugin) + """With no invocation open, the filter leaves the record unmodified.""" + _create_plugin() + log_filter = OtelContextLogFilter() record = _make_record() log_filter.filter(record) @@ -118,7 +166,7 @@ def test_filter_injects_trace_context_from_invocation_span(): """The filter stamps the invocation span context for top-level code.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - log_filter = OtelContextLogFilter(plugin) + log_filter = OtelContextLogFilter() record = _make_record() log_filter.filter(record) @@ -134,24 +182,151 @@ def test_filter_uses_attempt_span_inside_user_function(): plugin.on_invocation_start(_invocation_start_info()) operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) + try: + record = _make_record() + OtelContextLogFilter().filter(record) + + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + expected_span_id = format(attempt_span.get_span_context().span_id, "016x") + assert record.spanId == expected_span_id + finally: + # Ends the invocation, which detaches the attempt scope this test + # attached to the running thread's OTel context. Left attached, it would + # stay current for every later test on this thread. + plugin.on_invocation_end(_invocation_end_info()) + + +def test_concurrent_invocations_each_stamp_their_own_span_context(): + """Two invocations open at once each correlate to their own trace. + + Logging handlers are process-global, so both invocations are served by one + filter instance. Each record must carry the trace and span identifiers of the + invocation that emitted it, not of whichever invocation started most + recently. + """ + shared_filter = OtelContextLogFilter() + both_started = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + own: dict[str, tuple[str, str]] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + + def invocation(owner: str) -> None: + try: + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + with lock: + own[owner] = _own_identifiers(plugin) + # Emit only once both invocations are open, so a filter holding + # one mutable plugin reference is guaranteed to have been + # overwritten by the other invocation. + both_started.wait() + record = _make_record() + shared_filter.filter(record) + with lock: + stamped[owner] = _stamped(record) + finally: + plugin.on_invocation_end(_invocation_end_info()) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_started.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures + assert own["a"] != own["b"] + assert stamped["a"] == own["a"] + assert stamped["b"] == own["b"] + + +def test_record_on_an_unclaimed_thread_uses_the_only_open_invocation(): + """A thread the plugin never ran on still correlates to the one invocation. + + The SDK runs the handler body on a worker thread it creates after the + invocation-start hook has run, and Python does not propagate context into a + new thread, so that thread carries no claim. With a single invocation open + there is no ambiguity to resolve. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info()) + try: + expected = _own_identifiers(plugin) + stamped: list[tuple[str | None, str | None]] = [] + + def emit() -> None: + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + + worker = threading.Thread(target=emit, name="unclaimed") + worker.start() + worker.join(timeout=10) + + assert stamped == [expected] + finally: + plugin.on_invocation_end(_invocation_end_info()) + + +def test_record_on_an_unclaimed_thread_is_left_alone_when_two_are_open(): + """An unattributable record is not correlated to an arbitrary invocation.""" + first, _ = _create_plugin(enrich_logger=False) + second, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + second.on_invocation_start(_invocation_start_info(suffix="second")) + try: + records: list[logging.LogRecord] = [] + + def emit() -> None: + record = _make_record() + OtelContextLogFilter().filter(record) + records.append(record) + + worker = threading.Thread(target=emit, name="unclaimed") + worker.start() + worker.join(timeout=10) + + assert len(records) == 1 + assert not hasattr(records[0], "traceId") + assert not hasattr(records[0], "spanId") + finally: + first.on_invocation_end(_invocation_end_info()) + second.on_invocation_end(_invocation_end_info()) + + +def test_finished_invocation_does_not_correlate_later_records(): + """A thread claimed by an invocation stops correlating once it ends. + + A pooled thread can outlive the invocation that claimed it, so liveness is + checked at emit time rather than assumed from the claim. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) record = _make_record() - OtelContextLogFilter(plugin).filter(record) + OtelContextLogFilter().filter(record) - attempt_span = plugin._get_span("step-1:attempt:1") - assert attempt_span is not None - expected_span_id = format(attempt_span.get_span_context().span_id, "016x") - assert record.spanId == expected_span_id + assert not hasattr(record, "traceId") + assert not hasattr(record, "spanId") def test_install_log_filter_attaches_to_handlers(): """install_log_filter adds the filter to each handler on the target logger.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.install") handler = logging.NullHandler() target.addHandler(handler) try: - installed = install_log_filter(plugin, target_logger=target) + installed = install_log_filter(target_logger=target) assert isinstance(installed, OtelContextLogFilter) assert any(isinstance(f, OtelContextLogFilter) for f in handler.filters) @@ -161,13 +336,12 @@ def test_install_log_filter_attaches_to_handlers(): def test_install_log_filter_is_idempotent(): """Repeated installs do not stack duplicate filters on a handler.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.idempotent") handler = logging.NullHandler() target.addHandler(handler) try: - install_log_filter(plugin, target_logger=target) - install_log_filter(plugin, target_logger=target) + install_log_filter(target_logger=target) + install_log_filter(target_logger=target) otel_filters = [ f for f in handler.filters if isinstance(f, OtelContextLogFilter) @@ -177,47 +351,92 @@ def test_install_log_filter_is_idempotent(): target.removeHandler(handler) -def test_install_log_filter_rebinds_to_the_current_invocations_plugin(): - """A later invocation's plugin takes over the already-installed filter. +def test_concurrent_first_time_installs_attach_one_filter(): + """Two invocations installing at once cannot both add a filter. - The handler outlives the invocation but a plugin instance does not, so - without rebinding the filter would keep reading the first invocation's - discarded plugin and stop correlating logs after that invocation. + The handler holds the first install open at the moment it attaches, so a + second caller that was not serialized behind it still sees a filterless + handler and attaches a second filter. """ - first_plugin, _ = _create_plugin() - second_plugin, _ = _create_plugin() + + class HandlerRacingOnAttach(logging.NullHandler): + """Holds the first attach open until a second caller reaches it.""" + + def __init__(self) -> None: + super().__init__() + self._barrier = threading.Barrier(2, timeout=0.2) + self._released = False + + def addFilter(self, filter) -> None: # noqa: A002 - stdlib signature + if not self._released: + try: + self._barrier.wait() + except threading.BrokenBarrierError: + # Installation was serialized, so no second caller arrived. + pass + self._released = True + super().addFilter(filter) + + target = logging.getLogger("test.install.race") + handler = HandlerRacingOnAttach() + target.addHandler(handler) + try: + threads = [ + threading.Thread(target=install_log_filter, args=(target,)) + for _ in range(2) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10) + + otel_filters = [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ] + assert len(otel_filters) == 1 + finally: + target.removeHandler(handler) + + +def test_a_later_invocation_takes_over_log_correlation(): + """The invocation that is open now owns correlation, not the first one. + + A logging handler lives as long as the Lambda environment while a plugin + instance lives for one invocation, so a filter tied to the first + invocation's plugin would stop correlating logs after that invocation ended. + """ + first, _ = _create_plugin() + second, _ = _create_plugin() target = logging.getLogger("test.rebind") handler = logging.NullHandler() target.addHandler(handler) try: - installed = install_log_filter(first_plugin, target_logger=target) - rebound = install_log_filter(second_plugin, target_logger=target) - - assert rebound is installed + installed = install_log_filter(target_logger=target) assert installed is not None - assert installed._plugin is second_plugin + assert install_log_filter(target_logger=target) is installed + + first.on_invocation_start(_invocation_start_info(suffix="first")) + first.on_invocation_end(_invocation_end_info()) + second.on_invocation_start(_invocation_start_info(suffix="second")) - second_plugin.on_invocation_start(_invocation_start_info()) record = _make_record() installed.filter(record) - expected = second_plugin.get_current_span_context() - assert expected is not None - assert record.spanId == format(expected.span_id, "016x") + assert _stamped(record) == _own_identifiers(second) finally: + second.on_invocation_end(_invocation_end_info()) target.removeHandler(handler) def test_install_log_filter_reuses_single_instance_across_handlers(): """A single filter instance is shared across all handlers.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.shared") handler_a = logging.NullHandler() handler_b = logging.NullHandler() target.addHandler(handler_a) target.addHandler(handler_b) try: - installed = install_log_filter(plugin, target_logger=target) + installed = install_log_filter(target_logger=target) filter_a = next( f for f in handler_a.filters if isinstance(f, OtelContextLogFilter) @@ -233,10 +452,9 @@ def test_install_log_filter_reuses_single_instance_across_handlers(): def test_install_log_filter_returns_none_without_handlers(): """With no handlers, install_log_filter has nothing to attach to.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.nohandlers") - assert install_log_filter(plugin, target_logger=target) is None + assert install_log_filter(target_logger=target) is None def test_plugin_installs_filter_on_root_logger_at_construction(): diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 6a21fd67..2d41136f 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -29,6 +29,7 @@ from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPluginFactory, PluginExecutor, + PluginHost, ) from aws_durable_execution_sdk_python.plugin_discovery import ( load_configured_plugins, @@ -191,10 +192,16 @@ def durable_execution( logger.debug("Starting durable execution handler...") - plugin_executor = PluginExecutor(load_configured_plugins(plugins)) + # Only the resolved factory list is handler-lifetime. The plugin instances, + # and the invocation metadata the hooks read, are built per invocation by + # PluginHost.invocation() and live in that invocation's frame -- see the + # plugin_executor parameter below. + plugin_host = PluginHost(load_configured_plugins(plugins)) - @plugin_executor.handle_durable_output - def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: + @plugin_host.handle_durable_output + def wrapper( + event: Any, context: LambdaContext, plugin_executor: PluginExecutor + ) -> MutableMapping[str, Any]: invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient 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 ed46429e..0d54de67 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 @@ -5,7 +5,7 @@ import datetime import functools import logging -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum @@ -477,18 +477,51 @@ def _factory_name(factory: object) -> str: class PluginExecutor: + """One invocation's plugin instances, metadata and dispatch. + + Scoped to a single invocation, not to the handler. Everything mutable here -- + the instances built for this invocation, the start info the end hook derives + from, the operations provider, the dispatch pool -- describes one invocation, + so a single instance shared by two of them would let each overwrite the + other's state. Concurrent executions in one environment (Lambda Managed + Instances) are exactly that case: they run in separate threads of one + process, against one decorated handler. :class:`PluginHost` therefore builds + a fresh executor per invocation and holds it only in that invocation's frame. + + Single-use by construction: :meth:`run` refuses a second entry, so the + lifetime is an invariant of the class rather than a convention its callers + have to keep. Only the factory list is handler-lifetime, and it is copied in + rather than shared mutably. + """ + def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): - # Factories live for the life of the handler; the instances they build do - # not. _plugins is populated in on_invocation_start and emptied when the - # invocation scope exits, so one instance never spans two invocations. + # Factories outlive this executor -- the list is copied, never aliased. + # The instances they build do not: _plugins is populated in + # on_invocation_start and emptied when the invocation scope exits. self._plugin_factories = list(plugins or []) self._plugins: list[DurableInstrumentationPlugin] = [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + self._run_entered = False @contextlib.contextmanager def run(self): + """Open this executor's one invocation scope. + + Raises: + RuntimeError: if entered more than once. A second entry would mean an + executor is serving two invocations, which is the shape this + class exists to prevent; failing loudly here keeps the bug from + reappearing as silent crosstalk. + """ + if self._run_entered: + msg = ( + "PluginExecutor.run() is single-use: this executor has already " + "served an invocation. Build one executor per invocation." + ) + raise RuntimeError(msg) + self._run_entered = True if self._plugin_factories: self._executor = ThreadPoolExecutor( max_workers=1, @@ -500,11 +533,14 @@ def run(self): self._invocation_status = None self._operations_provider = None # Shut down the thread pool, waiting for pending tasks to complete. + # The pool belongs to this invocation, so this drains only this + # invocation's queued dispatches and cannot cut short a concurrent + # invocation's. if self._executor: self._executor.shutdown(wait=True) # Drop this invocation's plugin instances. After the pool has - # drained, so no queued dispatch still holds one: nothing reachable - # from this handler-lifetime executor outlives the invocation. + # drained, so no queued dispatch still holds one: nothing outlives + # the invocation. self._plugins = [] def _create_plugins(self, info: InvocationStartInfo) -> None: @@ -895,21 +931,67 @@ def _is_terminal_status(status): OperationStatus.STOPPED, ] + +class PluginHost: + """Handler-lifetime owner of the configured plugin factories. + + The factory list is the only plugin state that may span invocations: a + factory is resolved once when the handler is initialized and is, by + definition, environment-lifetime. Everything a factory produces is + invocation-lifetime, so this class never holds an instance, a start info or a + dispatch pool -- :meth:`invocation` hands out a fresh + :class:`PluginExecutor` and the caller keeps it in the invocation's own + frame. + + Mirrors the JS SDK's ``createInvocationPluginRunner`` and the Java SDK's + per-invocation ``PluginRunner``: the handler holds factories, the invocation + holds instances. + """ + + def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): + self._plugin_factories = list(plugins or []) + + @contextlib.contextmanager + def invocation(self) -> Iterator[PluginExecutor]: + """Open one invocation's plugin scope and yield its executor. + + The executor is created here rather than at handler-initialization time + so that two invocations sharing this process -- concurrent executions on + a Lambda Managed Instance, or successive executions on a warm + environment -- never write to the same slot. Teardown on scope exit + touches only the executor yielded here. + """ + executor = PluginExecutor(self._plugin_factories) + with executor.run(): + yield executor + @property def handle_durable_output(self): - def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): + """Wrap an invocation body so plugins see its outcome. + + The wrapped function receives this invocation's :class:`PluginExecutor` + as a third argument. Passing it in, rather than closing over one, is what + keeps the instances out of handler-lifetime state: the executor is + reachable only from the frames of the invocation it belongs to. + """ + + def decorator( + func: Callable[ + [Any, LambdaContext, PluginExecutor], MutableMapping[str, Any] + ], + ): @functools.wraps(func) def wrapper(event: Any, context: LambdaContext): - with self.run(): + with self.invocation() as plugin_executor: try: - output = func(event, context) + output = func(event, context, plugin_executor) - self.on_invocation_end( + plugin_executor.on_invocation_end( output=DurableExecutionInvocationOutput.from_dict(output), ) return output except Exception as e: - self.on_invocation_end( + plugin_executor.on_invocation_end( output=DurableExecutionInvocationOutput.create_retry( ErrorObject.from_exception(e) ), 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 0b6a5fcf..1f50ccc9 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 @@ -282,6 +282,9 @@ def __init__( self._current_checkpoint_token: str = initial_checkpoint_token self._operations: dict[str, Operation] = dict(operations) self._service_client: DurableServiceClient = service_client + # Invocation-scoped, like this state object itself: PluginHost builds one + # executor per invocation, so holding it here cannot make one + # invocation's plugin instances visible to another. self._plugin_executor: PluginExecutor = plugin_executor self._operations_lock: Lock = Lock() diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 7d8b6898..0cba5fb3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -2,6 +2,7 @@ import datetime import json +import threading import time import warnings from collections.abc import Sequence @@ -2729,10 +2730,10 @@ def _make_invocation_input(mock_client, next_marker="", input_payload="{}"): ) -def _make_lambda_context(): +def _make_lambda_context(request_id: str = "test-request"): """Helper to create a standard mock Lambda context.""" ctx = Mock() - ctx.aws_request_id = "test-request" + ctx.aws_request_id = request_id ctx.client_context = None ctx.identity = None ctx._epoch_deadline_time_in_ms = 1000000 # noqa: SLF001 @@ -3837,6 +3838,122 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert plugin.calls.count("invocation_end:SUCCEEDED") == 1 +class _TaggedRecordingPlugin(DurableInstrumentationPlugin): + """Records the hooks it receives, tagging each with its own identity.""" + + def __init__(self) -> None: + self.invocation_starts: list[str | None] = [] + self.invocation_ends: list[str] = [] + self.operation_names: list[str | None] = [] + + def on_invocation_start(self, info): + self.invocation_starts.append(info.request_id) + + def on_invocation_end(self, info): + self.invocation_ends.append(f"{info.request_id}:{info.status.value}") + + def on_operation_start(self, info): + self.operation_names.append(info.name) + + def on_operation_end(self, info): + self.operation_names.append(info.name) + + def on_user_function_start(self, info): + self.operation_names.append(info.name) + + def on_user_function_end(self, info): + self.operation_names.append(info.name) + + +def test_durable_execution_keeps_overlapping_invocations_isolated(): + """Two invocations in flight at once never see each other's hooks. + + This is the Lambda Managed Instances case: one decorated handler, one + process, two concurrent executions in separate threads. A barrier holds both + invocations inside their user function at the same time, so both + ``on_invocation_start`` hooks have already fired before either operation + runs, and neither invocation returns until the other has run its operation. + Running the two invocations sequentially would not pin this -- the defect it + guards against is a per-invocation slot being overwritten while both + invocations are live. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + built: dict[str, _TaggedRecordingPlugin] = {} + built_lock = threading.Lock() + + def build_plugin(info) -> _TaggedRecordingPlugin: + plugin = _TaggedRecordingPlugin() + with built_lock: + built[str(info.request_id)] = plugin + return plugin + + timeout = 30 + # Released only once both invocations are inside their user function, so both + # invocation-start hooks have fired and both invocations are live. + both_in_user_code = threading.Barrier(2, timeout=timeout) + # b runs its operation first; a runs its operation afterwards, while b is + # still inside its user function waiting on a. + b_ran_operation = threading.Event() + a_ran_operation = threading.Event() + + @durable_execution(plugins=[build_plugin]) + def test_handler(event: Any, context: DurableContext) -> dict: + tag = event["tag"] + both_in_user_code.wait() + if tag == "b": + context.step(lambda _: "ok", name="step-b") + b_ran_operation.set() + assert a_ran_operation.wait(timeout) + else: + assert b_ran_operation.wait(timeout) + context.step(lambda _: "ok", name="step-a") + a_ran_operation.set() + return {"result": tag} + + results: dict[str, Any] = {} + + def invoke(tag: str) -> None: + results[tag] = test_handler( + _make_invocation_input(mock_client, input_payload=f'{{"tag": "{tag}"}}'), + _make_lambda_context(request_id=f"request-{tag}"), + ) + + threads = [ + threading.Thread(target=invoke, args=(tag,), name=f"invocation-{tag}") + for tag in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=timeout) + for thread in threads: + assert not thread.is_alive(), f"{thread.name} did not finish" + + assert results["a"]["Status"] == InvocationStatus.SUCCEEDED.value + assert results["b"]["Status"] == InvocationStatus.SUCCEEDED.value + + # One instance per invocation, and each one keyed to its own request. + assert sorted(built) == ["request-a", "request-b"] + + for tag in ("a", "b"): + plugin = built[f"request-{tag}"] + # Exactly its own invocation hooks: not the other invocation's request + # id, not two copies of its own, not zero because the other invocation's + # teardown got there first. + assert plugin.invocation_starts == [f"request-{tag}"] + assert plugin.invocation_ends == [f"request-{tag}:SUCCEEDED"] + # Only its own operation. The other invocation's step ran while this one + # was live, so a shared plugin slot would show up here. + assert f"step-{tag}" in plugin.operation_names + other = "b" if tag == "a" else "a" + assert f"step-{other}" not in plugin.operation_names + + def test_durable_execution_with_failing_plugin_factory_does_not_break_execution(): """A factory that raises is contained exactly as a failing hook is.""" mock_client = Mock(spec=DurableServiceClient) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 77cb54dd..55095d0e 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -32,6 +32,7 @@ OperationStartInfo, OperationType, PluginExecutor, + PluginHost, UserFunctionEndInfo, UserFunctionOutcome, UserFunctionStartInfo, @@ -580,18 +581,6 @@ def test_init_records_factories_without_building_plugins(self): class TestPluginLifetime(unittest.TestCase): """The per-invocation plugin lifetime.""" - @staticmethod - def _run_invocation(executor: PluginExecutor, request_id: str) -> None: - lambda_context = MagicMock() - lambda_context.aws_request_id = request_id - with executor.run(): - executor.on_invocation_start( - execution_arn="arn:exec", - lambda_context=lambda_context, - execution_start_time=START_TS, - is_first_invocation=False, - ) - def test_each_invocation_gets_its_own_instance(self): """Two invocations of one handler never share a plugin instance.""" built: list[_TrackingPlugin] = [] @@ -601,10 +590,20 @@ def build(info: InvocationStartInfo) -> _TrackingPlugin: built.append(plugin) return plugin - executor = PluginExecutor(plugins=[build]) + # One host for the handler, one executor per invocation -- the shape + # durable_execution() uses. + host = PluginHost(plugins=[build]) - self._run_invocation(executor, "req-1") - self._run_invocation(executor, "req-2") + for request_id in ("req-1", "req-2"): + lambda_context = MagicMock() + lambda_context.aws_request_id = request_id + with host.invocation() as executor: + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=lambda_context, + execution_start_time=START_TS, + is_first_invocation=False, + ) self.assertEqual(len(built), 2) self.assertIsNot(built[0], built[1]) @@ -612,6 +611,29 @@ def build(info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_host_hands_out_a_new_executor_per_invocation(self): + """The host itself holds no per-invocation state to overwrite.""" + host = PluginHost(plugins=[plugin_factory(_TrackingPlugin())]) + + with host.invocation() as first, host.invocation() as second: + self.assertIsNot(first, second) + + def test_executor_refuses_a_second_invocation(self): + """An executor that has served an invocation cannot serve another. + + The lifetime is enforced by the class, not just by how + ``durable_execution()`` happens to call it, so reintroducing a shared + executor fails loudly instead of silently mixing two invocations. + """ + executor = PluginExecutor(plugins=[plugin_factory(_TrackingPlugin())]) + + with executor.run(): + pass + + with self.assertRaisesRegex(RuntimeError, "single-use"): + with executor.run(): + pass + def test_instances_are_dropped_when_the_invocation_returns(self): """Nothing on the handler-lifetime executor still references the instance.""" plugin = _TrackingPlugin() From 86e0eef25ac5b75a2a99513c73268495b5f2945f Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 17 Sep 2026 15:26:02 -0700 Subject: [PATCH 04/28] fix(otel): correlate top-level logs under concurrency A log record emitted at the top of a handler got no OTel correlation at all when two invocations were open, which contradicts the documented behaviour of the default filter. The chain: `on_invocation_start` runs on the Lambda invocation thread, so the plugin's `ContextVar` claim lands there, but the SDK then submits the handler body to a worker thread, and Python does not propagate a `ContextVar` into a thread started that way. No hook runs on that worker before the handler's first statement -- `on_user_function_start` is dispatched for durable operations only -- so the claim never reached it. The filter fell back to "the single open invocation", which is unavailable when two are open, and left both records unstamped. The SDK now carries the invocation thread's context into the worker that runs the handler body, with `contextvars.copy_context().run`. The copy is taken at submit time, after `on_invocation_start` has run on the same thread, so the claim is inside it. This is the correct layer: the SDK creates that thread to run user code, and any contextvar-based instrumentation is otherwise invisible there. It also matches `asyncio.to_thread`, which runs its callable in a copy of the caller's context. The background checkpointing thread is deliberately left starting from an empty context. It runs SDK checkpointing rather than user code, and it dispatches operation-end hooks, so copying the context there would make the invocation thread's ambient OTel context current on a thread that starts and ends spans. Sharing one copy between the two submissions is impossible anyway, because a `Context` cannot be entered twice concurrently. Propagation had one side effect worth closing rather than documenting. `get_current_span_context` prefers a same-trace span that is current over the plugin's own `Invocation` span, so once the ambient context reached the handler thread, a top-level record named the ambient span instead. Under X-Ray active tracing with the ADOT layer that ambient span is the parent of `Invocation`, so records moved one level up the tree. The plugin now records the span that enclosed the invocation and excludes that one span from the preference, which restores the previous behaviour. Anything that becomes current later is inside the invocation -- an attempt span, a child context span, or a span the handler starts itself -- and still wins, because it is more specific. One case remains uncorrelated by design: a record emitted on a thread that carries no claim. That covers a thread customer code starts itself and the SDK's checkpointing thread. The README states this rather than claiming every record is stamped. BREAKING CHANGE: the handler body now runs in a copy of the invocation thread's context rather than in an empty one, so two behaviours change for every customer, not only for plugin authors. A span started at the top of the handler body now has the ambient span as its parent instead of being a root span; under X-Ray active tracing with the ADOT layer that ambient span is the Lambda invocation span, so existing traces gain a level after upgrading. And a ContextVar set before the decorator is now visible inside the handler body. Both follow from the propagation, which matches what asyncio.to_thread does with the caller's context. --- .../README.md | 32 +- .../execution_plugin.py | 56 +++- .../invocation_plugin.py | 61 +++- .../log_filter.py | 32 +- .../test_concurrent_log_correlation_int.py | 290 ++++++++++++++++++ .../tests/test_execution_plugin.py | 73 +++++ .../tests/test_invocation_plugin.py | 153 +++++++++ .../tests/test_log_filter.py | 190 +++++++++++- .../execution.py | 32 +- .../handler_context_propagation_int_test.py | 263 ++++++++++++++++ 10 files changed, 1130 insertions(+), 52 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 4cbe4850..5efd8947 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -312,15 +312,39 @@ each durable span in the same invocation. ### Log Correlation When `enrich_logger=True` (the default), the plugin installs a logging filter on -the root logger at invocation start. The filter stamps the active OTel trace -context onto every emitted log record using these attributes: +the root logger at invocation start. The filter stamps the trace context that is +active for the emitting invocation onto log records, using these attributes: - `traceId`: 32-char hex trace identifier - `spanId`: 16-char hex span identifier - `otelTraceSampled`: boolean indicating if the trace is sampled -These attributes are only set when a valid span context is active, so any log -formatter or schema must treat the fields as optional. +Which span a record carries follows the emitting thread: the operation attempt +span inside a step, the context span inside a child context, and the Invocation +span for top-level handler code. A span you start yourself on the execution +trace is used for records emitted inside it. A span the runtime already had +active when the invocation started — the Lambda invocation span the ADOT layer +creates under X-Ray active tracing — is the parent of the Invocation span rather +than a substitute for it, so a top-level record still names the Invocation span. +Correlation holds from the first statement of the handler, before any durable +operation, and holds when several invocations run concurrently in one +environment (Lambda Managed Instances): the plugin claims the invocation thread +at invocation start and the SDK runs the handler body in a copy of that thread's +context, so each record resolves to the invocation that emitted it rather than +to whichever invocation started last. + +Two cases are left unstamped, so any log formatter or schema must treat the +fields as optional: + +- No invocation is open — for example during environment initialization or + teardown. +- The record is emitted on a thread that carries no invocation claim, while more + than one invocation is open in the environment. A thread your code starts + itself is such a thread, since Python does not copy context into a new thread, + as is the SDK's background checkpointing thread. With exactly one invocation + open, such a record is correlated to it. With several open there is no way to + tell which one it belongs to, and an uncorrelated record is preferred over one + attributed to another execution. ## Verification 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 7af7e452..3663b3be 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 @@ -164,6 +164,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._invocation_span: Span | None = None + # The span that was already current when this invocation's body began, + # recorded by _record_enclosing_span. Used only to resolve log + # correlation; see get_current_span_context. + self._enclosing_span_context: SpanContext | None = None self._operation_spans: dict[str, Span] = {} # CONTEXT operations that emitted a durable START hook this invocation. # A context absent from this set is checkpointless (for example, a FLAT @@ -299,9 +303,21 @@ def _detach_remaining_contexts(self) -> None: self._detach_context(key) def get_current_span_context(self) -> SpanContext | None: - """Return the active span context for log correlation (see log_filter).""" + """Return the active span context for log correlation (see log_filter). + + A span that became current *inside* this invocation wins: the attempt + span inside a step, the context span inside a child context, or a span + the handler body starts itself. The span that enclosed this invocation is + excluded, so top-level handler records resolve to the Invocation span + rather than to the Workflow span this plugin makes current at invocation + start, which the SDK carries into the thread running the handler body. + """ span_context = trace.get_current_span().get_span_context() - if span_context and span_context.is_valid: + if ( + span_context + and span_context.is_valid + and not self._is_enclosing_span(span_context) + ): return span_context for candidate in (self._invocation_span, self._workflow_span): if candidate is not None: @@ -310,6 +326,26 @@ def get_current_span_context(self) -> SpanContext | None: return ctx return None + def _record_enclosing_span(self) -> None: + """Record the span that is current now, as this invocation's body begins. + + Called at the end of ``on_invocation_start``, on the invocation thread, + which is the context the SDK copies into the thread that runs the handler + body. That is the Workflow span this plugin just attached, so it is less + specific than the Invocation span for a top-level record. + """ + span_context = trace.get_current_span().get_span_context() + self._enclosing_span_context = span_context if span_context.is_valid else None + + def _is_enclosing_span(self, span_context: SpanContext) -> bool: + """Whether ``span_context`` is the span that enclosed this invocation.""" + enclosing = self._enclosing_span_context + return ( + enclosing is not None + and enclosing.trace_id == span_context.trace_id + and enclosing.span_id == span_context.span_id + ) + # ------------------------------------------------------------------ # Links # ------------------------------------------------------------------ @@ -512,6 +548,9 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ), ) + # Last, so that everything this hook makes current is accounted for. + self._record_enclosing_span() + def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -693,13 +732,12 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return - # Runs on the thread that drives the durable operation, which is the - # thread running the handler body and not the thread the - # invocation-start hook claimed. Claim it too, so records emitted from - # top-level handler code are correlated to this invocation even while - # another invocation is open in the same process. Claimed after the - # tracing-enabled gate, so a hook arriving after the invocation ended - # cannot re-register a finished invocation. + # Runs on the thread that drives the durable operation. The thread + # running the handler body already carries this invocation's claim, + # propagated from the invocation thread, but a branch of a map or + # parallel runs on a pool thread that does not, so claim it here. + # Claimed after the tracing-enabled gate, so a hook arriving after the + # invocation ended cannot re-register a finished invocation. bind_invocation(self) if info.operation_type is OperationType.CONTEXT: with self._lock: 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 b0043c6f..686f87bd 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 @@ -162,6 +162,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._span_time_floor_ns: int | None = None + # The span that was already current when this invocation's body began, + # recorded by _record_enclosing_span. Used only to resolve log + # correlation; see get_current_span_context. + self._enclosing_span_context: SpanContext | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} # Replay state supplied by CONTEXT operation START hooks. Missing @@ -300,11 +304,13 @@ def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. Resolution order: - 1. The same-trace span attached to the OTel thread-local context. + 1. A same-trace span that became current *inside* this invocation. Inside a step this is the active attempt span, and inside a child context this is the active context span (attached in - on_user_function_start). Unrelated ambient spans are ignored so logs - stay correlated to the durable execution trace. + on_user_function_start); a span the handler body starts itself also + lands here. Such a span is more specific than the Invocation span, so + it wins. Unrelated ambient spans are ignored so logs stay correlated + to the durable execution trace. 2. The invocation span from the plugin registry. This is the path used for top-level handler code: the invocation span is never attached to the worker thread's context, so the registry is the only way to @@ -312,6 +318,16 @@ def get_current_span_context(self) -> SpanContext | None: detaching the operation scope restores a context with no durable span. + The span that enclosed this invocation is deliberately excluded from + step 1. Under X-Ray active tracing with the ADOT layer, the layer's + Lambda invocation span is current before this invocation starts and is + on the execution trace, and the SDK carries it into the thread running + the handler body. It is also the parent of this plugin's Invocation + span, so preferring it would point top-level records one level up the + tree from the invocation they were emitted by. Anything that becomes + current after the enclosing span was recorded is inside the invocation + and still takes precedence. + Returns: A valid SpanContext, or None if no span is active. """ @@ -320,6 +336,7 @@ def get_current_span_context(self) -> SpanContext | None: span_context and span_context.is_valid and span_context.trace_id == self._execution_trace_id + and not self._is_enclosing_span(span_context) ): return span_context @@ -331,6 +348,28 @@ def get_current_span_context(self) -> SpanContext | None: return None + def _record_enclosing_span(self) -> None: + """Record the span that is current now, as this invocation's body begins. + + Called at the end of ``on_invocation_start``, on the invocation thread, + which is the context the SDK copies into the thread that runs the + handler body. Whatever span is current at that moment existed before the + invocation did -- the ADOT layer's Lambda invocation span, in the X-Ray + active tracing shape -- and is therefore less specific than this + plugin's own Invocation span. + """ + span_context = trace.get_current_span().get_span_context() + self._enclosing_span_context = span_context if span_context.is_valid else None + + def _is_enclosing_span(self, span_context: SpanContext) -> bool: + """Whether ``span_context`` is the span that enclosed this invocation.""" + enclosing = self._enclosing_span_context + return ( + enclosing is not None + and enclosing.trace_id == span_context.trace_id + and enclosing.span_id == span_context.span_id + ) + # ------------------------------------------------------------------ # Context resolution # ------------------------------------------------------------------ @@ -603,6 +642,9 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: attributes=self._extract_attributes(info), ) + # Last, so that everything this hook makes current is accounted for. + self._record_enclosing_span() + def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -760,13 +802,12 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return - # Runs on the thread that drives the durable operation, which is the - # thread running the handler body and not the thread the - # invocation-start hook claimed. Claim it too, so records emitted from - # top-level handler code are correlated to this invocation even while - # another invocation is open in the same process. Claimed after the - # tracing-enabled gate, so a hook arriving after the invocation ended - # cannot re-register a finished invocation. + # Runs on the thread that drives the durable operation. The thread + # running the handler body already carries this invocation's claim, + # propagated from the invocation thread, but a branch of a map or + # parallel runs on a pool thread that does not, so claim it here. + # Claimed after the tracing-enabled gate, so a hook arriving after the + # invocation ended cannot re-register a finished invocation. bind_invocation(self) if info.operation_type is OperationType.CONTEXT: # The user-function hook owns the span, but this durable START hook 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 65cf71fe..68162ea8 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 @@ -32,24 +32,28 @@ which is per-thread and per-task and so cannot be overwritten by a concurrent invocation. - ``unbind_invocation`` marks the invocation closed. - - A record emitted on a thread no invocation has claimed resolves to the one - open invocation, if exactly one is open. That covers the SDK's user-code - worker threads: the invocation-start hook runs on the Lambda handler - thread, the handler body runs on a pool thread the plugin has not been - given control on yet, and Python does not propagate context into new - threads. - - If several invocations are open and the thread is unclaimed, the record is - left unstamped. An unattributed record is a smaller defect than one + - The claim reaches the thread running the handler body because the SDK + submits that work with a copy of the invocation thread's context, taken + after the invocation-start hook has run. Records from top-level handler + code therefore resolve to their own invocation, including the first + statement of the handler, before any durable operation has claimed the + thread directly. + - A record emitted on a thread that carries no claim resolves to the one + open invocation, if exactly one is open. Threads the SDK does not submit + the invocation's context into land here: a thread customer code starts + itself, since Python does not copy context into a new thread, and the + SDK's background checkpointing thread. + - If several invocations are open and the thread carries no claim, the record + is left unstamped. An unattributed record is a smaller defect than one attributed to another customer execution. Reading the active span straight from the OTel context (as the Java plugin's static ``MdcSpanEnricher`` does) is not sufficient here: the invocation span is -never attached to the OTel context, and the context of the handler thread -- the -only place the plugin attaches anything at invocation scope -- is not visible on -the worker thread that runs the handler body. Top-level records would silently -lose correlation. The plugin's ``get_current_span_context()`` still reads the -OTel context first, so records emitted inside a step or child context resolve to -the active operation span exactly as before. +never attached to the OTel context, so a record emitted between durable +operations would find no durable span current and silently lose correlation. The +plugin's ``get_current_span_context()`` still reads the OTel context first, so +records emitted inside a step or child context resolve to the active operation +span exactly as before. """ from __future__ import annotations diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py new file mode 100644 index 00000000..198b648d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py @@ -0,0 +1,290 @@ +"""End-to-end log correlation coverage for concurrent invocations. + +Drives the decorated handler, not the filter in isolation. The SDK runs the +handler body on a worker thread it submits to a pool, so only a test that goes +through ``durable_execution`` exercises the path from the invocation-start hook +to a record emitted by top-level handler code. + +Both invocations are held at a barrier until each has started, so a record is +only emitted while two invocations are open. That is the case the log filter +cannot resolve from the number of open invocations alone. +""" + +from __future__ import annotations + +import json +import logging +import threading +from dataclasses import replace +from datetime import UTC, datetime +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + ExecutionDetails, + Operation, + OperationAction, + OperationStatus, + OperationType, + StepDetails, +) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + InvocationOtelPluginFactory, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +EXECUTION_START = datetime(2026, 8, 27, 5, 11, 47, tzinfo=UTC) +OWNERS = ("first", "second") +TOP_LEVEL_MESSAGE = "top-level-handler-log" + + +def _execution_arn(owner: str) -> str: + return f"test-arn/{_execution_operation_id(owner)}" + + +def _execution_operation_id(owner: str) -> str: + """The execution operation is keyed by the last segment of the ARN.""" + return f"concurrent-log-correlation-{owner}" + + +def _lambda_context(owner: str) -> Mock: + context = Mock() + context.aws_request_id = f"request-{owner}" + context.client_context = None + context.identity = None + context._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + context.invoked_function_arn = "test-arn" + context.tenant_id = None + return context + + +def _execution_operation(owner: str) -> Operation: + return Operation( + operation_id=_execution_operation_id(owner), + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=EXECUTION_START, + # The handler receives the parsed input payload, so the owner travels + # through the execution input rather than the invocation event. + execution_details=ExecutionDetails(input_payload=json.dumps({"owner": owner})), + ) + + +def _event(owner: str) -> dict[str, Any]: + return { + "DurableExecutionArn": _execution_arn(owner), + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [_execution_operation(owner).to_json_dict()], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _shared_checkpoint_store(): + """Return a checkpoint callable serving several concurrent executions. + + One mocked Lambda client is shared by both invocations, so operations are + kept per execution ARN and mutated under a lock. + """ + operations: dict[str, dict[str, Operation]] = { + _execution_arn(owner): { + _execution_operation_id(owner): _execution_operation(owner) + } + for owner in OWNERS + } + lock = threading.Lock() + + def checkpoint( + durable_execution_arn, + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + with lock: + execution_operations = operations.setdefault(durable_execution_arn, {}) + for update in updates: + now = datetime.now(UTC) + previous = execution_operations.get(update.operation_id) + if update.action is OperationAction.START: + execution_operations[update.operation_id] = Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + elif update.action is OperationAction.SUCCEED: + base = previous or Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + execution_operations[update.operation_id] = replace( + base, + status=OperationStatus.SUCCEEDED, + end_timestamp=now, + step_details=( + StepDetails(result=update.payload, attempt=1) + if update.operation_type is OperationType.STEP + else base.step_details + ), + ) + snapshot = list(execution_operations.values()) + + return CheckpointOutput( + checkpoint_token="new-token", + new_execution_state=CheckpointUpdatedExecutionState(operations=snapshot), + ) + + return checkpoint + + +class _RecordCollector(logging.Handler): + """Collects the top-level handler records the filter has stamped.""" + + def __init__(self) -> None: + super().__init__() + self.records: list[logging.LogRecord] = [] + self._lock = threading.Lock() + + def emit(self, record: logging.LogRecord) -> None: + if record.getMessage().startswith(TOP_LEVEL_MESSAGE): + with self._lock: + self.records.append(record) + + +def _remove_otel_filters(handler: logging.Handler) -> None: + for installed in [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ]: + handler.removeFilter(installed) + + +def test_overlapping_invocations_stamp_top_level_logs_with_their_own_span() -> None: + """A handler's first log carries its own invocation's trace and span ids. + + Both invocations are open when either record is emitted, and neither record + is emitted from a thread the plugin hooks have run on, so the record can + only be correlated if the SDK carried the invocation's context into the + thread it runs the handler body on. + """ + collector = _RecordCollector() + root = logging.getLogger() + root.addHandler(collector) + # The record must reach the root handlers, so the emitting logger opts in to + # INFO explicitly rather than inheriting the root level. + probe_logger = logging.getLogger("probe.handler") + previous_level = probe_logger.level + probe_logger.setLevel(logging.INFO) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + # enrich_logger is the documented default: the plugin installs the filter on + # the root logger's handlers, including the collector added above. + factory = InvocationOtelPluginFactory( + OtelPluginConfig(tracer_provider=provider, enrich_logger=True) + ) + + both_started = threading.Barrier(len(OWNERS), timeout=30) + + @durable_step + def after_log(_step_context) -> str: + return "done" + + def handler_impl(event: Any, context: DurableContext) -> str: + owner = event["owner"] + # Hold here until every invocation has started, so the log below is + # emitted with more than one invocation open. Nothing durable has run. + both_started.wait() + logging.getLogger("probe.handler").info("%s %s", TOP_LEVEL_MESSAGE, owner) + return context.step(after_log(), name=f"after-log-{owner}") + + handler = durable_execution(handler_impl, plugins=[factory]) + results: dict[str, Any] = {} + failures: list[BaseException] = [] + results_lock = threading.Lock() + + def invoke(owner: str) -> None: + try: + result = handler(_event(owner), _lambda_context(owner)) + with results_lock: + results[owner] = result + except BaseException as error: # noqa: BLE001 + with results_lock: + failures.append(error) + both_started.abort() + + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _shared_checkpoint_store() + mock_client_class.initialize_client.return_value = mock_client + + threads = [ + threading.Thread(target=invoke, args=(owner,), name=f"invoke-{owner}") + for owner in OWNERS + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + + assert not failures, failures + assert not any(thread.is_alive() for thread in threads) + for owner in OWNERS: + assert results[owner]["Status"] == InvocationStatus.SUCCEEDED.value + + # Each invocation's Invocation span identifies its execution by ARN, so + # the expected identifiers are read back from the exported spans rather + # than recomputed. + expected: dict[str, tuple[str, str]] = {} + for span in exporter.get_finished_spans(): + if span.name != "Invocation": + continue + assert span.attributes is not None + arn = span.attributes["durable.execution.arn"] + assert span.context is not None + expected[str(arn)] = ( + format(span.context.trace_id, "032x"), + format(span.context.span_id, "016x"), + ) + assert set(expected) == {_execution_arn(owner) for owner in OWNERS} + + stamped = { + record.getMessage().rsplit(" ", 1)[1]: ( + getattr(record, "traceId", None), + getattr(record, "spanId", None), + ) + for record in collector.records + } + assert set(stamped) == set(OWNERS), collector.records + assert stamped == { + owner: expected[_execution_arn(owner)] for owner in OWNERS + }, f"records carried {stamped}" + finally: + for installed_on in list(root.handlers): + _remove_otel_filters(installed_on) + root.removeHandler(collector) + probe_logger.setLevel(previous_level) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 8c94d257..3fd0d9e9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import threading from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta @@ -32,6 +33,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( NonRecordingSpan, + Span, SpanContext, SpanKind, TraceFlags, @@ -48,6 +50,7 @@ from aws_durable_execution_sdk_python_otel.durable_parent_span import ( DurableParentSpan, ) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig @@ -1644,3 +1647,73 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + + +# --------------------------------------------------------------------------- +# Log correlation +# --------------------------------------------------------------------------- +def _stamped_span_id() -> str: + """Return the span ID the log filter stamps on a record emitted right here.""" + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="message", + args=(), + exc_info=None, + ) + OtelContextLogFilter().filter(record) + return str(getattr(record, "spanId", None)) + + +def _span_id_hex(span: Span) -> str: + """Return a span's ID in the hex form the log filter stamps.""" + return format(span.get_span_context().span_id, "016x") + + +def test_top_level_log_names_invocation_span_not_the_attached_workflow_span(): + """A top-level handler record names the Invocation span, not the Workflow span. + + This plugin makes the Workflow span current at invocation start so + auto-instrumented spans join the execution trace, and the SDK carries the + invocation thread's context into the thread that runs the handler body. The + Workflow span spans the whole execution, so it is less specific than the + Invocation span for a record emitted by one invocation's top-level code. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + try: + assert plugin._invocation_span is not None + assert plugin._workflow_span is not None + # Confirm the shape under test: the Workflow span is the current span. + assert trace.get_current_span() is plugin._workflow_span + + stamped = _stamped_span_id() + + assert stamped == _span_id_hex(plugin._invocation_span) + assert stamped != _span_id_hex(plugin._workflow_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) + + +def test_step_attempt_log_names_the_attempt_span(): + """A record inside a step names the attempt span, not the Invocation span. + + Excluding the Workflow span from log correlation must not also exclude spans + that become current inside the invocation. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + try: + plugin.on_user_function_start(_step_start_info("step-1")) + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + assert plugin._invocation_span is not None + + stamped = _stamped_span_id() + + assert stamped == _span_id_hex(attempt_span) + assert stamped != _span_id_hex(plugin._invocation_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 1c778cda..efc765d3 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -35,6 +35,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( NonRecordingSpan, + Span, SpanContext, SpanKind, StatusCode, @@ -134,6 +135,45 @@ def _invocation_end_info( ) +def _same_trace_ambient_context() -> SpanContext: + """Return the span context of an ADOT-style enclosing Lambda span. + + Under X-Ray active tracing the execution trace is the ambient trace, so the + span the ADOT layer has current before the invocation starts is on the + execution trace and becomes the parent of the Invocation span. + """ + return SpanContext( + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME), + span_id=int("1234567890abcdef", 16), + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + + +def _stamped_span_id(plugin: InvocationOtelPlugin) -> str: + """Return the span ID the log filter stamps on a record emitted right here.""" + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="message", + args=(), + exc_info=None, + ) + OtelContextLogFilter().filter(record) + assert getattr(record, "traceId", None) == format( + plugin._execution_trace_id or 0, "032x" + ) + return str(getattr(record, "spanId", None)) + + +def _span_id_hex(span: Span) -> str: + """Return a span's ID in the hex form the log filter stamps.""" + return format(span.get_span_context().span_id, "016x") + + def _user_function_start_info( operation_id: str, attempt: int = 1, @@ -426,6 +466,119 @@ def test_invocation_span_parents_to_same_trace_ambient_span(): assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) +def test_top_level_log_names_invocation_span_not_the_enclosing_ambient_span(): + """A top-level handler record names the Invocation span, not its parent. + + In the X-Ray active tracing plus ADOT shape the layer's Lambda span is + current before the invocation starts, is on the execution trace, and is the + parent of the Invocation span. The SDK carries the invocation thread's + context into the thread that runs the handler body, so that span is current + where top-level handler code runs. It is one level up the tree from the + invocation that emitted the record, so the Invocation span is the more + specific answer and the one a top-level record must carry. + """ + plugin, _ = _create_plugin() + ambient_context = _same_trace_ambient_context() + ambient = NonRecordingSpan(ambient_context) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + # Confirm the shape under test: same trace, and the ambient span really + # is the parent of the Invocation span. + assert plugin._execution_trace_id == ambient_context.trace_id + assert invocation_span.parent is not None + assert invocation_span.parent.span_id == ambient_context.span_id + + stamped = _stamped_span_id(plugin) + + assert stamped == _span_id_hex(invocation_span) + assert stamped != format(ambient_context.span_id, "016x") + + # The same holds between top-level operations: once a step's scope is + # released the enclosing ambient span is current again. + plugin.on_user_function_start(_user_function_start_info("step-1")) + plugin.on_user_function_end(_user_function_end_info("step-1")) + assert trace.get_current_span().get_span_context().span_id == ( + ambient_context.span_id + ) + assert _stamped_span_id(plugin) == _span_id_hex(invocation_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + +def test_durable_operation_spans_win_over_the_enclosing_ambient_span(): + """A record inside a durable operation names that operation's span. + + The enclosing ambient span is excluded from log correlation, but a span that + becomes current inside the invocation is more specific than the Invocation + span and still wins. Both operation shapes that attach a scope are covered: + a STEP attempt and a child CONTEXT. + """ + plugin, _ = _create_plugin() + ambient = NonRecordingSpan(_same_trace_ambient_context()) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + + plugin.on_user_function_start(_user_function_start_info("step-1")) + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + assert _stamped_span_id(plugin) == _span_id_hex(attempt_span) + assert _stamped_span_id(plugin) != _span_id_hex(invocation_span) + plugin.on_user_function_end(_user_function_end_info("step-1")) + + plugin.on_user_function_start( + _user_function_start_info("ctx-1", operation_type=OperationType.CONTEXT) + ) + context_span = plugin._get_span("ctx-1") + assert context_span is not None + assert _stamped_span_id(plugin) == _span_id_hex(context_span) + assert _stamped_span_id(plugin) != _span_id_hex(invocation_span) + finally: + # The child context never ends, so invocation cleanup releases its scope. + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + +def test_customer_span_started_in_the_handler_wins_over_the_invocation_span(): + """A record inside a span the handler body started names that span. + + A span the customer creates while the invocation is running is on the + execution trace, because the enclosing ambient span is current when it + starts, and it is not the enclosing span itself, so it is the most specific + span for a record emitted inside it. + """ + plugin, _ = _create_plugin() + ambient = NonRecordingSpan(_same_trace_ambient_context()) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + + customer_span = plugin._provider.get_tracer("customer").start_span( + "customer-work" + ) + customer_token = otel_context.attach( + trace.set_span_in_context(customer_span, otel_context.get_current()) + ) + try: + stamped = _stamped_span_id(plugin) + assert stamped == _span_id_hex(customer_span) + assert stamped != _span_id_hex(invocation_span) + finally: + otel_context.detach(customer_token) + customer_span.end() + finally: + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + def test_pre_terminal_placeholder_preserves_same_trace_tracestate(): """The Workflow placeholder and operation links carry ambient tracestate.""" plugin, exporter = _create_plugin() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index df2b6af0..f46d5999 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -2,10 +2,12 @@ from __future__ import annotations +import contextvars import logging import threading from datetime import UTC, datetime +import opentelemetry.context as otel_context import pytest from aws_durable_execution_sdk_python.lambda_service import ( OperationStatus, @@ -17,11 +19,22 @@ OperationType, UserFunctionStartInfo, ) +from opentelemetry import trace +from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import ( + NonRecordingSpan, + SpanContext, + TraceFlags, + TraceState, +) from aws_durable_execution_sdk_python_otel import log_filter as log_filter_module +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + _to_otel_trace_id, +) from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, @@ -250,12 +263,12 @@ def invocation(owner: str) -> None: def test_record_on_an_unclaimed_thread_uses_the_only_open_invocation(): - """A thread the plugin never ran on still correlates to the one invocation. + """A thread carrying no claim still correlates to the one open invocation. - The SDK runs the handler body on a worker thread it creates after the - invocation-start hook has run, and Python does not propagate context into a - new thread, so that thread carries no claim. With a single invocation open - there is no ambiguity to resolve. + A thread that carries no claim -- one customer code started itself, since + ``threading.Thread`` does not copy the starting thread's context -- has no + invocation of its own. With a single invocation open there is no ambiguity to + resolve, so the record is correlated to it. """ plugin, _ = _create_plugin(enrich_logger=False) plugin.on_invocation_start(_invocation_start_info()) @@ -277,27 +290,178 @@ def emit() -> None: plugin.on_invocation_end(_invocation_end_info()) -def test_record_on_an_unclaimed_thread_is_left_alone_when_two_are_open(): - """An unattributable record is not correlated to an arbitrary invocation.""" +def test_context_propagated_into_a_worker_resolves_the_claiming_invocation(): + """A worker started from a copy of the claiming thread's context resolves it. + + This is what the SDK does for the thread it runs the handler body on: the + invocation-start hook claims the invocation thread, and the handler body runs + in a copy of that thread's context. Both invocations are open when either + record is emitted, so the number of open invocations cannot resolve them and + only the propagated claim can. + """ + both_emitted = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + own: dict[str, tuple[str, str]] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + + def invocation(owner: str) -> None: + """Run one invocation the way the SDK does, on its own thread.""" + try: + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + with lock: + own[owner] = _own_identifiers(plugin) + + def emit() -> None: + record = _make_record() + # Emit only once both invocations are open, so the + # single-open-invocation fallback cannot resolve the record. + both_emitted.wait() + OtelContextLogFilter().filter(record) + with lock: + stamped[owner] = _stamped(record) + + # A fresh copy per submission: one Context cannot be entered + # twice concurrently, which is why the SDK copies at each submit. + worker = threading.Thread( + target=contextvars.copy_context().run, + args=(emit,), + name=f"worker-{owner}", + ) + worker.start() + worker.join(timeout=10) + finally: + plugin.on_invocation_end(_invocation_end_info()) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_emitted.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures + assert own["a"] != own["b"] + assert stamped == own + + +def test_concurrent_invocations_with_enclosing_ambient_spans_stay_separate(): + """Excluding the enclosing ambient span does not blur two open invocations. + + Each invocation runs with its own ADOT-style enclosing span current, on its + own execution trace, and both are open when either record is emitted. Each + record must carry its own invocation's Invocation span: not the other + invocation's span, and not its own enclosing span. + """ + both_started = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + invocation_spans: dict[str, tuple[str, str]] = {} + ambient_span_ids: dict[str, str] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + ambient_span_id_by_owner = { + "a": int("1111aaaa1111aaaa", 16), + "b": int("2222bbbb2222bbbb", 16), + } + + def invocation(owner: str) -> None: + try: + plugin, _ = _create_plugin(enrich_logger=False) + # The enclosing span sits on this execution's own trace, which is + # what the ADOT layer produces under X-Ray active tracing. + ambient_context = SpanContext( + trace_id=_to_otel_trace_id(f"{EXECUTION_ARN}{owner}", START_TIME), + span_id=ambient_span_id_by_owner[owner], + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + token = otel_context.attach( + trace.set_span_in_context(NonRecordingSpan(ambient_context), Context()) + ) + try: + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + invocation_span = plugin._get_span(None) + assert invocation_span is not None + span_context = invocation_span.get_span_context() + with lock: + invocation_spans[owner] = ( + format(span_context.trace_id, "032x"), + format(span_context.span_id, "016x"), + ) + ambient_span_ids[owner] = format( + ambient_context.span_id, "016x" + ) + both_started.wait() + record = _make_record() + OtelContextLogFilter().filter(record) + with lock: + stamped[owner] = _stamped(record) + finally: + plugin.on_invocation_end(_invocation_end_info(suffix=owner)) + finally: + otel_context.detach(token) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_started.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures, failures + assert invocation_spans["a"] != invocation_spans["b"] + assert stamped == invocation_spans + for owner in ("a", "b"): + assert stamped[owner][1] != ambient_span_ids[owner] + + +def test_record_on_a_customer_thread_is_not_attributed_to_another_invocation(): + """A thread customer code starts itself is never given the wrong invocation. + + ``threading.Thread`` does not copy the starting thread's context, so a thread + a handler creates directly carries no claim. This case is not the SDK's + handler worker, which is submitted with a copy of the invocation's context. + With two invocations open there is nothing to resolve such a record against, + and it is left uncorrelated rather than attributed to either invocation. + """ first, _ = _create_plugin(enrich_logger=False) second, _ = _create_plugin(enrich_logger=False) first.on_invocation_start(_invocation_start_info(suffix="first")) second.on_invocation_start(_invocation_start_info(suffix="second")) try: - records: list[logging.LogRecord] = [] + wrong_identifiers = {_own_identifiers(first), _own_identifiers(second)} + stamped: list[tuple[str | None, str | None]] = [] def emit() -> None: record = _make_record() OtelContextLogFilter().filter(record) - records.append(record) + stamped.append(_stamped(record)) - worker = threading.Thread(target=emit, name="unclaimed") + worker = threading.Thread(target=emit, name="customer-created") worker.start() worker.join(timeout=10) - assert len(records) == 1 - assert not hasattr(records[0], "traceId") - assert not hasattr(records[0], "spanId") + assert len(stamped) == 1 + # The rule that matters: never another execution's trace. + assert stamped[0] not in wrong_identifiers + # And with no claim and no single open invocation, nothing is stamped. + assert stamped[0] == (None, None) finally: first.on_invocation_end(_invocation_end_info()) second.on_invocation_end(_invocation_end_info()) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 2d41136f..1e7b9578 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import contextvars import functools import json import logging @@ -316,14 +317,41 @@ def wrapper( operations_provider=lambda: execution_state.operations, updated_operation_ids=invocation_input.updated_operation_ids, ) - # Thread 1: Run background checkpoint processing + # Thread 1: Run background checkpoint processing. + # + # Submitted without the invocation's context, deliberately. This + # thread runs SDK checkpointing, never user code, and nothing it + # does needs a contextvar the invocation thread set. Copying the + # context here would also make the invocation thread's ambient OTel + # context current on a background thread that starts and ends + # plugin spans, which widens the change with no caller-visible + # benefit. executor.submit(execution_state.checkpoint_batches_forever) # Thread 2: Execute user function logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - user_future = executor.submit(func, input_event, durable_context) + # Carry this thread's context into the worker that runs the handler + # body. A callable submitted to a ThreadPoolExecutor runs on a worker + # thread, whose context is not the submitting thread's, so without + # this the handler body starts from a context in which no contextvar + # set before the handler is visible -- neither those set by the + # plugins that ran in on_invocation_start just above, nor those set + # by customer code around the decorator. Log correlation depends on + # it, since a plugin that claims the invocation for the calling + # thread at invocation start has no other way to reach the thread the + # handler body runs on. + # + # The copy is taken here, per submission: a Context cannot be entered + # twice concurrently, so the checkpoint thread above could not share + # one with this. Copying does not couple the two threads -- the + # worker mutates its own copy, and the invocation thread's context is + # unchanged either way, exactly as it was when the worker started + # from an unrelated context. + user_future = executor.submit( + contextvars.copy_context().run, func, input_event, durable_context + ) logger.debug( "%s waiting for user code completion...", 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 new file mode 100644 index 00000000..37e2614b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py @@ -0,0 +1,263 @@ +"""Integration tests for context propagation into the user-function thread. + +The SDK runs the handler body on a worker thread it submits to a pool. A +submitted callable is given a context of its own, so anything a contextvar +carries -- set by a plugin in ``on_invocation_start``, or by customer code +around the decorator -- is only visible to top-level handler code if the +invocation thread's context is carried into that worker. + +Log-correlating plugins depend on this: a plugin that claims the invocation for +the calling thread at invocation start has no hook that runs on the worker +before the handler body, so the claim reaches top-level handler code by context +propagation or not at all. +""" + +from __future__ import annotations + +import contextvars +import threading +from dataclasses import replace +from datetime import UTC, datetime +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationAction, + OperationStatus, + OperationType, + StepDetails, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from tests.test_helpers import plugin_factory + + +UNSET = "unset" +_probe: contextvars.ContextVar[str] = contextvars.ContextVar( + "handler_context_propagation_probe", default=UNSET +) + + +def _lambda_context() -> Mock: + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +def _event() -> dict: + return { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +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. + """ + operations: dict[str, Operation] = {} + + def mock_checkpoint( + durable_execution_arn, # noqa: ARG001 + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + for update in updates: + previous = operations.get(update.operation_id) + base = previous or Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=datetime.now(UTC), + ) + if update.action is OperationAction.SUCCEED: + operations[update.operation_id] = replace( + base, + status=OperationStatus.SUCCEEDED, + end_timestamp=datetime.now(UTC), + step_details=( + StepDetails(result=update.payload, attempt=1) + if update.operation_type is OperationType.STEP + else base.step_details + ), + ) + else: + operations[update.operation_id] = base + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=list(operations.values()) + ), + ) + + return mock_checkpoint + + +def _run(handler, event: dict | None = None) -> dict: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + return handler(event if event is not None else _event(), _lambda_context()) + + +class _ContextvarSettingPlugin(DurableInstrumentationPlugin): + """Sets a contextvar on the invocation thread, as a log-correlating plugin does.""" + + def __init__(self, value: str) -> None: + self._value = value + self.operation_end_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. + with self._lock: + self.operation_end_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. + + The hook runs on the invocation thread and the handler body runs on a worker + thread, and no hook runs on that worker before the handler's first + statement, so this only holds if the invocation's context is propagated. + """ + plugin = _ContextvarSettingPlugin("claimed-by-plugin") + seen: list[str] = [] + + @durable_execution(plugins=[plugin_factory(plugin)]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + seen.append(_probe.get()) + return "ok" + + result = _run(my_handler) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert seen == ["claimed-by-plugin"] + + +def test_handler_body_sees_a_contextvar_set_by_the_caller(): + """A contextvar set before the handler is visible inside the handler body. + + This matches ``asyncio.to_thread``, which also runs the callable in a copy + of the caller's context. + """ + seen: list[str] = [] + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + seen.append(_probe.get()) + return "ok" + + token = _probe.set("set-by-caller") + try: + result = _run(my_handler) + finally: + _probe.reset(token) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert seen == ["set-by-caller"] + + +def test_handler_body_contextvar_writes_do_not_leak_to_the_caller(): + """The worker mutates its own copy, so the caller's context is untouched. + + A worker thread has a context of its own whether it starts empty or from a + copy, so propagation adds no path from the handler back to the invocation + thread. + """ + observed_in_step: list[str] = [] + + @durable_step + def read_probe(_step_context) -> str: + observed_in_step.append(_probe.get()) + return _probe.get() + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + _probe.set("set-inside-handler") + return context.step(read_probe(), name="read-probe") + + token = _probe.set("set-by-caller") + try: + result = _run(my_handler) + after = _probe.get() + finally: + _probe.reset(token) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + # The write is visible to code the handler drives, and nowhere else. + assert observed_in_step == ["set-inside-handler"] + assert after == "set-by-caller" + + +def test_checkpointing_thread_does_not_carry_the_invocation_context(): + """The background checkpointing thread is submitted without the 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. + """ + 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 + off_handler_thread = [ + value + for thread_name, value in plugin.operation_end_observations + if thread_name not in handler_thread_names + ] + assert off_handler_thread, plugin.operation_end_observations + assert set(off_handler_thread) == {UNSET} From 1b060b14477cb3ec458d95d80170be94c6f373c5 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 17 Sep 2026 16:19:25 -0700 Subject: [PATCH 05/28] fix(insight): contain exporter BaseException; bump to 3.0.0 Four review findings. The first is the most serious defect in this branch so far, because it exhausts the environment rather than only stalling it. An exporter whose flush() raises a BaseException on every call used to hang the invocation without bound. A drain is released only by a flush that completed. The flush raised, so nothing advanced, and the worker died. A drain replaces a dead worker, so the drain started a replacement, which ran the same failing flush and died the same way. Measured over three seconds on one such exporter: 4842 flush attempts, 3682 live threads, and the drain never returned. So the earlier fix that let a replacement start converted a permanent park into a thread storm. The same finally block also advanced the completion bookkeeping when the first exporter aborted the fan-out with a BaseException. Measured with a failing exporter followed by a healthy one: the drain returned, the healthy exporter received nothing, and coverage read as delivered. The snapshot was already out of the pending slot, so the loss was permanent. Both exporter call sites now contain BaseException, so the containment that already held for Exception holds for every shape an exporter can raise. That is safe here specifically because none of the three signalling cases can reach these sites as a signal: the interpreter raises KeyboardInterrupt only in the main thread, threading discards a SystemExit raised in a worker thread, and nothing cancels the export worker because nothing outside the scheduler knows it exists. A BaseException seen there was raised by the exporter, so it reports a defective exporter rather than instructing this thread to stop. asyncio.CancelledError is the reachable case, since an exporter that touches asyncio raises it without writing `raise`. Consecutive worker deaths are now bounded at three, as a backstop for the scheduler's own code rather than for exporters. At the bound the scheduler latches asynchronous export off, wakes every parked waiter and drops what is queued, with a warning. Releasing the waiter is the priority because the waiter is an invocation thread inside on_invocation_end: parking it turns an instrumentation defect into a stalled customer execution, while dropping records loses instrumentation data only. Any completed export or flush resets the count, so a worker making progress never approaches it. Second, the decorated handler advertised a phantom argument. functools.wraps was applied to a three-argument internal function, so inspect.signature() on the handler reported a required plugin_executor parameter and signature.bind(event, context) raised TypeError. A signature-aware runtime or test harness can reject such a handler. The user function's metadata is now copied at the head of the wrapper chain, so the handler reports (event, context) again. Third, explicit plugins entries were never validated. The entry-point path rejects a non-callable target, but explicit entries were copied through, so plugins=[MyPlugin()] passed configuration and then raised TypeError inside _create_plugins on every invocation, which logged it and continued without that plugin. A breaking migration therefore produced repeated telemetry loss instead of one configuration failure. Explicit entries are now validated at load time and name the offending position. A plugin class is still accepted, because calling a class in Python constructs an instance, so plugins=[MyPlugin] is a valid factory. Fourth, both plugin packages allowed a core version that cannot work. They declared aws-durable-execution-sdk-python>=2.0.0 while requiring the factory contract, so pip accepted a resolution that fails at handler initialization. RELEASING.md makes version numbers this repository's responsibility, so this commit sets them: core 2.0.0 to 3.0.0 for the removed provider contract, otel 1.0.0 to 2.0.0 for the replaced plugin constructors, and insight 0.0.1 to 0.1.0 rather than 0.0.2 so a consumer pinning ~=0.0.1 cannot silently accept a breaking release. Both plugin bounds move to >=3.0.0, the OTel layer's sdk-version pin follows, and the two conformance packages' exact pins follow. Two new metadata tests assert that the declared bound tracks the core major this repository builds and that the layer pin matches the core version, so the next bump cannot leave one of them behind. --- .github/lambda-layer-publish.toml | 2 +- .../pyproject.toml | 4 +- .../pyproject.toml | 2 +- .../pyproject.toml | 7 +- .../__about__.py | 2 +- .../_export_scheduler.py | 180 +++++++++++++---- .../types.py | 13 +- .../tests/test_export_scheduler.py | 182 ++++++++++++++++++ .../tests/test_package_metadata.py | 78 ++++++++ .../README.md | 2 +- .../pyproject.toml | 6 +- .../__about__.py | 2 +- .../tests/test_package_metadata.py | 79 +++++++- .../__about__.py | 2 +- .../execution.py | 10 + .../plugin_discovery.py | 34 +++- .../tests/execution_test.py | 80 ++++++++ .../tests/plugin_discovery_test.py | 115 +++++++++++ pyproject.toml | 2 +- 19 files changed, 744 insertions(+), 58 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py diff --git a/.github/lambda-layer-publish.toml b/.github/lambda-layer-publish.toml index f141c2a0..4fd02156 100644 --- a/.github/lambda-layer-publish.toml +++ b/.github/lambda-layer-publish.toml @@ -1,2 +1,2 @@ [layer] -sdk-version = "2.0.0" +sdk-version = "3.0.0" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml index 465a475d..3deffbdf 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml @@ -8,8 +8,8 @@ version = "0.0.0" description = "OpenTelemetry conformance test handlers for the AWS Durable Execution SDK for Python, exercised by the aws-durable-execution-conformance-tests OTel suites." requires-python = ">=3.11" dependencies = [ - "aws-durable-execution-sdk-python==2.0.0", - "aws-durable-execution-sdk-python-otel==1.0.0", + "aws-durable-execution-sdk-python==3.0.0", + "aws-durable-execution-sdk-python-otel==2.0.0", ] [tool.hatch.build.targets.wheel] diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml b/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml index 49a64f55..a8901c90 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml @@ -8,7 +8,7 @@ version = "0.0.0" description = "Cross-SDK conformance test handlers for the AWS Durable Execution SDK for Python, exercised by the aws-durable-execution-conformance-tests runner." requires-python = ">=3.11" dependencies = [ - "aws-durable-execution-sdk-python==2.0.0", + "aws-durable-execution-sdk-python==3.0.0", ] [tool.hatch.build.targets.wheel] diff --git a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml index 0691bc32..0b25d4e1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml @@ -21,9 +21,12 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", ] dependencies = [ - # >=2.0.0: first published release carrying the plugin invocation-hook fields + # >=3.0.0: the first release whose `plugins` argument takes factories. + # `workflow_insight()` returns a factory, which core 2.x cannot call, so an + # install resolved against 2.x fails at handler initialization. 3.0.0 also + # carries the invocation-hook fields this plugin reads # (InvocationInfo.execution_input / InvocationEndInfo.execution_result). - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=3.0.0", ] [project.optional-dependencies] diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py index c7c5adad..4faf00f1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "0.0.1" +__version__ = "0.1.0" diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 06a52b67..f03c5da0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -31,6 +31,23 @@ _logger = logging.getLogger("aws_durable_execution_sdk_python_insight") +# Consecutive export-worker deaths tolerated before asynchronous export is +# disabled for good. +# +# A replacement worker is started by whoever is waiting, so a worker that dies on +# every attempt is retried as fast as threads can be created, and each retry +# leaves the waiting drain -- an invocation thread -- exactly where it was. The +# bound converts that unbounded retry into a bounded one. +# +# The bound is not 1, because a single death can come from a transient condition +# that the next attempt would not hit, and disabling instrumentation for the rest +# of the environment's life on one transient is too coarse. A deterministic defect +# reproduces on every attempt, so a small constant separates the two cases. Every +# completed export attempt and every completed flush resets the count, so only +# deaths with no work completed in between accumulate. +_MAX_CONSECUTIVE_WORKER_FAULTS = 3 + + class _ExportState: """One execution's export bookkeeping, and its slot in the export queue. @@ -122,6 +139,10 @@ def __init__(self, exporters: list[InsightExporter]) -> None: self._flush_barrier = 0 self._worker: threading.Thread | None = None self._disabled = False + # Export workers that died without completing any work, counted since the + # last completed export attempt or flush. Only deaths accumulate here, so + # a worker that keeps making progress never approaches the bound. + self._worker_faults = 0 def schedule(self, execution: _ExportState, record: dict[str, Any]) -> None: """Replace this execution's pending snapshot; never runs exporters inline.""" @@ -169,8 +190,10 @@ def drain(self, execution: _ExportState) -> None: Two paths return without exporting or flushing anything, because the permanent ``_disabled`` latch means no record will ever be exported: the latch was already set when this call started, or it is set while this call - is parked. Failing to start the export worker sets that latch, so a drain - that hits a worker-start failure also returns without a flush. + is parked. Two things set that latch, and a drain that meets either + returns without a flush: failing to start the export worker, and an export + worker that has died ``_MAX_CONSECUTIVE_WORKER_FAULTS`` times without + completing any work. """ failed_pending: _Dropped | None = None start_error: Exception | None = None @@ -237,6 +260,34 @@ def drain(self, execution: _ExportState) -> None: # -- internals ------------------------------------------------------------ + def _disable_locked(self) -> _Dropped: + """Latch asynchronous export off for good and surrender everything queued. + + The latch is permanent, so no record the scheduler still holds will ever + be exported: keeping any of them would pin customer objects for the + remaining life of the environment. Empty the queue, which is the + scheduler's only per-execution structure, and hand what came out back to + the CALLER to release once it is outside the lock -- a record can carry + customer objects whose finalizers run arbitrary code, and so can an + execution whose hook state the plugin has already discarded. + + Every parked waiter is woken, because the latch means the export and the + flush it is waiting for are never going to happen. + + Callers hold ``self._condition``. + """ + self._disabled = True + self._worker = None + dropped = [(execution, execution.pending_record) for execution in self._pending] + for execution, _ in dropped: + execution.pending_record = None + self._pending = {} + self._flush_requested = False + self._flush_barrier = 0 + self._flush_in_flight = None + self._condition.notify_all() + return dropped + def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: if self._worker is not None and self._worker.is_alive(): return None, None @@ -249,27 +300,7 @@ def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: try: worker.start() except Exception as exc: # noqa: BLE001 - instrumentation must not escape hooks - self._disabled = True - self._worker = None - # Nothing is retained once the plugin has given up on asynchronous - # export for good: the queue is the scheduler's only per-execution - # structure, so emptying it drops every reference it holds. Both the - # records and the execution objects go back to the CALLER to release - # outside the lock -- a record can carry customer objects whose - # finalizers run arbitrary code, and so can an execution whose hook - # state the plugin has already discarded. - failed_pending = [ - (execution, execution.pending_record) for execution in self._pending - ] - for execution, _ in failed_pending: - execution.pending_record = None - self._pending = {} - self._flush_requested = False - self._flush_barrier = 0 - self._flush_in_flight = None - # Release every waiter; the permanent disable latch means no record - # will ever be exported. - self._condition.notify_all() + failed_pending = self._disable_locked() return failed_pending, exc return None, None @@ -281,20 +312,57 @@ def _blocking_pending_locked(self) -> bool: def _run(self) -> None: # The worker slot must be empty whenever no worker is running, or # _ensure_worker_locked() never starts a replacement and every later - # record sits pending forever. The loop's own exits clear it, but a - # BaseException from a customer exporter -- asyncio.CancelledError is one, - # so an exporter that merely touches asyncio can raise it without writing - # `raise` -- unwinds past them, and a thread that is unwinding still - # reports is_alive(), so the slot would stay occupied by a dead thread. - # Vacate it here, on every exit path, and wake anyone parked so they can - # ask for the replacement. + # record sits pending forever. The loop's own exits clear it, but an + # exception that unwinds out of the loop passes them by, and a thread that + # is unwinding still reports is_alive(), so the slot would stay occupied + # by a dead thread. Vacate it here, on every exit path, and wake anyone + # parked so they can ask for the replacement. + # + # A replacement alone is not enough when the death repeats. The waiter + # that starts the replacement runs the same work again, so a fault the + # work reproduces every time is retried as fast as threads can be + # created, and the drain that keeps starting them never returns: the + # invocation hangs and the environment fills with dead threads. + # _export() and _flush() contain everything a customer exporter can + # raise, so a fault reaching here comes from the scheduler's own code or + # from a failure-reporting call that a customer object subverted, and + # neither is something a retry can be expected to clear. Count + # consecutive faults and give up on asynchronous export at the bound. + faulted = True + dropped: _Dropped | None = None + gave_up = False try: self._run_loop() + faulted = False finally: with self._condition: if self._worker is threading.current_thread(): self._worker = None + if faulted: + self._worker_faults += 1 + if self._worker_faults >= _MAX_CONSECUTIVE_WORKER_FAULTS: + # Releasing the waiters matters more than delivering the + # records. A waiter is an invocation thread inside + # on_invocation_end, so leaving it parked turns an + # instrumentation defect into a stalled customer + # execution; dropping records loses instrumentation data + # only. The drop is reported below, so the scheduler + # never claims delivery it did not make. + dropped = self._disable_locked() + gave_up = True self._condition.notify_all() + # A dropped record can run customer finalizers, so release it outside + # the lock. The exception that brought us here keeps unwinding once + # this block finishes, into the thread's traceback, with nothing + # swallowed. + del dropped + if gave_up: + _logger.warning( + "workflow-insight: export worker died %d times without " + "completing any work; disabling asynchronous export and " + "dropping every record still queued", + self._worker_faults, + ) def _run_loop(self) -> None: while True: @@ -331,16 +399,21 @@ def _run_loop(self) -> None: # nothing will ever export that snapshot again. The bookkeeping # must therefore advance whatever export() did: skip it and # execution.exported_seq never reaches a waiter's want_seq, so a - # drain parked on this execution is never released. _export() - # already contains every Exception, but a BaseException from a - # customer exporter unwinds through here. Count the attempt in a - # finally and let the exception continue out to the wrapper -- - # and into the thread's traceback -- with nothing swallowed. + # drain parked on this execution is never released. Count the + # attempt in a finally so that holds even if _export() raises. + # + # Advancing here means the record was OFFERED to every exporter, + # not that every exporter accepted it. _export() reports each + # exporter's own failure and moves to the next, so no exporter is + # skipped because another one failed, and coverage never stands + # for a delivery that was never attempted. # # The record and its bookkeeping are one object, so there is no # second lookup left to come back empty: publishing cannot miss. + exported = False try: self._export(record) + exported = True finally: # Release the exported record before re-locking: a custom # finalizer may re-enter schedule(). @@ -350,6 +423,11 @@ def _run_loop(self) -> None: if seq > execution.exported_seq: execution.exported_seq = seq execution.exported_at = self._export_count + if exported: + # This worker completed work, so any earlier worker + # death was not the start of a fault the work + # reproduces every time. + self._worker_faults = 0 self._condition.notify_all() continue @@ -367,12 +445,40 @@ def _run_loop(self) -> None: self._flushes_completed += 1 if flush_covers > self._flushed_through: self._flushed_through = flush_covers + self._worker_faults = 0 self._condition.notify_all() with self._condition: if not self._pending and not self._flush_requested: self._worker = None return + # Isolation of one exporter's failure from the others is what both loops below + # exist for, and it holds for every exception type a customer exporter can + # raise, BaseException included. + # + # A BaseException is normally not caught, because KeyboardInterrupt, + # SystemExit and asyncio.CancelledError each mean that the operation the + # current thread is running must stop. None of those three can arrive here + # that way. This is the export worker thread: the interpreter raises + # KeyboardInterrupt only in the main thread, threading discards a SystemExit + # raised in a worker thread, and nothing cancels this thread because nothing + # outside the scheduler knows it exists. A BaseException seen at these call + # sites was therefore raised by the exporter itself, which makes it a report + # of a defective exporter rather than an instruction to this thread. An + # exporter that touches asyncio can raise CancelledError without writing + # `raise`, so the case is reachable without a customer intending it. + # + # Containing it here is what keeps the two guarantees the worker owes. Every + # remaining exporter still receives the record, so one exporter cannot make + # the others miss a snapshot that is then discarded. And the flush the waiting + # drain asked for still completes, so the drain is released by this worker + # instead of by a replacement that runs the same failing exporter and dies the + # same way. + # + # The containment is confined to these two call sites. Nowhere else does the + # scheduler catch a BaseException, and neither loop runs on an invocation + # thread. + def _export(self, record: dict[str, Any]) -> None: for exporter in self._exporters: try: @@ -380,7 +486,7 @@ def _export(self, record: dict[str, Any]) -> None: record, exporter.max_record_size_bytes, exporter.render ) exporter.export(shaped) - except Exception as exc: # noqa: BLE001 - one exporter must not break others + except BaseException as exc: # noqa: BLE001 - one exporter must not break others _logger.warning( "workflow-insight: exporter %s failed: %s", type(exporter).__name__, @@ -391,7 +497,7 @@ def _flush(self) -> None: for exporter in self._exporters: try: exporter.flush() - except Exception as exc: # noqa: BLE001 - one exporter must not break others + except BaseException as exc: # noqa: BLE001 - one exporter must not break others _logger.warning( "workflow-insight: exporter %s flush failed: %s", type(exporter).__name__, diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py index e8410388..5743d5ee 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py @@ -89,13 +89,12 @@ def flush(self) -> None: # pragma: no cover It must return promptly. The invocation that triggered it cannot return until it does, so a slow flush is billed to the customer's invocation. - Failures are isolated: an ``Exception`` is logged, never retried, never - propagated into the execution, and never prevents another exporter from - flushing. A ``BaseException`` (``asyncio.CancelledError`` is one) is not - contained -- it skips the remaining exporters for that flush and ends the - export worker -- but it still never reaches the execution, and the - waiting invocation is released by a replacement worker running the flush - it asked for. + Failures are isolated, whatever is raised. An exception is logged, never + retried, never propagated into the execution, and never prevents another + exporter from flushing. That holds for a ``BaseException`` too + (``asyncio.CancelledError`` is one): the export worker is not the thread + such a signal is ever addressed to, so one raised here is a report of a + defective exporter and is contained exactly like any other failure. """ diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 477a18b1..e5ced2e8 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -9,6 +9,7 @@ from typing import Any from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _MAX_CONSECUTIVE_WORKER_FAULTS, _ExportScheduler, _ExportState, ) @@ -117,6 +118,62 @@ def export(self, record: dict[str, Any]) -> None: super().export(record) +class AlwaysBaseExceptionExportExporter(CaptureExporter): + """Raises a ``BaseException`` out of every ``export()`` call, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self.lock = threading.Lock() + self.export_attempts = 0 + + def export(self, record: dict[str, Any]) -> None: + with self.lock: + self.export_attempts += 1 + raise ExporterBaseException("export exploded") + + def exports(self) -> int: + with self.lock: + return self.export_attempts + + +class AlwaysBaseExceptionFlushExporter(CaptureExporter): + """Raises a ``BaseException`` out of every ``flush()`` call, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self.lock = threading.Lock() + self.flush_attempts = 0 + + def flush(self) -> None: + with self.lock: + self.flush_attempts += 1 + raise ExporterBaseException("flush exploded") + + def flushes(self) -> int: + with self.lock: + return self.flush_attempts + + +def _drain_off_thread( + scheduler: _ArnScheduler, arn: str +) -> tuple[threading.Thread, threading.Event]: + """Start drain() on its own thread and return it with its completion event. + + A regression that parks the drain would block whichever thread called it, so + no test may call drain() on the thread it asserts from. Waiting on the event + with a timeout turns such a regression into a failure instead of a hung run. + """ + returned = threading.Event() + + def drain() -> None: + scheduler.drain(arn) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + return thread, returned + + def test_latest_pending_coalesces_within_one_execution() -> None: exporter = BlockingExporter() scheduler = _ArnScheduler([exporter]) @@ -179,6 +236,131 @@ def drain() -> None: assert _wait_until(lambda: not scheduler._worker_alive()) +def test_always_failing_flush_releases_the_drain_without_a_respawn() -> None: + # A drain is only released by a flush that COMPLETED, and a worker that dies + # is replaced by whoever is waiting. An exporter whose flush() raises a + # BaseException on every call therefore used to make the drain start a + # replacement worker, which ran the same flush and died the same way, without + # bound: the flush was attempted thousands of times per second, thousands of + # threads were created, and the invocation parked on that drain never + # returned. The flush has to be attempted once, the failure reported, and the + # flush counted as completed so the drain returns. + # + # The wait is bounded, so the regression this pins fails the test rather than + # blocking the run. + exporter = AlwaysBaseExceptionFlushExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while flush() kept failing" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + # One attempt, not one per replacement worker. + assert exporter.flushes() == 1 + # The record still reached the exporter, and the failing flush was not retried + # after the drain returned either. + assert exporter.calls == [("export", "terminal")] + assert exporter.flushes() == 1 + + +def test_base_exception_from_one_exporter_never_skips_the_next() -> None: + # Consuming a record from the pending slot is what advances the completion + # bookkeeping, and nothing re-exports a consumed snapshot. A first exporter + # that raised a BaseException used to abort the fan-out loop, so every + # exporter after it missed that record permanently while the drain was + # released as though the record had been delivered. Each exporter's failure + # has to be contained at its own call, so the next exporter still receives the + # record that the bookkeeping counts as offered. + failing = AlwaysBaseExceptionExportExporter() + healthy = CaptureExporter() + scheduler = _ArnScheduler([failing, healthy]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned after export() raised" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert healthy.calls == [("export", "terminal"), ("flush", None)] + # Offered once. The snapshot is gone from the pending slot, so a retry is not + # available and must not be implied. + assert failing.exports() == 1 + + +def test_base_exception_from_one_exporters_flush_never_skips_the_next() -> None: + # The same containment at the flush call. A first exporter whose flush() + # raises a BaseException must not stop a later exporter from flushing, and the + # flush must still count as completed so the waiting drain is released. + failing = AlwaysBaseExceptionFlushExporter() + healthy = CaptureExporter() + scheduler = _ArnScheduler([failing, healthy]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while flush() kept failing" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert healthy.calls == [("export", "terminal"), ("flush", None)] + assert failing.flushes() == 1 + + +class _FaultingScheduler(_ArnScheduler): + """Kills every export worker with a ``BaseException`` before it does any work. + + Stands in for a fault in the scheduler's own code rather than in an exporter: + exporter failures are contained at the exporter call, so they can no longer + reach the worker's exit path, and this is the only way left to drive it. + """ + + def __init__(self, exporters: list[Any]) -> None: + super().__init__(exporters) + self.runs = 0 + + def _run_loop(self) -> None: + with self._condition: + self.runs += 1 + raise ExporterBaseException("worker exploded") + + +def test_worker_deaths_stop_at_the_bound_instead_of_respawning_forever() -> None: + # A waiting drain starts a replacement worker for every worker that dies, so a + # fault the worker reproduces on every attempt is retried as fast as threads + # can be created and the drain never returns. Releasing the waiter matters + # more than delivering the records, because the waiter is an invocation + # thread: bound the replacements, then latch asynchronous export off, which + # wakes every waiter and drops what is queued. + exporter = CaptureExporter() + scheduler = _FaultingScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while the worker kept dying" + thread.join(5.0) + assert not thread.is_alive() + + with scheduler._condition: + assert scheduler.runs == _MAX_CONSECUTIVE_WORKER_FAULTS + assert scheduler._worker_faults == _MAX_CONSECUTIVE_WORKER_FAULTS + # The latch is what released the drain, and it retains nothing. + assert scheduler._disabled + assert scheduler._pending == {} + assert scheduler._flush_requested is False + assert scheduler._flush_in_flight is None + # No further worker is started once the latch is set, so the count cannot + # creep up after the drain returned. + scheduler.schedule(ARN_B, _record("after-the-latch")) + scheduler.drain(ARN_B) + with scheduler._condition: + assert scheduler.runs == _MAX_CONSECUTIVE_WORKER_FAULTS + assert exporter.calls == [] + + def test_drain_flushes_after_export() -> None: capture = CaptureExporter() scheduler = _ArnScheduler([capture]) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py new file mode 100644 index 00000000..e9130b3b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Packaging checks for the Workflow Insight plugin. + +The plugin and the core SDK are separate distributions, so pip resolves their +versions independently. A declared bound that admits a core release without the +contract this plugin uses is an install that succeeds and then fails at handler +initialization, which is why the bound is asserted here rather than left to +review. +""" + +from __future__ import annotations + +import re +import tomllib +from pathlib import Path + +from packaging.version import Version + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] +CORE_DISTRIBUTION = "aws-durable-execution-sdk-python" + + +def _core_version() -> str: + """The core SDK version this repository builds, read from its source. + + Read from the file rather than imported, so the check keeps describing this + repository even when a published core is installed alongside these sources. + """ + about = ( + REPOSITORY_ROOT + / "packages" + / "aws-durable-execution-sdk-python" + / "src" + / "aws_durable_execution_sdk_python" + / "__about__.py" + ).read_text() + match = re.search(r'^__version__ = "([^"]+)"', about, re.MULTILINE) + assert match is not None, "core __about__.py has no __version__ assignment" + return match.group(1) + + +def _core_dependency_lower_bound() -> str: + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + + bounds = [ + dependency.removeprefix(CORE_DISTRIBUTION + ">=") + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION + ">=") + ] + assert len(bounds) == 1, f"expected one {CORE_DISTRIBUTION} bound, got {bounds}" + return bounds[0] + + +def _major(version: str) -> int: + return int(version.split(".", 1)[0]) + + +def test_core_dependency_bound_matches_the_core_major_in_this_repository() -> None: + """The declared bound must not admit a core major that predates the factory contract. + + ``workflow_insight()`` returns a plugin factory, and only the core major that + introduced factories calls it. A lower bound naming an earlier major is a + resolution pip accepts and that then fails at handler initialization, so the + bound tracks the core major this repository builds. The bound may lag within + that major -- a later core minor still satisfies the contract -- which is why + only the major is compared and the bound is required not to exceed the core + version. + """ + core_version = _core_version() + lower_bound = _core_dependency_lower_bound() + + assert _major(lower_bound) == _major(core_version) + assert Version(lower_bound) <= Version(core_version) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 5efd8947..0c9c07a1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -462,7 +462,7 @@ invocation the environment runs, including concurrent ones. ## Requirements - Python >= 3.11 -- `aws-durable-execution-sdk-python` >= 2.0.0 +- `aws-durable-execution-sdk-python` >= 3.0.0 - An ADOT/community OpenTelemetry Lambda layer, or the `standalone` extra ## License diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index 60e835af..de1d940d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,7 +22,11 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + # >=3.0.0: the first release whose `plugins` argument takes factories. + # DurableInstrumentationPluginProvider was removed in 3.0.0, and this package's + # entry points resolve to factories, so core 2.x accepts the install and then + # fails at handler initialization. + "aws-durable-execution-sdk-python>=3.0.0", ] [project.entry-points."aws_durable_execution.plugins"] diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py index fc21e413..82fa7f7e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "1.0.0" +__version__ = "2.0.0" diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 22b41638..59ef0a13 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -1,10 +1,14 @@ +import re import tomllib from pathlib import Path +from packaging.version import Version + PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.0.0" +CORE_DISTRIBUTION = "aws-durable-execution-sdk-python" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=3.0.0" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -80,3 +84,76 @@ def test_pypi_compatibility_environment_uses_compatible_core_sdk() -> None: ]["test-pypi-otel"]["dependencies"] assert CORE_DEPENDENCY in dependencies + + +def _core_version() -> str: + """The core SDK version this repository builds, read from its source. + + Read from the file rather than imported. The ``test-pypi-otel`` environment + installs a *published* core alongside this package's source, so an import + would report that release's version and the checks below would stop saying + anything about this repository's version story. + """ + about = ( + REPOSITORY_ROOT + / "packages" + / "aws-durable-execution-sdk-python" + / "src" + / "aws_durable_execution_sdk_python" + / "__about__.py" + ).read_text() + match = re.search(r'^__version__ = "([^"]+)"', about, re.MULTILINE) + assert match is not None, "core __about__.py has no __version__ assignment" + return match.group(1) + + +def _core_dependency_lower_bound() -> str: + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + bounds = [ + dependency.removeprefix(CORE_DISTRIBUTION + ">=") + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION + ">=") + ] + assert len(bounds) == 1, f"expected one {CORE_DISTRIBUTION} bound, got {bounds}" + return bounds[0] + + +def _major(version: str) -> int: + return int(version.split(".", 1)[0]) + + +def test_core_dependency_bound_matches_the_core_major_in_this_repository() -> None: + """The declared bound must not admit a core major that predates this plugin's contract. + + This package's entry points resolve to plugin factories, and the ``plugins`` + argument only accepts factories from the core major that introduced them. A + lower bound naming an earlier major is a resolution pip accepts and that then + fails at handler initialization, so the bound has to track the core major this + repository builds. The bound may lag within that major -- a later core minor + still satisfies the contract -- which is why only the major is compared and + the bound is required not to exceed the core version. + """ + core_version = _core_version() + lower_bound = _core_dependency_lower_bound() + + assert _major(lower_bound) == _major(core_version) + assert Version(lower_bound) <= Version(core_version) + + +def test_layer_sdk_pin_matches_the_core_version_in_this_repository() -> None: + """The OTel layer pin selects the core wheel bundled into the published layer. + + A combined SDK and OTel release fails outright when the pin disagrees with the + released SDK version, and an OTel-only release downloads exactly the pinned + version from PyPI. A stale pin therefore either blocks the release or ships a + layer whose core cannot run this plugin, so the pin tracks the core version + this repository builds. + """ + metadata_path = REPOSITORY_ROOT / ".github" / "lambda-layer-publish.toml" + + with metadata_path.open("rb") as metadata_file: + pinned_version = tomllib.load(metadata_file)["layer"]["sdk-version"] + + assert pinned_version == _core_version() diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py index 82fa7f7e..92b4bb99 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "2.0.0" +__version__ = "3.0.0" diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 1e7b9578..4fb80fbc 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -200,6 +200,16 @@ def durable_execution( plugin_host = PluginHost(load_configured_plugins(plugins)) @plugin_host.handle_durable_output + # The metadata of whatever function the host wrapper wraps becomes the + # decorated handler's own, because the host wrapper applies + # functools.wraps() to it. This function takes a third argument the returned + # handler does not accept, so without this line inspect.signature() on the + # handler advertises a required `plugin_executor` parameter, and a + # signature-aware runtime or test harness rejects or misinvokes a handler + # that in fact takes (event, context). Copying the user function's metadata + # here puts it at the head of the chain the host wrapper then extends, so + # the handler reports the user function's signature, name and docstring. + @functools.wraps(func) def wrapper( event: Any, context: LambdaContext, plugin_executor: PluginExecutor ) -> MutableMapping[str, Any]: diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index 57d1d969..ac70d520 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -85,6 +85,38 @@ def _load_factory( return cast(DurableInstrumentationPluginFactory, factory) +def _validate_explicit_factories( + explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, +) -> list[DurableInstrumentationPluginFactory]: + """Check that every explicitly passed plugin entry is callable. + + Each entry is called once per invocation to build that invocation's plugin + instance. An entry that is not callable can never be called, so + :meth:`PluginExecutor._create_plugins` raises ``TypeError`` on every + invocation, logs it and continues without that plugin -- telemetry is lost + for the lifetime of the function, and nothing fails. Raising here converts + that into one configuration failure while the handler is being initialized. + The position is named because a caller passing several entries cannot + otherwise tell which one is wrong. + + A plugin *class* is callable and stays valid: calling it constructs an + instance, so ``plugins=[MyPlugin]`` is a factory whenever ``MyPlugin`` + accepts the info argument. Only a plugin *instance*, or any other + non-callable value, is rejected. + """ + factories = list(explicit_plugins or []) + for index, factory in enumerate(factories): + if not callable(factory): + raise PluginLoadError( + f"Durable instrumentation plugin at plugins[{index}] must be a " + "callable plugin factory taking an InvocationStartInfo, but is " + f"{_qualified_type_name(factory)}. Pass a factory rather than a " + "plugin instance, for example " + "plugins=[lambda info: MyPlugin(...)]." + ) + return factories + + def load_configured_plugins( explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, *, @@ -105,7 +137,7 @@ def load_configured_plugins( build the same plugin type will now both be registered. """ - resolved_factories = list(explicit_plugins or []) + resolved_factories = _validate_explicit_factories(explicit_plugins) resolved_environment = os.environ if environment is None else environment plugin_names = _parse_configured_plugin_names(resolved_environment) if not plugin_names: diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 0cba5fb3..bf19b8d2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -1,6 +1,7 @@ """Tests for execution.""" import datetime +import inspect import json import threading import time @@ -4051,3 +4052,82 @@ def test_handler(event: Any, context: DurableContext) -> dict: # endregion Plugin Integration Tests + + +# region Handler Metadata Tests + + +def _bare_decorated_handler(): + """The ``@durable_execution`` form, desugared so the test keeps both objects.""" + + def test_handler(event: Any, context: DurableContext) -> dict: + """Handler docstring.""" + return {"result": "success"} + + return durable_execution(test_handler), test_handler + + +def _plugin_decorated_handler(): + """The ``@durable_execution(plugins=[...])`` form, desugared the same way. + + This form takes a different path through the decorator: the first call + returns a ``functools.partial``, which the second call applies to the user + function. + """ + + def test_handler(event: Any, context: DurableContext) -> dict: + """Handler docstring.""" + return {"result": "success"} + + decorator = durable_execution(plugins=[plugin_factory(_RecordingPlugin())]) + return decorator(test_handler), test_handler + + +@pytest.mark.parametrize( + "build_handler", + [_bare_decorated_handler, _plugin_decorated_handler], + ids=["bare", "with_plugins"], +) +def test_decorated_handler_signature_is_event_and_context(build_handler) -> None: + """The decorated handler accepts exactly the two Lambda handler arguments. + + A Lambda runtime, or a test harness that inspects a handler before calling + it, reads ``inspect.signature()``. Any parameter reported there that the + callable does not accept makes the handler look uninvokable, or invites a + caller to pass an argument that raises. The SDK's own invocation body takes a + third argument -- this invocation's ``PluginExecutor`` -- so this pins the + public signature to the two arguments the handler really takes, for both + decorator forms. + """ + handler, _user_function = build_handler() + + signature = inspect.signature(handler) + + assert list(signature.parameters) == ["event", "context"] + assert "plugin_executor" not in signature.parameters + # Binding is the operation a signature-aware caller actually performs. + signature.bind({}, _make_lambda_context()) + + +@pytest.mark.parametrize( + "build_handler", + [_bare_decorated_handler, _plugin_decorated_handler], + ids=["bare", "with_plugins"], +) +def test_decorated_handler_reports_user_function_metadata(build_handler) -> None: + """The handler identifies itself as the user's function, not as SDK internals. + + Logging, ``help()`` and error messages read ``__name__`` and ``__doc__``. The + SDK wraps the user function twice, so without the user function's metadata at + the head of the chain the handler names an internal wrapper instead. + ``inspect.unwrap()`` reaching the user function is what makes + ``inspect.signature()`` report that function's parameters. + """ + handler, user_function = build_handler() + + assert handler.__name__ == user_function.__name__ + assert handler.__doc__ == user_function.__doc__ + assert inspect.unwrap(handler) is user_function + + +# endregion Handler Metadata Tests diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index a5153132..7936bd13 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -439,3 +439,118 @@ def __init__(self, info: InvocationStartInfo) -> None: plugin = result[0](INVOCATION_START_INFO) assert isinstance(plugin, _InfoAwarePlugin) assert plugin.info is INVOCATION_START_INFO + + +def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: + """A plugin instance in ``plugins`` fails configuration, not every invocation. + + An instance is not callable, so the per-invocation factory call raises + ``TypeError``, which the executor logs and swallows -- the plugin silently + never runs. The position is asserted because a caller passing several entries + has no other way to tell which one is wrong. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins( + [_plugin_a_factory, _PluginB(), _plugin_b_factory], # type: ignore[list-item] + environment={}, + ) + + assert "plugins[1]" in str(error.value) + assert "must be a callable plugin factory" in str(error.value) + assert "_PluginB" in str(error.value) + + +@pytest.mark.parametrize( + ("non_callable", "expected_type_name"), + [ + (_PluginA(), "_PluginA"), + (object(), "builtins.object"), + ("not-a-factory", "builtins.str"), + (None, "builtins.NoneType"), + ], +) +def test_explicit_non_callable_entries_are_rejected( + non_callable: object, + expected_type_name: str, +) -> None: + """Callability is the only property checkable without building an instance.""" + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([non_callable], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert expected_type_name in str(error.value) + + +def test_explicit_plugin_class_is_accepted_as_a_factory() -> None: + """Calling a class constructs an instance, so a class is a factory in Python. + + This is the language difference against the TypeScript SDK, which rejects a + class because calling one there throws. Rejecting classes here would break + ``plugins=[MyPlugin]``, which works whenever ``__init__`` takes the info. + """ + + class _InfoAwarePlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + result = load_configured_plugins([_InfoAwarePlugin], environment={}) + + assert result == [_InfoAwarePlugin] + plugin = result[0](INVOCATION_START_INFO) + assert isinstance(plugin, _InfoAwarePlugin) + assert plugin.info is INVOCATION_START_INFO + + +def test_explicit_plain_function_is_accepted_as_a_factory() -> None: + result = load_configured_plugins([_plugin_a_factory], environment={}) + + assert result == [_plugin_a_factory] + assert isinstance(result[0](INVOCATION_START_INFO), _PluginA) + + +def test_explicit_callable_object_is_accepted_as_a_factory() -> None: + """A ``__call__`` instance is the documented way to hold handler-lifetime state.""" + + class _CallableFactory: + def __init__(self) -> None: + self.calls: list[InvocationStartInfo] = [] + + def __call__(self, info: InvocationStartInfo) -> _PluginA: + self.calls.append(info) + return _PluginA() + + factory = _CallableFactory() + + result = load_configured_plugins([factory], environment={}) + + assert result == [factory] + assert isinstance(result[0](INVOCATION_START_INFO), _PluginA) + assert factory.calls == [INVOCATION_START_INFO] + + +def test_explicit_entries_are_validated_before_entry_points_are_imported() -> None: + """Explicit entries are checked first, so a valid provider is not imported. + + Importing a provider runs third-party module code. A configuration that is + already invalid should fail before that happens, and the failure should name + the invalid entry rather than whatever the import did. + """ + entry_point = _FakeEntryPoint("a", _plugin_a_factory) + + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ) as entry_points, + patch.object( + _FakeEntryPoint, "load", side_effect=AssertionError("must not import") + ) as load, + pytest.raises(PluginLoadError, match=r"plugins\[0\]"), + ): + load_configured_plugins( + [_PluginA()], # type: ignore[list-item] + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + entry_points.assert_not_called() + load.assert_not_called() diff --git a/pyproject.toml b/pyproject.toml index 428d680e..9c2e4972 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" [tool.hatch.envs.test-pypi-otel] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=3.0.0", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", From f8b4f31cd6b59232656902dd87141c21614a6951 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Thu, 17 Sep 2026 22:12:42 -0700 Subject: [PATCH 06/28] fix(otel): stamp a log record only from a live claim Two reviewers found two orderings where the log filter attributed a record to the wrong execution. Both went through the same fallback, which resolved a record to the only open invocation when the emitting thread carried no live claim. The first ordering is a stale claim. Invocation A claims thread T. A ends, and unbind_invocation cannot reset T's ContextVar, because only the thread that set a ContextVar can reset it. B starts and is the only open invocation. A thread the customer started under A still carries A's claim, so the claim is stale, the fallback runs, and A's record is stamped with B's trace. Measured: the record carried B's trace and span exactly. The second ordering is a missing claim. A is open. B emits a log on its own wrapper thread before B reaches on_invocation_start, so B's thread has no claim and the registry holds only A. Measured: B's record carried A's trace and span exactly. Both outcomes are the one the design chose against, which is that an unattributed record is better than a record attributed to another customer's execution. So the fallback is deleted and a live claim is now required. One reviewer proposed keeping the fallback for threads that never held a claim, which fixes the first ordering and leaves the second. Requiring a claim fixes both. Deleting it is affordable because the threads that run customer code all carry a claim. The core SDK copies the invocation's context into the thread running the handler body, and the operation hooks claim the map and parallel branch threads before the branch body runs. Measured over one execution driving a step, a parallel, a map and a customer-started thread: every record from user code carried a claim. Two threads now go unstamped whatever the invocation count. The SDK's background checkpointing thread is submitted without the invocation's context deliberately. A thread customer code starts itself does not inherit the starting thread's context. The README states both rather than claiming every record is stamped. A second defect made the registry untrustworthy, so it is fixed here too. on_invocation_end ended the spans and exported the workflow span before releasing the invocation scope, with no finally. Both calls run customer tracer and span-processor code, which can raise. The plugin executor contains that exception, so the invocation survives, but the release was skipped: the plugin stayed registered as open for the life of the environment and its OTel context stayed attached. Measured with a span processor whose on_end raises: the plugin remained in the open registry for both plugins, and for ExecutionOtelPlugin a later record on the still-claimed thread was stamped with the finished invocation. The release and the flush now run in a finally, in both plugins. force_flush is contained there as well, because an exception escaping a finally would replace the exception that ended the invocation and hide its cause. Two documentation corrections from the same review. The factory docs presented a plugin class as the example factory, which encourages __init__(self, info) to do setup work, and CONTRIBUTING.md:246 asks for light constructors. A class is still a valid factory, because calling a class constructs an instance, so only the emphasis changes: the recommended forms are now a lambda or a classmethod, with class-as- factory noted as permitted. And the PluginHost docstring cited two cross-SDK symbols as if they were on the default branches. Checked against those repositories: the JS function there is createPluginRunner, and createInvocationPluginRunner exists only in aws/aws-durable- execution-sdk-js#924; the Java PluginRunner does exist on main, but takes plugin instances, and aws/aws-durable-execution-sdk-java#721 is what makes it per-invocation. The docstring now says that. --- .../README.md | 18 +- .../execution_plugin.py | 44 ++++- .../invocation_plugin.py | 42 ++++- .../log_filter.py | 48 ++++-- .../tests/test_log_filter.py | 157 +++++++++++++++++- .../README.md | 11 ++ .../plugin.py | 29 +++- .../plugin_discovery.py | 9 +- 8 files changed, 306 insertions(+), 52 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 0c9c07a1..81304132 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -331,20 +331,22 @@ operation, and holds when several invocations run concurrently in one environment (Lambda Managed Instances): the plugin claims the invocation thread at invocation start and the SDK runs the handler body in a copy of that thread's context, so each record resolves to the invocation that emitted it rather than -to whichever invocation started last. +to whichever invocation started last. A branch of a `map` or `parallel` runs on +a pool thread the invocation's context was not copied into, and the plugin +claims that thread from the hooks that run on it before your branch body does. Two cases are left unstamped, so any log formatter or schema must treat the fields as optional: - No invocation is open — for example during environment initialization or teardown. -- The record is emitted on a thread that carries no invocation claim, while more - than one invocation is open in the environment. A thread your code starts - itself is such a thread, since Python does not copy context into a new thread, - as is the SDK's background checkpointing thread. With exactly one invocation - open, such a record is correlated to it. With several open there is no way to - tell which one it belongs to, and an uncorrelated record is preferred over one - attributed to another execution. +- The record is emitted on a thread that carries no invocation claim. A thread + your code starts itself is such a thread, since Python does not copy context + into a new thread, as is the SDK's background checkpointing thread. The number + of invocations open in the environment does not change this: attributing an + unclaimed record to the single open invocation would be wrong whenever the + emitting thread belongs to a different invocation, and an uncorrelated record + is preferred over one carrying another execution's trace. ## Verification 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 3663b3be..4654a826 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 @@ -654,11 +654,28 @@ def _start_invocation_span(self, info: InvocationStartInfo) -> None: self._set_span(_INVOCATION_KEY, self._invocation_span) def on_invocation_end(self, info: InvocationEndInfo) -> None: + """End this invocation's spans, then release its scope and flush. + + Ending a span and exporting the Workflow span call the configured tracer + and span processors, which are customer-supplied and can raise. The + invocation must still be released: until it is, the log filter counts it + as open and the OTel scope this plugin attached at invocation start stays + current on a warm environment's thread. The release and the flush + therefore run in a ``finally``. + """ logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: self._release_invocation_scope() return + try: + self._end_invocation_spans(info) + finally: + self._release_invocation_scope() + self._force_flush() + + def _end_invocation_spans(self, info: InvocationEndInfo) -> None: + """End this invocation's open spans and export the Workflow span.""" # End the invocation span regardless of terminal status. Record the # invocation status and map it to a span status: # SUCCEEDED/PENDING -> OK (this invocation did its work, whether it @@ -692,13 +709,20 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._release_invocation_scope() + def _force_flush(self) -> None: + """Flush pending spans, containing any error the flush raises. - if hasattr(self._provider, "force_flush"): - try: - self._provider.force_flush() - except Exception: # noqa: BLE001 - logger.exception("force_flush failed at invocation end") + A flush calls the configured span processors and exporters, which are + customer-supplied and can raise. This runs in a ``finally`` block, where + an escaping exception would replace the exception that ended the + invocation and hide its cause, so the error is logged and dropped. + """ + if not hasattr(self._provider, "force_flush"): + return + try: + self._provider.force_flush() + except Exception: # noqa: BLE001 + logger.exception("force_flush failed at invocation end") def _release_invocation_scope(self) -> None: """Release what this invocation attached, and stop instrumenting. @@ -713,9 +737,11 @@ def _release_invocation_scope(self) -> None: detached here or it would stay current on a warm environment's thread after the invocation returns. * The log filter's record of open invocations is process-global, so this - invocation must be removed from it. Until it is, a record emitted on an - unclaimed thread could still be correlated to this finished - invocation's spans. + invocation must be removed from it. A thread this invocation claimed + keeps that claim, because a context variable can only be reset by the + thread that set it, so until the invocation is removed a record emitted + 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 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 686f87bd..923733f8 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 @@ -715,12 +715,29 @@ def _export_workflow_span(self, info: InvocationEndInfo) -> None: workflow_span.end() def on_invocation_end(self, info: InvocationEndInfo) -> None: - """Called at the end of each invocation. Ends the invocation span and flushes.""" + """Called at the end of each invocation. Ends the invocation span and flushes. + + Ending a span and exporting the Workflow span call the configured tracer + and span processors, which are customer-supplied and can raise. The + invocation must still be released: until it is, the log filter counts it + as open and any OTel scope this plugin attached stays current on a warm + environment's thread. The release and the flush therefore run in a + ``finally``. + """ logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: self._release_invocation_scope() return + try: + self._end_invocation_spans(info) + finally: + self._release_invocation_scope() + # Flush before Lambda freeze. + self._force_flush() + + def _end_invocation_spans(self, info: InvocationEndInfo) -> None: + """End this invocation's open spans and export the Workflow span.""" # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. with self._operation_spans_lock: @@ -766,11 +783,20 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._release_invocation_scope() + def _force_flush(self) -> None: + """Flush pending spans, containing any error the flush raises. - # Flush before Lambda freeze - if hasattr(self._provider, "force_flush"): + A flush calls the configured span processors and exporters, which are + customer-supplied and can raise. This runs in a ``finally`` block, where + an escaping exception would replace the exception that ended the + invocation and hide its cause, so the error is logged and dropped. + """ + if not hasattr(self._provider, "force_flush"): + return + try: self._provider.force_flush() + except Exception: # noqa: BLE001 + logger.exception("force_flush failed at invocation end") def _release_invocation_scope(self) -> None: """Release what this invocation attached, and stop instrumenting. @@ -785,9 +811,11 @@ def _release_invocation_scope(self) -> None: detached here or it would stay current on a warm environment's thread after the invocation returns. * The log filter's record of open invocations is process-global, so this - invocation must be removed from it. Until it is, a record emitted on an - unclaimed thread could still be correlated to this finished - invocation's spans. + invocation must be removed from it. A thread this invocation claimed + keeps that claim, because a context variable can only be reset by the + thread that set it, so until the invocation is removed a record emitted + 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 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 68162ea8..87293990 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 @@ -38,14 +38,25 @@ code therefore resolve to their own invocation, including the first statement of the handler, before any durable operation has claimed the thread directly. - - A record emitted on a thread that carries no claim resolves to the one - open invocation, if exactly one is open. Threads the SDK does not submit - the invocation's context into land here: a thread customer code starts - itself, since Python does not copy context into a new thread, and the - SDK's background checkpointing thread. - - If several invocations are open and the thread carries no claim, the record - is left unstamped. An unattributed record is a smaller defect than one - attributed to another customer execution. + - A branch of a map or parallel runs on a pool thread the invocation's + context was not copied into, so the plugins claim that thread from the + hooks that run on it before the branch body does. + - 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. + +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 +keeps that claim after A ends, because a :class:`contextvars.ContextVar` can only +be reset by the thread that set it, so a record A's thread emits while B is the +only open invocation would be stamped with B's trace. A record emitted on +invocation B's own thread before B reaches its invocation-start hook carries no +claim at all, so while A is the only open invocation it would be stamped with A's +trace. Both stamp one customer's execution onto another's, which is a larger +defect than an unattributed record, so a live claim is required and the count of +open invocations is never consulted. Reading the active span straight from the OTel context (as the Java plugin's static ``MdcSpanEnricher`` does) is not sufficient here: the invocation span is @@ -137,13 +148,26 @@ def unbind_invocation(provider: _SpanContextProvider) -> None: def _resolve_provider() -> _SpanContextProvider | None: - """Return the invocation to correlate a record emitted right here against.""" + """Return the invocation to correlate a record emitted right here against. + + A claim is only honoured while the invocation naming it is still open. A + :class:`contextvars.ContextVar` can only be reset by the thread that set it, + so ``unbind_invocation`` leaves the claim in place on every other thread the + invocation claimed; without the liveness check a pooled thread would keep + correlating records to a finished invocation. + + A thread with no live claim resolves to nothing, even when exactly one + invocation is open. Deciding by count would attribute the record to that + invocation, and two orderings make that the wrong one: a thread still + carrying a finished invocation's claim, and an invocation's own thread that + has not reached its invocation-start hook yet. + """ claimed = _current_invocation.get() + if claimed is None: + return None with _registry_lock: - if claimed is not None and claimed in _open_invocations: + if claimed in _open_invocations: return claimed - if len(_open_invocations) == 1: - return next(iter(_open_invocations)) return None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index f46d5999..583bcdcf 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -22,7 +22,7 @@ from opentelemetry import trace from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( NonRecordingSpan, @@ -35,6 +35,7 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( _to_otel_trace_id, ) +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, @@ -262,18 +263,22 @@ def invocation(owner: str) -> None: assert stamped["b"] == own["b"] -def test_record_on_an_unclaimed_thread_uses_the_only_open_invocation(): - """A thread carrying no claim still correlates to the one open invocation. +def test_record_on_an_unclaimed_thread_is_not_stamped_with_the_open_invocation(): + """A thread carrying no claim is not correlated, even to a lone invocation. A thread that carries no claim -- one customer code started itself, since ``threading.Thread`` does not copy the starting thread's context -- has no - invocation of its own. With a single invocation open there is no ambiguity to - resolve, so the record is correlated to it. + invocation of its own. Resolving it to the single open invocation would be + right only when that invocation is the one that emitted the record, and two + orderings make it the wrong one: a thread still carrying a finished + invocation's claim, and an invocation's own thread before its start hook has + run. Correlation for such records is given up so neither ordering can stamp + one execution's record with another's trace. """ plugin, _ = _create_plugin(enrich_logger=False) plugin.on_invocation_start(_invocation_start_info()) try: - expected = _own_identifiers(plugin) + open_identifiers = _own_identifiers(plugin) stamped: list[tuple[str | None, str | None]] = [] def emit() -> None: @@ -285,11 +290,99 @@ def emit() -> None: worker.start() worker.join(timeout=10) - assert stamped == [expected] + assert stamped == [(None, None)] + assert stamped[0] != open_identifiers finally: plugin.on_invocation_end(_invocation_end_info()) +def test_a_stale_claim_is_not_resolved_to_the_only_open_invocation(): + """A thread still claimed by a finished invocation is not given a live one. + + ``unbind_invocation`` cannot reset the claim on any thread but its own, so a + thread the first invocation claimed still names that invocation after it + ends. A second invocation is then the only open one. Resolving by the number + of open invocations would stamp the first invocation's record with the + second's trace, which is one customer execution's identifiers on another's + record. + """ + first, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + first_identifiers = _own_identifiers(first) + + # A worker running in a copy of the claiming thread's context, which is what + # the SDK submits the handler body as, so the claim reaches it. + claimed_context = contextvars.copy_context() + second_is_open = threading.Event() + stamped: list[tuple[str | None, str | None]] = [] + + def emit() -> None: + assert second_is_open.wait(timeout=10) + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + + worker = threading.Thread( + target=claimed_context.run, args=(emit,), name="claimed-by-first" + ) + worker.start() + try: + first.on_invocation_end(_invocation_end_info(suffix="first")) + second, _ = _create_plugin(enrich_logger=False) + second.on_invocation_start(_invocation_start_info(suffix="second")) + try: + second_identifiers = _own_identifiers(second) + second_is_open.set() + worker.join(timeout=10) + + assert len(stamped) == 1 + assert stamped[0] != second_identifiers + assert stamped[0] != first_identifiers + assert stamped[0] == (None, None) + finally: + second.on_invocation_end(_invocation_end_info(suffix="second")) + finally: + second_is_open.set() + worker.join(timeout=10) + + +def test_a_record_emitted_before_its_invocation_binds_is_not_given_the_open_one(): + """A record emitted before its own invocation binds is not correlated. + + An invocation emits records on its own thread before its invocation-start + hook runs -- while the SDK fetches initial state, for instance -- and that + thread carries no claim yet. Another invocation can be open at that moment. + Resolving by the number of open invocations would stamp the starting + invocation's record with the open invocation's trace. + """ + first, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + second, _ = _create_plugin(enrich_logger=False) + stamped: list[tuple[str | None, str | None]] = [] + try: + first_identifiers = _own_identifiers(first) + + def start_second_invocation() -> None: + # This thread belongs to the second invocation, which has not bound + # itself yet, so nothing here carries a claim. + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + second.on_invocation_start(_invocation_start_info(suffix="second")) + + worker = threading.Thread(target=start_second_invocation, name="second-wrapper") + worker.start() + worker.join(timeout=10) + + assert len(stamped) == 1 + assert stamped[0] != first_identifiers + assert stamped[0] != _own_identifiers(second) + assert stamped[0] == (None, None) + finally: + first.on_invocation_end(_invocation_end_info(suffix="first")) + second.on_invocation_end(_invocation_end_info(suffix="second")) + + def test_context_propagated_into_a_worker_resolves_the_claiming_invocation(): """A worker started from a copy of the claiming thread's context resolves it. @@ -484,6 +577,56 @@ def test_finished_invocation_does_not_correlate_later_records(): assert not hasattr(record, "spanId") +class _SpanProcessorFailingOnEnd(SpanProcessor): + """Raises when a span ends, standing in for customer span-processor code.""" + + def on_start(self, span, parent_context=None) -> None: + pass + + def on_end(self, span) -> None: + raise RuntimeError("span processor on_end failed") + + +@pytest.mark.parametrize( + "plugin_class", + [InvocationOtelPlugin, ExecutionOtelPlugin], + ids=lambda c: c.__name__, +) +def test_invocation_is_released_when_span_shutdown_fails(plugin_class): + """A failing span shutdown still releases the invocation. + + Ending a span calls the configured span processors, which are customer code + and can raise. The plugin executor contains that exception, so the invocation + survives it. An invocation left registered would stay open for the life of + the environment, and a thread still carrying its claim would keep correlating + records to its finished spans. + """ + provider = TracerProvider() + provider.add_span_processor(_SpanProcessorFailingOnEnd()) + plugin = plugin_class( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + plugin.on_invocation_start(_invocation_start_info()) + + with pytest.raises(RuntimeError, match="span processor on_end failed"): + plugin.on_invocation_end(_invocation_end_info()) + + assert plugin not in log_filter_module._open_invocations + # This thread carried the claim, so it is the thread a later record would be + # mis-attributed on. + assert log_filter_module._resolve_provider() is not plugin + record = _make_record() + OtelContextLogFilter().filter(record) + assert _stamped(record) == (None, None) + # The OTel context stack belongs to the thread, so a scope the plugin left + # attached would stay current after the invocation returned. + assert plugin._context_tokens == {} + + def test_install_log_filter_attaches_to_handlers(): """install_log_filter adds the filter to each handler on the target logger.""" target = logging.getLogger("test.install") diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index a241e23b..d542c8f6 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -49,6 +49,17 @@ callable taking the invocation's `InvocationStartInfo` and returning a instance it returns serves that one invocation only and can hold per-execution state in ordinary attributes. +Write the factory as a function, a `lambda`, or a `@classmethod`, and construct +the plugin inside it. A constructor should only assign fields, so a plugin class +used directly as a factory invites setup work into `__init__`. A plugin class is +callable and is accepted as a factory whenever its `__init__` takes the info, but +prefer the explicit form: + +```python +plugins=[lambda info: AuditPlugin(sink)] # construct explicitly +plugins=[AuditPlugin.create] # a @classmethod factory +``` + Provider packages expose such a factory: ```python 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 0d54de67..e24bb13a 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 @@ -465,9 +465,23 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: alias is also the more permissive of the two, because ``Callable`` parameters are positional-only -- a factory may name its parameter whatever reads best (``lambda info: ...``, ``def build(invocation): ...``), where a ``__call__`` -Protocol would pin that name. Anything callable satisfies it: a lambda, a -module-level function, a ``functools.partial``, or a class whose ``__init__`` -takes the info. +Protocol would pin that name. + +Prefer a factory that constructs the plugin explicitly:: + + plugins=[lambda info: MyPlugin(exporter)] + plugins=[MyPlugin.create] # a @classmethod factory + +A constructor should only assign fields, so setup work that can fail or that +reads the environment belongs in a factory rather than in ``__init__`` (see +``CONTRIBUTING.md``, "Initialization and conversion"). An explicit factory is +where that work goes, and it also lets the plugin take its own collaborators +rather than deriving them from the hook info. + +Anything callable satisfies the alias: a lambda, a module-level function, a +``functools.partial``, a ``@classmethod``, or a plugin class itself, since calling +a class in Python constructs an instance. ``plugins=[MyPlugin]`` is therefore +permitted whenever ``MyPlugin.__init__`` takes the info, and it stays permitted. """ @@ -943,9 +957,12 @@ class PluginHost: :class:`PluginExecutor` and the caller keeps it in the invocation's own frame. - Mirrors the JS SDK's ``createInvocationPluginRunner`` and the Java SDK's - per-invocation ``PluginRunner``: the handler holds factories, the invocation - holds instances. + The handler holds factories and the invocation holds instances. The other + SDKs are moving to the same split, in aws/aws-durable-execution-sdk-js#924 + (``createInvocationPluginRunner``) and aws/aws-durable-execution-sdk-java#721 + (``PluginRunner`` constructed from factories). Both are open pull requests, so + neither shape is on those repositories' default branches: ``createPluginRunner`` + on JS and ``PluginRunner`` on Java both still hold plugin instances directly. """ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index ac70d520..185becb1 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -100,9 +100,12 @@ def _validate_explicit_factories( otherwise tell which one is wrong. A plugin *class* is callable and stays valid: calling it constructs an - instance, so ``plugins=[MyPlugin]`` is a factory whenever ``MyPlugin`` - accepts the info argument. Only a plugin *instance*, or any other - non-callable value, is rejected. + instance, so ``plugins=[MyPlugin]`` is accepted whenever ``MyPlugin`` accepts + the info argument. It is permitted rather than recommended, because a + constructor should only assign fields and a class used directly as a factory + invites setup work into ``__init__``. ``plugins=[lambda info: MyPlugin(...)]`` + or a ``@classmethod`` factory keeps that work out of the constructor. Only a + plugin *instance*, or any other non-callable value, is rejected. """ factories = list(explicit_plugins or []) for index, factory in enumerate(factories): From 717ef06d47ee2badd14806bf7126bc08d4fc7848 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 10:29:55 -0700 Subject: [PATCH 07/28] refactor(plugin): make the factory an object with create_plugin A reviewer asked for the registration type to be an object with a method rather than a bare callable, and the reason is future extensibility rather than taste. A callable type alias has no member to add anything to. So adding a process-level hook later -- a flush when the execution environment shuts down, for example -- would have to change the registration type from callable to object, which is a second breaking change on the same public surface. An object with one method can gain an optional second member additively, so doing this now costs one break instead of two. `DurableInstrumentationPluginFactory` is therefore a Protocol declaring `create_plugin(info)`, not `Callable[[InvocationStartInfo], ...]`. The type's name and the returned type's name are both unchanged, so no identifier changes meaning. The method is `create_plugin` rather than `new_invocation` because what it returns is still called a plugin, so the verb and the noun agree. This is the shape the Java SDK already had, as `DurableExecutionPluginFactory.createPlugin(InvocationInfo)`, so the three SDKs now describe one contract. The protocol is deliberately not `@runtime_checkable`, for three reasons. An `isinstance` check against a runtime-checkable protocol tests only that the member name exists, so it accepts an object whose `create_plugin` is a string. The SDK also has to name what an invalid entry actually was, which a boolean cannot supply. And such a check requires every declared member, which would make the optional second member above non-additive for any caller who wrote one. `_is_plugin_factory` tests for a callable `create_plugin` instead, structurally, so a factory need not import the protocol. `info` is positional-only on the protocol method. A protocol parameter that is positional-or-keyword is part of the structural contract, so a factory naming it `invocation` would otherwise be a type error. One consequence reverses guidance added earlier in this branch. A plugin class used to be a valid factory, because calling a class constructs an instance, and the docs said so. A class carries no `create_plugin`, so `plugins=[MyPlugin]` now fails at handler initialization, and the test that pinned the old behaviour is replaced by one pinning the new. A class that declares `create_plugin` itself, as a classmethod, is still valid, because the requirement is the member and not the kind of object. The reversal is an improvement for a reason the earlier round raised separately: `CONTRIBUTING.md:246` asks for light constructors, and class-as-factory encouraged `__init__(self, info)` to do setup work. A factory object puts that work in the factory's own constructor, where it can fail before any invocation depends on it. Migrated with it: both bundled plugins' factory classes, whose `__call__` becomes `create_plugin`; 23 conformance plugin handlers, each gaining a small factory class beside its plugin; two examples; one testing-package e2e test that previously aborted collection of that whole suite; and the prose in four READMEs. No handler's observable behaviour changed, and that was measured rather than assumed. The pre-migration handlers were extracted against the pre- migration SDK, every handler was driven through the local runner in both trees, and the emitted records were compared after dropping the execution ARN and wall-clock fields: 23 handlers, 89 records, 23 identical streams, 0 differing. Registration order is preserved in both handlers that register two plugins. --- .../README.md | 4 +- .../src/common.py | 5 +- .../plugin/plugin_attempt_hooks_retry.py | 17 +- .../plugin/plugin_attempt_info_shape.py | 17 +- .../plugin/plugin_context_info_shape.py | 17 +- .../handlers/plugin/plugin_error_isolation.py | 17 +- .../plugin_external_update_on_invoke.py | 17 +- .../plugin/plugin_faulty_and_healthy.py | 43 +++- .../plugin/plugin_first_invocation_flag.py | 17 +- .../plugin/plugin_invocation_info_shape.py | 17 +- .../plugin/plugin_invocation_lifecycle.py | 24 ++- .../plugin/plugin_multiple_plugins.py | 36 +++- .../plugin/plugin_nested_parent_linkage.py | 17 +- .../plugin/plugin_operation_change.py | 17 +- .../plugin/plugin_operation_change_shape.py | 17 +- .../plugin/plugin_operation_info_shape.py | 17 +- .../plugin/plugin_operation_lifecycle.py | 17 +- .../plugin/plugin_parallel_branch_hooks.py | 17 +- .../handlers/plugin/plugin_replay_flags.py | 17 +- .../plugin/plugin_retry_exhaustion.py | 17 +- .../plugin_suspension_invocation_end.py | 17 +- .../plugin/plugin_terminal_failure.py | 17 +- .../plugin/plugin_terminal_payloads.py | 17 +- .../plugin/plugin_wait_operation_hooks.py | 17 +- .../plugin/plugin_wait_replay_flag.py | 17 +- .../src/plugin/execution_with_plugin.py | 18 +- .../src/plugin/execution_with_wait_plugin.py | 17 +- .../README.md | 8 +- .../plugin.py | 24 ++- .../tests/test_plugin.py | 6 +- .../README.md | 18 +- .../plugin_factory.py | 21 +- .../e2e/test_invocation_wait_resume_int.py | 5 +- .../test_execution_plugin_integration.py | 4 +- .../test_invocation_plugin_integration.py | 4 +- .../tests/test_plugin_factory.py | 41 ++-- .../tests/e2e/wait_suspend_replay_test.py | 20 +- .../README.md | 57 ++++-- .../execution.py | 8 +- .../plugin.py | 112 +++++++--- .../plugin_discovery.py | 86 +++++--- .../plugin_invocation_operations_int_test.py | 11 +- .../tests/execution_test.py | 35 ++-- .../tests/plugin_discovery_test.py | 192 ++++++++++++------ .../tests/plugin_test.py | 126 +++++++++--- .../tests/test_helpers.py | 21 +- 46 files changed, 985 insertions(+), 301 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index b36a1b2c..9c23dd9b 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -182,8 +182,8 @@ write access; the runner identity needs list, read, and cleanup access. 1. Find or add the requirement in the conformance repository under `test-requirements//.yaml`. New requirement IDs must be registered there first. -2. Add `src/otel__.py` exporting `handler`. Select the plugin with - `common.otel_plugin_factory()` and guard the input with +2. Add `src/otel__.py` exporting `handler`. Select the plugin factory + with `common.otel_plugin_factory()` and guard the input with `common.require_scenario()`. Use the SDK's real API; never hand-roll behavior to force an expected result. 3. Register the function in `template.yaml` (or `template-long-running.yaml`) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py index bc8cbc37..19ecbfd0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py @@ -23,8 +23,9 @@ def otel_plugin_factory() -> DurableInstrumentationPluginFactory: """Select the telemetry view configured for this deployed function. Returns a factory, which is what ``durable_execution(plugins=[...])`` takes: - the SDK calls it once per invocation to build that invocation's plugin. The - view is still resolved once, when the handler module is imported. + the SDK calls its ``create_plugin`` once per invocation to build that + invocation's plugin. The view is still resolved once, when the handler module + is imported. """ if os.environ.get("OTEL_PLUGIN_MODE") == "execution": diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py index ef684516..4b218725 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py @@ -79,6 +79,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: ) +class AttemptPluginFactory: + """Builds one :class:`AttemptPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> AttemptPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return AttemptPlugin() + + @durable_step def unreliable_operation(step_context: StepContext) -> str: # Fail on the first attempt, succeed on the second, using the SDK's built-in @@ -89,7 +104,7 @@ def unreliable_operation(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[lambda _info: AttemptPlugin()]) +@durable_execution(plugins=[AttemptPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py index a03e98c7..0bb85e96 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py @@ -79,6 +79,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: _emit(record, self._execution_arn) +class AttemptInfoShapePluginFactory: + """Builds one :class:`AttemptInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> AttemptInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return AttemptInfoShapePlugin() + + @durable_step def flaky(step_context: StepContext) -> str: if step_context.attempt < 2: @@ -86,7 +101,7 @@ def flaky(step_context: StepContext) -> str: return "ok" -@durable_execution(plugins=[lambda _info: AttemptInfoShapePlugin()]) +@durable_execution(plugins=[AttemptInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py index 1c9f8873..9e40bb39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py @@ -79,6 +79,21 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: _emit(record, self._execution_arn) +class ContextInfoShapePluginFactory: + """Builds one :class:`ContextInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ContextInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ContextInfoShapePlugin() + + @durable_step def inner(_step_context: StepContext) -> str: return "x" @@ -94,7 +109,7 @@ def branch_b(_context: DurableContext) -> str: return "b-done" -@durable_execution(plugins=[lambda _info: ContextInfoShapePlugin()]) +@durable_execution(plugins=[ContextInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> list[str]: result: BatchResult[str] = context.parallel( [ diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py index ac0a44a8..b378607c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py @@ -82,12 +82,27 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError("faulty attempt-end") +class FaultyPluginFactory: + """Builds one :class:`FaultyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FaultyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FaultyPlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[lambda _info: FaultyPlugin()]) +@durable_execution(plugins=[FaultyPluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py index 580100d5..50da44b4 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py @@ -65,7 +65,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[lambda _info: ExternalUpdatePlugin()]) +class ExternalUpdatePluginFactory: + """Builds one :class:`ExternalUpdatePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ExternalUpdatePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ExternalUpdatePlugin() + + +@durable_execution(plugins=[ExternalUpdatePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py index c45dee84..1c178404 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py @@ -1,9 +1,10 @@ """10-17: Faulty plugin does not affect a healthy plugin. -Two plugins are registered together, in order: a faulty plugin whose every -exercised hook logs a line and then raises, and a healthy plugin that logs -normally. The exercised hooks span the full lifecycle: invocation-start, -operation-start, attempt-start, attempt-end, operation-end, and invocation-end. +Two plugins are registered together, through their factories, in order: a faulty +plugin whose every exercised hook logs a line and then raises, and a healthy +plugin that logs normally. The exercised hooks span the full lifecycle: +invocation-start, operation-start, attempt-start, attempt-end, operation-end, and +invocation-end. In the Python SDK the per-attempt hooks are the real ``on_user_function_start`` / ``on_user_function_end`` callbacks (the latter carries the attempt ``outcome``), and the operation hooks are ``on_operation_start`` / ``on_operation_end``. @@ -103,6 +104,21 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: raise RuntimeError("faulty invocation-end") +class FaultyPluginFactory: + """Builds one :class:`FaultyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FaultyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FaultyPlugin() + + class HealthyPlugin(DurableInstrumentationPlugin): def __init__(self) -> None: # Operation/attempt hooks do not carry the execution ARN, so capture it @@ -182,14 +198,27 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class HealthyPluginFactory: + """Builds one :class:`HealthyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> HealthyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return HealthyPlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution( - plugins=[lambda _info: FaultyPlugin(), lambda _info: HealthyPlugin()] -) +@durable_execution(plugins=[FaultyPluginFactory(), HealthyPluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py index 55b422a1..4b6ea43a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py @@ -47,7 +47,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[lambda _info: FirstInvocationPlugin()]) +class FirstInvocationPluginFactory: + """Builds one :class:`FirstInvocationPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FirstInvocationPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FirstInvocationPlugin() + + +@durable_execution(plugins=[FirstInvocationPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py index f6ec386e..19a20ddb 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py @@ -58,7 +58,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: _emit(record, info.execution_arn) -@durable_execution(plugins=[lambda _info: InvocationInfoShapePlugin()]) +class InvocationInfoShapePluginFactory: + """Builds one :class:`InvocationInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> InvocationInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return InvocationInfoShapePlugin() + + +@durable_execution(plugins=[InvocationInfoShapePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return f"done-{event}" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py index 9e40133f..4deed16d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py @@ -1,9 +1,10 @@ """10-1: Plugin invocation lifecycle hooks (start and end on a single invocation). Registers an instrumentation plugin through the SDK's real ``plugins=[...]`` -parameter on ``durable_execution``. The plugin emits its lines from the SDK's -``on_invocation_start`` / ``on_invocation_end`` hooks; the step body logs its -running line via the SDK-provided step context logger (mirrors handler 1-7). +parameter on ``durable_execution``, which takes the plugin's factory. The plugin +emits its lines from the SDK's ``on_invocation_start`` / ``on_invocation_end`` +hooks; the step body logs its running line via the SDK-provided step context +logger (mirrors handler 1-7). """ import json @@ -50,13 +51,28 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class LifecyclePluginFactory: + """Builds one :class:`LifecyclePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> LifecyclePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return LifecyclePlugin() + + @durable_step def greet(step_context: StepContext, name: str) -> str: step_context.logger.info(f"Greeting step running for: {name}") return f"Hello, {name}!" -@durable_execution(plugins=[lambda _info: LifecyclePlugin()]) +@durable_execution(plugins=[LifecyclePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py index c5f80166..2600ed32 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py @@ -1,8 +1,8 @@ """10-5: Multiple registered plugins all receive lifecycle hooks. Two instrumentation plugins are registered together, in order A then B, through -the SDK's real ``plugins=[...]`` parameter. Each emits its own prefixed lines -from the invocation-start / invocation-end hooks. +the SDK's real ``plugins=[...]`` parameter, which takes their factories. Each +emits its own prefixed lines from the invocation-start / invocation-end hooks. """ import json @@ -45,6 +45,21 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class PluginAFactory: + """Builds one :class:`PluginA` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> PluginA: + """Return this invocation's plugin. ``info`` is unused.""" + return PluginA() + + class PluginB(DurableInstrumentationPlugin): def on_invocation_start(self, info: InvocationStartInfo) -> None: _emit( @@ -60,12 +75,27 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class PluginBFactory: + """Builds one :class:`PluginB` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> PluginB: + """Return this invocation's plugin. ``info`` is unused.""" + return PluginB() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[lambda _info: PluginA(), lambda _info: PluginB()]) +@durable_execution(plugins=[PluginAFactory(), PluginBFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py index ec49bd39..ccb58191 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py @@ -56,6 +56,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class ParentLinkagePluginFactory: + """Builds one :class:`ParentLinkagePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ParentLinkagePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ParentLinkagePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" @@ -66,7 +81,7 @@ def child_operation(ctx: DurableContext, name: str) -> str: return ctx.step(greet(name)) -@durable_execution(plugins=[lambda _info: ParentLinkagePlugin()]) +@durable_execution(plugins=[ParentLinkagePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.run_in_child_context(child_operation(str(event))) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py index 366d3ebe..f3220bf7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py @@ -58,12 +58,27 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: ) +class OperationChangePluginFactory: + """Builds one :class:`OperationChangePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationChangePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationChangePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[lambda _info: OperationChangePlugin()]) +@durable_execution(plugins=[OperationChangePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py index e4f86021..c1982a33 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py @@ -78,12 +78,27 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: _emit(record, self._execution_arn) +class OperationChangeShapePluginFactory: + """Builds one :class:`OperationChangeShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationChangeShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationChangeShapePlugin() + + @durable_step def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[lambda _info: OperationChangeShapePlugin()]) +@durable_execution(plugins=[OperationChangeShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py index a0294291..11de1570 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py @@ -74,12 +74,27 @@ def on_operation_end(self, info: OperationEndInfo) -> None: _emit(_operation_record("operation-end", info), self._execution_arn) +class OperationInfoShapePluginFactory: + """Builds one :class:`OperationInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationInfoShapePlugin() + + @durable_step def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[lambda _info: OperationInfoShapePlugin()]) +@durable_execution(plugins=[OperationInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py index 734177c8..10b1a943 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py @@ -72,12 +72,27 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class OperationLifecyclePluginFactory: + """Builds one :class:`OperationLifecyclePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationLifecyclePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationLifecyclePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[lambda _info: OperationLifecyclePlugin()]) +@durable_execution(plugins=[OperationLifecyclePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py index 7c642d47..f8ff4c01 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py @@ -75,6 +75,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: ) +class ParallelBranchPluginFactory: + """Builds one :class:`ParallelBranchPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ParallelBranchPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ParallelBranchPlugin() + + def branch0(_ctx: DurableContext) -> str: return "task-1" @@ -83,7 +98,7 @@ def branch1(_ctx: DurableContext) -> str: return "task-2" -@durable_execution(plugins=[lambda _info: ParallelBranchPlugin()]) +@durable_execution(plugins=[ParallelBranchPluginFactory()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [branch0, branch1], diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py index e0490f69..e7e17372 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py @@ -75,6 +75,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class ReplayFlagPluginFactory: + """Builds one :class:`ReplayFlagPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ReplayFlagPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ReplayFlagPlugin() + + @durable_step def step_a(_step_context: StepContext) -> str: return "a" @@ -90,7 +105,7 @@ def step_b(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[lambda _info: ReplayFlagPlugin()]) +@durable_execution(plugins=[ReplayFlagPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) retry_config = RetryStrategyConfig( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py index f851fce1..877b16af 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py @@ -93,13 +93,28 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class RetryExhaustionPluginFactory: + """Builds one :class:`RetryExhaustionPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RetryExhaustionPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RetryExhaustionPlugin() + + @durable_step def always_fail(_step_context: StepContext) -> str: msg = "boom" raise RuntimeError(msg) -@durable_execution(plugins=[lambda _info: RetryExhaustionPlugin()]) +@durable_execution(plugins=[RetryExhaustionPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=2, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py index 496dca5b..50f80739 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py @@ -57,7 +57,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[lambda _info: SuspensionPlugin()]) +class SuspensionPluginFactory: + """Builds one :class:`SuspensionPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> SuspensionPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return SuspensionPlugin() + + +@durable_execution(plugins=[SuspensionPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py index 237fe227..3f3cec44 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py @@ -52,13 +52,28 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class TerminalFailurePluginFactory: + """Builds one :class:`TerminalFailurePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> TerminalFailurePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return TerminalFailurePlugin() + + @durable_step def failing_step(_step_context: StepContext) -> str: msg = "Something went wrong" raise RuntimeError(msg) -@durable_execution(plugins=[lambda _info: TerminalFailurePlugin()]) +@durable_execution(plugins=[TerminalFailurePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step( failing_step(), diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py index 849e533e..ff3af95d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py @@ -62,6 +62,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class TerminalPayloadPluginFactory: + """Builds one :class:`TerminalPayloadPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> TerminalPayloadPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return TerminalPayloadPlugin() + + @durable_step def step_a(_step_context: StepContext) -> str: return "task-a" @@ -73,7 +88,7 @@ def step_b(_step_context: StepContext) -> str: raise RuntimeError(msg) -@durable_execution(plugins=[lambda _info: TerminalPayloadPlugin()]) +@durable_execution(plugins=[TerminalPayloadPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) result: str = context.step( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py index d18e7958..c8e8e1b3 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py @@ -68,7 +68,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[lambda _info: WaitOperationPlugin()]) +class WaitOperationPluginFactory: + """Builds one :class:`WaitOperationPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> WaitOperationPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return WaitOperationPlugin() + + +@durable_execution(plugins=[WaitOperationPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py index 882cf92e..59deac3a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py @@ -81,6 +81,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class WaitReplayFlagPluginFactory: + """Builds one :class:`WaitReplayFlagPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> WaitReplayFlagPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return WaitReplayFlagPlugin() + + def wait_short(ctx: DurableContext) -> str: ctx.wait(Duration.from_seconds(2), name="short") return "short-done" @@ -91,7 +106,7 @@ def wait_long(ctx: DurableContext) -> str: return "long-done" -@durable_execution(plugins=[lambda _info: WaitReplayFlagPlugin()]) +@durable_execution(plugins=[WaitReplayFlagPluginFactory()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [wait_short, wait_long], diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py index 56e908b4..903cef96 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py @@ -12,6 +12,7 @@ from aws_durable_execution_sdk_python.execution import durable_execution from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, + InvocationStartInfo, ) @@ -37,6 +38,21 @@ def on_user_function_end(self, info) -> None: self.logger.info(f"User function ended: {info}") +class MyPluginFactory: + """Builds one :class:`MyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> MyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return MyPlugin() + + @durable_step def add_numbers(_step_context: StepContext, a: int, b: int) -> int: return a + b @@ -51,7 +67,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[lambda _info: MyPlugin()]) +@durable_execution(plugins=[MyPluginFactory()]) def handler(_event: Any, context: DurableContext) -> int: result: int = context.run_in_child_context( add_numbers_in_child(6, 4), diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py index 72491e90..9851977c 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py @@ -42,7 +42,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[lambda _info: RecordingWaitPlugin()]) +class RecordingWaitPluginFactory: + """Builds one :class:`RecordingWaitPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RecordingWaitPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RecordingWaitPlugin() + + +@durable_execution(plugins=[RecordingWaitPluginFactory()]) def handler(_event: Any, context: DurableContext) -> dict[str, Any]: context.wait(Duration.from_seconds(1), name="plugin-wait") return { diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 1d8efb26..5aafb9ac 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -43,10 +43,10 @@ def handler(event, context): ``` `workflow_insight()` returns a plugin *factory*, which is what the SDK's -`plugins` argument takes: the SDK calls it once per invocation to build that -invocation's plugin instance. The factory holds the resolved configuration and -the exporters, so configuration is per handler while record state is per -invocation. +`plugins` argument takes: the SDK calls its `create_plugin` once per invocation to +build that invocation's plugin instance. The factory holds the resolved +configuration and the exporters, so configuration is per handler while record +state is per invocation. With no exporter configured, records are written to the function's own CloudWatch log group as single JSON lines (the `LambdaLogExporter` default), diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 35f8e024..f80bca97 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -14,8 +14,8 @@ ``@durable_execution(plugins=[...])`` takes. The factory lives as long as the handler and owns everything that is not per-execution: the resolved immutable config, the exporters, and the ``_ExportScheduler`` that serializes export - across executions. The SDK calls it once per invocation and drops the instance - it returns when the invocation scope exits, so a + across executions. The SDK calls its ``create_plugin`` once per invocation and + drops the instance it returns when the invocation scope exits, so a :class:`WorkflowInsightPlugin` instance serves exactly one invocation of one execution. Everything this environment holds for that execution is therefore ordinary instance state: no ARN-keyed registry, and no hook can reach an @@ -500,19 +500,20 @@ def _emit( class _WorkflowInsightFactory: """The handler-lifetime half of the plugin: what is NOT per-execution. - Satisfies the SDK's ``DurableInstrumentationPluginFactory`` -- it is called - with an ``InvocationStartInfo`` and returns the plugin instance for that - invocation. Everything it holds is either immutable after construction (the - resolved config) or deliberately shared across executions: + Satisfies the SDK's ``DurableInstrumentationPluginFactory`` -- its + ``create_plugin`` is called with an ``InvocationStartInfo`` and returns the + plugin instance for that invocation. Everything it holds is either immutable + after construction (the resolved config) or deliberately shared across + executions: * the exporters, which are customer objects registered once, and * the ``_ExportScheduler``, because export serialization is cross-execution: one worker, one ``export()`` at a time, whatever the instance that scheduled the record. - A callable class rather than a closure so the resolved config stays - inspectable (``factory._emit_mode``, ``factory._exporters``) instead of being - buried in cell variables. + A class rather than a closure so the resolved config stays inspectable + (``factory._emit_mode``, ``factory._exporters``) instead of being buried in + cell variables. """ def __init__(self, config: WorkflowInsightConfig) -> None: @@ -549,7 +550,7 @@ def __init__(self, config: WorkflowInsightConfig) -> None: ) self._scheduler = _ExportScheduler(self._exporters) - def __call__(self, info: InvocationStartInfo) -> WorkflowInsightPlugin: + def create_plugin(self, info: InvocationStartInfo) -> WorkflowInsightPlugin: return WorkflowInsightPlugin(self, info) @@ -557,6 +558,7 @@ def workflow_insight(config: WorkflowInsightConfig) -> _WorkflowInsightFactory: """Create a Workflow Insight plugin factory. Mirrors the JS ``workflowInsight()``. Pass the result straight to ``@durable_execution(plugins=[...])``: the SDK - calls it once per invocation to build that invocation's plugin instance. + calls its ``create_plugin`` once per invocation to build that invocation's + plugin instance. """ return _WorkflowInsightFactory(config) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index baed5b7e..8cbf23ea 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -156,7 +156,7 @@ def _invocation(factory, info: InvocationStartInfo) -> WorkflowInsightPlugin: start info and dispatches the very same object to its first hook. A test that drives hooks directly does both. """ - plugin = factory(info) + plugin = factory.create_plugin(info) plugin.on_invocation_start(info) return plugin @@ -827,7 +827,7 @@ def __del__(self) -> None: ) op = _step("s", op_id="1") start = _start(operations={}) - plugin = factory(start) + plugin = factory.create_plugin(start) holder["plugin"] = plugin holder["ops"] = _ops(op) change = OperationChangeInfo( @@ -1093,7 +1093,7 @@ def reentering_input(value: Any) -> Any: ) ) start = _start(operations={}) - plugin = factory(start) + plugin = factory.create_plugin(start) holder["plugin"] = plugin # On a bounded thread, so a regression that makes the lock non-reentrant diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 81304132..6aa8d1ed 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -43,10 +43,10 @@ processors, and exporter. 4. Add X-Ray write permissions The SDK's `plugins` list takes plugin *factories*, not plugin instances: it calls -each factory once per invocation and the plugin it returns serves that one -invocation. `InvocationOtelPluginFactory` and `ExecutionOtelPluginFactory` are -the factories for the two bundled plugins; each takes the optional -`OtelPluginConfig` that every plugin it builds will use. +each factory's `create_plugin` once per invocation and the plugin it returns +serves that one invocation. `InvocationOtelPluginFactory` and +`ExecutionOtelPluginFactory` are the factories for the two bundled plugins; each +takes the optional `OtelPluginConfig` that every plugin it builds will use. Alternatively, install this package in the function artifact or a Lambda layer and select either OTel plugin by entry-point name: @@ -58,9 +58,9 @@ DURABLE_EXECUTION_PLUGINS=otel-execution `otel-invocation` names a default-configured `InvocationOtelPluginFactory`; `otel-execution` names a default-configured `ExecutionOtelPluginFactory`. The SDK -discovers the selected package entry point at cold start and calls it once per -invocation, so the handler does not need to import or explicitly register the -plugin. +discovers the selected package entry point at cold start and calls its +`create_plugin` once per invocation, so the handler does not need to import or +explicitly register the plugin. ### 1. ADOT Lambda Layer @@ -381,8 +381,8 @@ After deploying your function with the plugin configured: Factory for the invocation-rooted plugin, and what belongs in the SDK's `plugins` list. Satisfies `DurableInstrumentationPluginFactory` from -`aws_durable_execution_sdk_python`: calling it with an `InvocationStartInfo` -returns the `InvocationOtelPlugin` for that invocation. +`aws_durable_execution_sdk_python`: its `create_plugin(info)` takes an +`InvocationStartInfo` and returns the `InvocationOtelPlugin` for that invocation. ```python InvocationOtelPluginFactory( diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py index ddda8dcd..e5bbfc51 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py @@ -1,14 +1,15 @@ """Plugin factories for the bundled durable-execution OTel plugins. -The SDK's plugin contract is a factory called once per invocation: -``DurableInstrumentationPluginFactory = Callable[[InvocationStartInfo], -DurableInstrumentationPlugin]``. The instance a factory returns serves exactly -that one invocation and is dropped when the invocation scope exits, so a plugin -keeps its per-invocation state in ordinary instance attributes. +The SDK's plugin contract is a factory object whose ``create_plugin`` is called +once per invocation: ``DurableInstrumentationPluginFactory`` declares +``create_plugin(info: InvocationStartInfo) -> DurableInstrumentationPlugin``. The +instance a factory returns serves exactly that one invocation and is dropped when +the invocation scope exits, so a plugin keeps its per-invocation state in ordinary +instance attributes. -Both factories are callable classes rather than closures so the configuration -they were built with stays inspectable (``factory.config``) and so the entry -points below name an object with a readable type. +Both factories are classes rather than closures so the configuration they were +built with stays inspectable (``factory.config``) and so the entry points below +name an object with a readable type. Everything else the plugins need is resolved per invocation inside the plugin itself: the tracer provider (which for the global-provider case may only be @@ -57,7 +58,7 @@ def handler(event, context): ... def __init__(self, config: OtelPluginConfig | None = None) -> None: self.config = config - def __call__(self, info: InvocationStartInfo) -> InvocationOtelPlugin: + def create_plugin(self, info: InvocationStartInfo) -> InvocationOtelPlugin: """Return this invocation's plugin. ``info`` is accepted because the SDK passes it, and is unused: the @@ -78,7 +79,7 @@ class ExecutionOtelPluginFactory: def __init__(self, config: OtelPluginConfig | None = None) -> None: self.config = config - def __call__(self, info: InvocationStartInfo) -> ExecutionOtelPlugin: + def create_plugin(self, info: InvocationStartInfo) -> ExecutionOtelPlugin: """Return this invocation's plugin. ``info`` is unused; see above.""" return ExecutionOtelPlugin(self.config) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index fafdfa7f..ebabc49b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -151,8 +151,9 @@ def test_otel_wait_resume_spans_share_default_xray_execution_trace( exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) - # The SDK takes a factory and calls it once per invocation, so the two - # invocations below are served by two plugin instances sharing this provider. + # The SDK takes a factory and calls its create_plugin once per invocation, so + # the two invocations below are served by two plugin instances sharing this + # provider. factory = factory_type( OtelPluginConfig( tracer_provider=provider, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index 6e8fd15c..d559625f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -316,7 +316,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() - plugin = factory(_invocation_start()) + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -325,7 +325,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () - plugin = factory(_invocation_start()) + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py index 12b0b6cb..7f71ee72 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -266,7 +266,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() - plugin = factory(_invocation_start()) + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -275,7 +275,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () - plugin = factory(_invocation_start()) + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py index 5999fc60..32931438 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py @@ -1,11 +1,11 @@ """Tests for the bundled OTel plugin factories and the entry points naming them. -The SDK plugin contract is a factory called once per invocation, so what these -tests have to establish is that the objects the package exposes -- and the ones -its entry points name -- are callables that build a plugin, and that each call -builds a NEW plugin. The old provider-shaped assertions (a declared -``plugin_type`` and an API version) have no counterpart: the contract carries -neither. +The SDK plugin contract is a factory object whose ``create_plugin`` is called once +per invocation, so what these tests have to establish is that the objects the +package exposes -- and the ones its entry points name -- carry ``create_plugin``, +that it builds a plugin, and that each call builds a NEW plugin. The old +provider-shaped assertions (a declared ``plugin_type`` and an API version) have no +counterpart: the contract carries neither. """ from __future__ import annotations @@ -96,7 +96,7 @@ def test_factory_builds_its_plugin_type( ) -> None: factory = factory_type(OtelPluginConfig(enrich_logger=False)) - assert isinstance(factory(_invocation_start_info()), plugin_type) + assert isinstance(factory.create_plugin(_invocation_start_info()), plugin_type) @pytest.mark.parametrize( @@ -113,8 +113,8 @@ def test_factory_builds_a_fresh_plugin_per_invocation(factory_type: type) -> Non """ factory = factory_type(OtelPluginConfig(enrich_logger=False)) - first = factory(_invocation_start_info()) - second = factory(_invocation_start_info()) + first = factory.create_plugin(_invocation_start_info()) + second = factory.create_plugin(_invocation_start_info()) assert first is not second @@ -128,8 +128,8 @@ def test_factory_passes_its_config_to_every_plugin(factory_type: type) -> None: factory = factory_type(config) assert factory.config is config - assert factory(_invocation_start_info())._config is config - assert factory(_invocation_start_info())._config is config + assert factory.create_plugin(_invocation_start_info())._config is config + assert factory.create_plugin(_invocation_start_info())._config is config @pytest.mark.parametrize( @@ -142,7 +142,7 @@ def test_factory_without_config_builds_a_default_configured_plugin( factory = factory_type() assert factory.config is None - assert factory(_invocation_start_info())._config == OtelPluginConfig() + assert factory.create_plugin(_invocation_start_info())._config == OtelPluginConfig() def test_module_level_factories_are_default_configured() -> None: @@ -158,12 +158,14 @@ def test_declared_entry_points_name_the_bundled_factories() -> None: assert _resolve(entry_points["otel-execution"]) is EXECUTION_OTEL_PLUGIN_FACTORY -def test_declared_entry_points_resolve_to_callables_that_build_plugins() -> None: - """The entry points must satisfy the SDK's factory contract, not a provider. +def test_declared_entry_points_resolve_to_factories_that_build_plugins() -> None: + """The entry points must satisfy the SDK's factory contract. - ``plugin_discovery._load_factory`` accepts anything callable, so a target - that resolved to a plugin class -- or to a plugin instance -- would load - without complaint and only fail at invocation time. + ``plugin_discovery._load_factory`` requires an object with a callable + ``create_plugin``, so a target resolving to a plugin class, to a plugin + instance, or to a plain function now fails at handler initialization. The + check here mirrors the SDK's own, then calls the method to confirm what it + builds. """ expected = { "otel-invocation": InvocationOtelPlugin, @@ -172,6 +174,7 @@ def test_declared_entry_points_resolve_to_callables_that_build_plugins() -> None for name, spec in _declared_entry_points().items(): factory = _resolve(spec) - assert callable(factory) + create_plugin = getattr(factory, "create_plugin", None) + assert callable(create_plugin) assert not isinstance(factory, type) - assert isinstance(factory(_invocation_start_info()), expected[name]) + assert isinstance(create_plugin(_invocation_start_info()), expected[name]) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py index be00a7f7..b7a4f137 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py @@ -60,15 +60,29 @@ def on_operation_end(self, info: OperationEndInfo) -> None: self.wait_end_infos.append(info) +class RecordingWaitPluginFactory: + """Builds one :class:`RecordingWaitPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, and this handler is built at + module import, so registering one would abort collection of this module. This + factory exists only to construct the plugin; the state the test asserts on is + class-level on the plugin, so the factory holds nothing. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RecordingWaitPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RecordingWaitPlugin() + + def _wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 """Suspend on a top-level wait, then finish.""" context.wait(Duration.from_seconds(1), name=_WAIT_NAME) return "done" -wait_handler = durable_execution( - _wait_handler, plugins=[lambda _info: RecordingWaitPlugin()] -) +wait_handler = durable_execution(_wait_handler, plugins=[RecordingWaitPluginFactory()]) def test_wait_completed_during_suspend_is_delivered_as_new() -> None: diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index d542c8f6..214f1422 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -43,23 +43,40 @@ variable preserves the existing behavior. The decorator's `plugins` argument remains supported; explicit factories run first, and a factory passed to the decorator is not registered a second time through the environment. -A plugin is registered as a *factory*, not as an instance. A factory is any -callable taking the invocation's `InvocationStartInfo` and returning a -`DurableInstrumentationPlugin`; the SDK calls it once per invocation, so the -instance it returns serves that one invocation only and can hold per-execution -state in ordinary attributes. - -Write the factory as a function, a `lambda`, or a `@classmethod`, and construct -the plugin inside it. A constructor should only assign fields, so a plugin class -used directly as a factory invites setup work into `__init__`. A plugin class is -callable and is accepted as a factory whenever its `__init__` takes the info, but -prefer the explicit form: +A plugin is registered as a *factory*, not as an instance. A factory is an object +with a `create_plugin(info)` method taking the invocation's `InvocationStartInfo` +and returning a `DurableInstrumentationPlugin`; the SDK calls that method once per +invocation, so the instance it returns serves that one invocation only and can +hold per-execution state in ordinary attributes. + +A factory is an object with a method rather than a plain callable so the +registration type can grow a second, optional member later -- a process-level +flush on execution-environment shutdown, for example -- without a second breaking +change to this surface. + +Write a small factory class and construct the plugin in its `create_plugin`. The +factory holds what outlives an invocation, such as an exporter or a resolved +configuration, and setup work that can fail belongs in the factory's own +constructor rather than the plugin's: ```python -plugins=[lambda info: AuditPlugin(sink)] # construct explicitly -plugins=[AuditPlugin.create] # a @classmethod factory +class AuditPluginFactory: + def __init__(self, sink): + self._sink = sink + + def create_plugin(self, info): + return AuditPlugin(self._sink) + + +plugins=[AuditPluginFactory(sink)] ``` +A plugin class is not a factory, and neither is a bare callable. `plugins=[MyPlugin]` +and `plugins=[lambda info: MyPlugin(sink)]` raise `PluginLoadError` during handler +initialization, because neither carries `create_plugin`. A class that declares +`create_plugin` as a `@classmethod` is accepted, since the requirement is the +member and not the kind of object. + Provider packages expose such a factory: ```python @@ -73,19 +90,23 @@ class AuditPlugin(DurableInstrumentationPlugin): pass -def audit_plugin_factory(info: InvocationStartInfo) -> AuditPlugin: - return AuditPlugin() +class AuditPluginFactory: + def create_plugin(self, info: InvocationStartInfo) -> AuditPlugin: + return AuditPlugin() + + +AUDIT_PLUGIN_FACTORY = AuditPluginFactory() ``` -Register the factory in the package's `pyproject.toml`: +Register the factory instance in the package's `pyproject.toml`: ```toml [project.entry-points."aws_durable_execution.plugins"] -example_audit = "example_audit:audit_plugin_factory" +example_audit = "example_audit:AUDIT_PLUGIN_FACTORY" ``` Provider names must be unique across installed distributions. Missing, -ambiguous, or non-callable providers raise `PluginLoadError` during handler +ambiguous, or wrongly shaped providers raise `PluginLoadError` during handler initialization with the provider and distribution details. ## 🚀 Quick Start diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 4fb80fbc..517eca30 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -179,10 +179,10 @@ def durable_execution( Args: func: The user function to decorate boto3_client: Optional boto3 Lambda client to use - plugins: Optional list of instrumentation plugin factories. Each factory - is called once per invocation with that invocation's - ``InvocationStartInfo``, and the instance it returns serves only that - invocation. + plugins: Optional list of instrumentation plugin factories. Each + factory's ``create_plugin`` is called once per invocation with that + invocation's ``InvocationStartInfo``, and the instance it returns + serves only that invocation. """ # Decorator called with parameters if func is None: 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 e24bb13a..e7acdd1f 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 @@ -9,7 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, MutableMapping, Protocol, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -449,40 +449,79 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: pass -DurableInstrumentationPluginFactory = Callable[ - [InvocationStartInfo], DurableInstrumentationPlugin -] -"""Builds one plugin instance for one invocation. +class DurableInstrumentationPluginFactory(Protocol): + """Builds one plugin instance for one invocation. + + An object with a method rather than a bare callable, because the SDK will + grow process-level plugin hooks -- a flush when the execution environment + shuts down, for example. A callable type has no member to add such a hook to, + so growing one would have to change the registration type from callable to + object, which is a second breaking change on the same public surface. One + method on an object leaves room for an optional second member, which is + additive. + + Not ``@runtime_checkable``. Three facts decide it. An ``isinstance`` check + against a runtime-checkable protocol tests only that the member name is + present, not that it is callable, so it would accept an object whose + ``create_plugin`` is a string; the SDK needs the stronger test. The SDK also + has to name what an invalid entry actually was, which a boolean + ``isinstance`` result cannot supply. And ``isinstance`` against a + runtime-checkable protocol requires *every* declared member, so publishing + one would make the optional second member above non-additive for any caller + who wrote such a check. :func:`plugin_discovery._is_plugin_factory` performs + the check instead. + + Register the factory, not a plugin:: + + class MyPluginFactory: + def __init__(self, exporter: Exporter) -> None: + self._exporter = exporter + + def create_plugin(self, info: InvocationStartInfo) -> MyPlugin: + return MyPlugin(self._exporter) + + plugins=[MyPluginFactory(exporter)] + + A plugin class is not a factory. ``plugins=[MyPlugin]`` used to work because + calling a class constructs an instance, and it now fails at handler + initialization because a class carries no ``create_plugin``. A class that + declares ``create_plugin`` itself -- as a ``@classmethod`` -- does satisfy the + shape, because the requirement is the member and not the kind of object. + + The factory holds what outlives an invocation: an exporter, a resolved + configuration, a shared worker. The plugin instance holds what does not. + Setup work that can fail or that reads the environment belongs in the + factory's own constructor rather than in the plugin's, because a plugin + constructor should only assign fields (see ``CONTRIBUTING.md``, + "Initialization and conversion"). + """ -Called once per invocation with that invocation's :class:`InvocationStartInfo` -- -the same object the instance's ``on_invocation_start`` then receives -- before any -hook fires. The instance serves only that invocation and is dropped when it -returns, so a plugin can hold per-execution state in ordinary instance -attributes without keying it by execution ARN. + def create_plugin( + self, info: InvocationStartInfo, / + ) -> DurableInstrumentationPlugin: + """Return the plugin instance that serves the described invocation. -A plain ``Callable`` alias rather than a ``Protocol``: the shape has exactly one -call signature and no other members, so a Protocol would only add a name. The -alias is also the more permissive of the two, because ``Callable`` parameters are -positional-only -- a factory may name its parameter whatever reads best -(``lambda info: ...``, ``def build(invocation): ...``), where a ``__call__`` -Protocol would pin that name. + Called once per invocation, with that invocation's + :class:`InvocationStartInfo` -- the same object the returned instance's + ``on_invocation_start`` then receives -- and before any hook fires. The + instance serves only that invocation and is dropped when it returns, so a + plugin can hold per-execution state in ordinary instance attributes + without keying it by execution ARN. -Prefer a factory that constructs the plugin explicitly:: + ``info`` is positional-only, so an implementation may name the parameter + whatever reads best; a named protocol parameter would pin that name for + every implementation. - plugins=[lambda info: MyPlugin(exporter)] - plugins=[MyPlugin.create] # a @classmethod factory + A call that raises, or that returns ``None``, is logged and skipped for + that invocation and never disrupts the execution. -A constructor should only assign fields, so setup work that can fail or that -reads the environment belongs in a factory rather than in ``__init__`` (see -``CONTRIBUTING.md``, "Initialization and conversion"). An explicit factory is -where that work goes, and it also lets the plugin take its own collaborators -rather than deriving them from the hook info. + Args: + info: The invocation the returned plugin instance will observe. -Anything callable satisfies the alias: a lambda, a module-level function, a -``functools.partial``, a ``@classmethod``, or a plugin class itself, since calling -a class in Python constructs an instance. ``plugins=[MyPlugin]`` is therefore -permitted whenever ``MyPlugin.__init__`` takes the info, and it stays permitted. -""" + Returns: + The plugin instance for this invocation. + """ + ... def _factory_name(factory: object) -> str: @@ -561,14 +600,21 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: """Build this invocation's plugin instances from its start info. Called once per invocation, before the first hook is dispatched. A - factory that raises or returns ``None`` is contained exactly as a failing - hook is -- logged and skipped -- so a broken plugin cannot disrupt the - execution. The remaining factories still produce their instances. + factory whose ``create_plugin`` raises or returns ``None`` is contained + exactly as a failing hook is -- logged and skipped -- so a broken plugin + cannot disrupt the execution. The remaining factories still produce their + instances. + + An entry without a usable ``create_plugin`` raises ``AttributeError`` + here, which this containment then swallows once per invocation. + :func:`plugin_discovery.load_configured_plugins` rejects such an entry + while the handler is being initialized, so the silent case is not + reachable through ``durable_execution()``. """ plugins: list[DurableInstrumentationPlugin] = [] for factory in self._plugin_factories: try: - plugin = factory(info) + plugin = factory.create_plugin(info) except Exception: # log and ignore the exception logger.exception( diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index 185becb1..c0c83209 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -54,17 +54,38 @@ def _qualified_type_name(value: object) -> str: return f"{value_type.__module__}.{value_type.__qualname__}" +def _is_plugin_factory(value: object) -> bool: + """Report whether a value has the shape of a plugin factory. + + :class:`DurableInstrumentationPluginFactory` declares one method, so the + shape is one member: a callable ``create_plugin``. The attribute is fetched + and tested for callability rather than merely for presence, because an + object carrying a non-callable ``create_plugin`` would otherwise pass here + and fail at invocation time. + + Structural rather than nominal, so a factory need not import the SDK + protocol to satisfy it. The protocol is deliberately not + ``@runtime_checkable``; see its docstring. + + The check cannot go further than one member. Whether ``create_plugin`` + accepts the info, and whether it returns a plugin, is only knowable by + calling it, and calling it at load time is what the per-invocation factory + design avoids: there is no invocation yet. + """ + return callable(getattr(value, "create_plugin", None)) + + def _load_factory( plugin_name: str, entry_point: metadata.EntryPoint ) -> DurableInstrumentationPluginFactory: """Resolve an entry point to a plugin factory. Only two things can still be checked here. The entry point has to import, - and what it resolves to has to be callable. Nothing more is knowable without - calling the factory, and calling it at load time is precisely what this - design avoids: the instance belongs to an invocation, and there is no - invocation yet. A factory that then misbehaves at invocation time is - contained by :meth:`PluginExecutor._create_plugins`. + and what it resolves to has to have the factory shape. Nothing more is + knowable without calling the factory, and calling it at load time is + precisely what this design avoids: the instance belongs to an invocation, + and there is no invocation yet. A factory that then misbehaves at invocation + time is contained by :meth:`PluginExecutor._create_plugins`. """ try: factory = entry_point.load() @@ -75,11 +96,14 @@ def _load_factory( f"({_distribution_name(entry_point)}): {error}" ) from error - if not callable(factory): + if not _is_plugin_factory(factory): raise PluginLoadError( f"Durable instrumentation plugin entry point '{plugin_name}' must " - "resolve to a callable plugin factory, but resolved to " - f"{_qualified_type_name(factory)}." + "resolve to a plugin factory -- an object with a " + "create_plugin(info) method returning a " + "DurableInstrumentationPlugin -- but resolved to " + f"{_qualified_type_name(factory)}. Name the factory instance, not a " + "plugin and not a plugin class." ) return cast(DurableInstrumentationPluginFactory, factory) @@ -88,34 +112,36 @@ def _load_factory( def _validate_explicit_factories( explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, ) -> list[DurableInstrumentationPluginFactory]: - """Check that every explicitly passed plugin entry is callable. + """Check that every explicitly passed plugin entry has the factory shape. - Each entry is called once per invocation to build that invocation's plugin - instance. An entry that is not callable can never be called, so - :meth:`PluginExecutor._create_plugins` raises ``TypeError`` on every + Each entry's ``create_plugin`` is called once per invocation to build that + invocation's plugin instance. An entry without one can never be called, so + :meth:`PluginExecutor._create_plugins` raises ``AttributeError`` on every invocation, logs it and continues without that plugin -- telemetry is lost for the lifetime of the function, and nothing fails. Raising here converts that into one configuration failure while the handler is being initialized. The position is named because a caller passing several entries cannot otherwise tell which one is wrong. - A plugin *class* is callable and stays valid: calling it constructs an - instance, so ``plugins=[MyPlugin]`` is accepted whenever ``MyPlugin`` accepts - the info argument. It is permitted rather than recommended, because a - constructor should only assign fields and a class used directly as a factory - invites setup work into ``__init__``. ``plugins=[lambda info: MyPlugin(...)]`` - or a ``@classmethod`` factory keeps that work out of the constructor. Only a - plugin *instance*, or any other non-callable value, is rejected. + A plugin *class* is rejected, and so is any bare callable. Both were + accepted while the registration type was ``Callable``: a lambda satisfied it + directly, and a class satisfied it because calling a class constructs an + instance. Neither carries ``create_plugin``, so ``plugins=[MyPlugin]`` and + ``plugins=[lambda info: MyPlugin()]`` now fail here. The replacement is a + small factory class, which is also where setup work that can fail belongs. A + class that declares ``create_plugin`` as a ``@classmethod`` is accepted, + because the requirement is the member and not the kind of object. """ factories = list(explicit_plugins or []) for index, factory in enumerate(factories): - if not callable(factory): + if not _is_plugin_factory(factory): raise PluginLoadError( f"Durable instrumentation plugin at plugins[{index}] must be a " - "callable plugin factory taking an InvocationStartInfo, but is " + "plugin factory -- an object with a create_plugin(info) method " + "returning a DurableInstrumentationPlugin -- but is " f"{_qualified_type_name(factory)}. Pass a factory rather than a " - "plugin instance, for example " - "plugins=[lambda info: MyPlugin(...)]." + "plugin, a plugin class, or a plain callable, for example " + "plugins=[MyPluginFactory(exporter)]." ) return factories @@ -128,16 +154,18 @@ def load_configured_plugins( """Combine explicit plugin factories with those selected through the environment. Explicit factories retain their order. Dynamically selected factories follow - in configured order. Every returned factory is called once per invocation. + in configured order. Every returned factory has its ``create_plugin`` called + once per invocation. A factory already registered explicitly is not registered a second time through the environment. The check is by factory identity, which is what is knowable here: the old shape declared a ``plugin_type`` and could dedup on - it, but a factory is opaque until called, and calling it at load time is what - this design avoids. Identity still covers the case the plugin packages - document -- the same provider callable both passed to the decorator and named - in ``DURABLE_EXECUTION_PLUGINS``. Two *different* factories that happen to - build the same plugin type will now both be registered. + it, but what a factory builds is unknown until ``create_plugin`` is called, + and calling it at load time is what this design avoids. Identity still covers + the case the plugin packages document -- the same factory object both passed + to the decorator and named in ``DURABLE_EXECUTION_PLUGINS``. Two *different* + factories that happen to build the same plugin type will now both be + registered. """ resolved_factories = _validate_explicit_factories(explicit_plugins) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py index 3cd9c806..311f8648 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -165,12 +165,13 @@ def test_operation_maps_across_suspend_and_replay(): built: list[_MapRecordingPlugin] = [] - def build_plugin(info: InvocationStartInfo) -> _MapRecordingPlugin: - plugin = _MapRecordingPlugin() - built.append(plugin) - return plugin + class _BuildingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _MapRecordingPlugin: + plugin = _MapRecordingPlugin() + built.append(plugin) + return plugin - @durable_execution(plugins=[build_plugin]) + @durable_execution(plugins=[_BuildingFactory()]) def wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 context.wait(Duration.from_seconds(60)) return "done" diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index bf19b8d2..7f30b255 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -3809,13 +3809,14 @@ def test_durable_execution_builds_a_plugin_per_invocation(): built: list[_RecordingPlugin] = [] factory_arns: list[str | None] = [] - def build_plugin(info) -> _RecordingPlugin: - factory_arns.append(info.execution_arn) - plugin = _RecordingPlugin() - built.append(plugin) - return plugin - - @durable_execution(plugins=[build_plugin]) + class _BuildingFactory: + def create_plugin(self, info) -> _RecordingPlugin: + factory_arns.append(info.execution_arn) + plugin = _RecordingPlugin() + built.append(plugin) + return plugin + + @durable_execution(plugins=[_BuildingFactory()]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3887,11 +3888,12 @@ def test_durable_execution_keeps_overlapping_invocations_isolated(): built: dict[str, _TaggedRecordingPlugin] = {} built_lock = threading.Lock() - def build_plugin(info) -> _TaggedRecordingPlugin: - plugin = _TaggedRecordingPlugin() - with built_lock: - built[str(info.request_id)] = plugin - return plugin + class _BuildingFactory: + def create_plugin(self, info) -> _TaggedRecordingPlugin: + plugin = _TaggedRecordingPlugin() + with built_lock: + built[str(info.request_id)] = plugin + return plugin timeout = 30 # Released only once both invocations are inside their user function, so both @@ -3902,7 +3904,7 @@ def build_plugin(info) -> _TaggedRecordingPlugin: b_ran_operation = threading.Event() a_ran_operation = threading.Event() - @durable_execution(plugins=[build_plugin]) + @durable_execution(plugins=[_BuildingFactory()]) def test_handler(event: Any, context: DurableContext) -> dict: tag = event["tag"] both_in_user_code.wait() @@ -3965,10 +3967,11 @@ def test_durable_execution_with_failing_plugin_factory_does_not_break_execution( recording_plugin = _RecordingPlugin() - def exploding_factory(info) -> DurableInstrumentationPlugin: - raise RuntimeError("factory boom") + class _ExplodingFactory: + def create_plugin(self, info) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") - @durable_execution(plugins=[exploding_factory, plugin_factory(recording_plugin)]) + @durable_execution(plugins=[_ExplodingFactory(), plugin_factory(recording_plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index 7936bd13..54b0f5be 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -33,12 +33,18 @@ class _PluginB(DurableInstrumentationPlugin): pass -def _plugin_a_factory(info: InvocationStartInfo) -> _PluginA: - return _PluginA() +class _PluginAFactory: + def create_plugin(self, info: InvocationStartInfo) -> _PluginA: + return _PluginA() + +class _PluginBFactory: + def create_plugin(self, info: InvocationStartInfo) -> _PluginB: + return _PluginB() -def _plugin_b_factory(info: InvocationStartInfo) -> _PluginB: - return _PluginB() + +_plugin_a_factory = _PluginAFactory() +_plugin_b_factory = _PluginBFactory() class _FakeDistribution: @@ -119,7 +125,8 @@ def test_discovery_returns_factories_without_calling_them() -> None: Nothing is constructed at load time, so no plugin instance exists outside the invocation that will use it. """ - factory = Mock(return_value=_PluginA()) + factory = Mock() + factory.create_plugin = Mock(return_value=_PluginA()) entry_point = _FakeEntryPoint("a", factory) with patch( @@ -132,7 +139,7 @@ def test_discovery_returns_factories_without_calling_them() -> None: ) assert result == [factory] - factory.assert_not_called() + factory.create_plugin.assert_not_called() def test_discovery_preserves_configured_order() -> None: @@ -151,7 +158,9 @@ def test_discovery_preserves_configured_order() -> None: ) assert result == [_plugin_a_factory, _plugin_b_factory] - assert [type(factory(INVOCATION_START_INFO)) for factory in result] == [ + assert [ + type(factory.create_plugin(INVOCATION_START_INFO)) for factory in result + ] == [ _PluginA, _PluginB, ] @@ -175,7 +184,7 @@ def test_explicit_factories_precede_discovered_factories() -> None: def test_explicit_registration_wins_over_the_same_discovered_factory( caplog: pytest.LogCaptureFixture, ) -> None: - """The same callable passed explicitly and named in the env registers once. + """The same factory passed explicitly and named in the env registers once. This is the narrowed form of the old type-based precedence rule. Dedup by declared plugin type is gone with the provider object; identity still covers @@ -206,14 +215,12 @@ def test_distinct_factories_for_one_plugin_type_are_both_registered() -> None: """Type-level dedup is gone: two distinct factories both register. Recorded deliberately. The provider object declared a ``plugin_type`` that - discovery could compare without constructing anything; a factory is opaque - until called, and calling it at load time would build an instance outside any - invocation. Callers that both pass a factory and name a different one in the - environment now get both plugins. + discovery could compare without constructing anything; what a factory builds + is unknown until ``create_plugin`` is called, and calling it at load time + would build an instance outside any invocation. Callers that both pass a + factory and name a different one in the environment now get both plugins. """ - - def another_plugin_a_factory(info: InvocationStartInfo) -> _PluginA: - return _PluginA() + another_plugin_a_factory = _PluginAFactory() entry_point = _FakeEntryPoint("a", another_plugin_a_factory) @@ -386,17 +393,18 @@ def test_discovery_names_unknown_distribution_in_load_failure() -> None: (object(), "builtins.object"), ("not-a-factory", "builtins.str"), (None, "builtins.NoneType"), + (lambda info: _PluginA(), "builtins.function"), ], ) -def test_discovery_rejects_non_callable_entry_point( +def test_discovery_rejects_entry_point_without_create_plugin( resolved_value: object, expected_type_name: str, ) -> None: """A plugin *instance* at the entry point is now the common mistake. - The old shape resolved to a provider object, so this replaces the - provider-type check with the only check that still means something: the - resolved value has to be callable. The message names what it actually was. + A bare callable is the other one, and it is rejected too: the registration + type is an object with ``create_plugin``, so a function that builds a plugin + no longer satisfies it. The message names what the target actually was. """ entry_point = _FakeEntryPoint("a", resolved_value) @@ -412,14 +420,13 @@ def test_discovery_rejects_non_callable_entry_point( environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "must resolve to a callable plugin factory, but resolved to" in str( - error.value - ) + assert "must resolve to a plugin factory" in str(error.value) + assert "create_plugin(info) method" in str(error.value) assert expected_type_name in str(error.value) -def test_discovery_accepts_a_plugin_class_as_factory() -> None: - """A class taking the info is callable, so it is a factory in its own right.""" +def test_discovery_rejects_a_plugin_class_at_the_entry_point() -> None: + """A plugin class carries no ``create_plugin``, so it is not a factory.""" class _InfoAwarePlugin(DurableInstrumentationPlugin): def __init__(self, info: InvocationStartInfo) -> None: @@ -427,27 +434,29 @@ def __init__(self, info: InvocationStartInfo) -> None: entry_point = _FakeEntryPoint("a", _InfoAwarePlugin) - with patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ), + pytest.raises(PluginLoadError) as error, ): - result = load_configured_plugins( + load_configured_plugins( None, environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - plugin = result[0](INVOCATION_START_INFO) - assert isinstance(plugin, _InfoAwarePlugin) - assert plugin.info is INVOCATION_START_INFO + assert "must resolve to a plugin factory" in str(error.value) + assert "not a plugin class" in str(error.value) def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: """A plugin instance in ``plugins`` fails configuration, not every invocation. - An instance is not callable, so the per-invocation factory call raises - ``TypeError``, which the executor logs and swallows -- the plugin silently - never runs. The position is asserted because a caller passing several entries - has no other way to tell which one is wrong. + An instance has no ``create_plugin``, so the per-invocation factory call + raises ``AttributeError``, which the executor logs and swallows -- the plugin + silently never runs. The position is asserted because a caller passing several + entries has no other way to tell which one is wrong. """ with pytest.raises(PluginLoadError) as error: load_configured_plugins( @@ -456,12 +465,12 @@ def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: ) assert "plugins[1]" in str(error.value) - assert "must be a callable plugin factory" in str(error.value) + assert "must be a plugin factory" in str(error.value) assert "_PluginB" in str(error.value) @pytest.mark.parametrize( - ("non_callable", "expected_type_name"), + ("invalid_entry", "expected_type_name"), [ (_PluginA(), "_PluginA"), (object(), "builtins.object"), @@ -469,62 +478,129 @@ def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: (None, "builtins.NoneType"), ], ) -def test_explicit_non_callable_entries_are_rejected( - non_callable: object, +def test_explicit_entries_without_create_plugin_are_rejected( + invalid_entry: object, expected_type_name: str, ) -> None: - """Callability is the only property checkable without building an instance.""" + """One callable member is all that is checkable without building an instance.""" with pytest.raises(PluginLoadError) as error: - load_configured_plugins([non_callable], environment={}) # type: ignore[list-item] + load_configured_plugins([invalid_entry], environment={}) # type: ignore[list-item] assert "plugins[0]" in str(error.value) assert expected_type_name in str(error.value) -def test_explicit_plugin_class_is_accepted_as_a_factory() -> None: - """Calling a class constructs an instance, so a class is a factory in Python. +def test_explicit_non_callable_create_plugin_is_rejected() -> None: + """The attribute is tested for callability, not merely for presence. + + An object whose ``create_plugin`` is data would otherwise pass here and raise + ``TypeError`` on every invocation, where it is logged and swallowed. + """ + + class _NotAFactory: + create_plugin = "not callable" + + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_NotAFactory()], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "_NotAFactory" in str(error.value) + + +@pytest.mark.parametrize( + "bare_callable", + [ + lambda info: _PluginA(), + _PluginAFactory.create_plugin, + ], +) +def test_explicit_bare_callable_is_rejected(bare_callable: object) -> None: + """A callable is no longer a factory, which reverses the previous rule. + + The registration type was ``Callable[[InvocationStartInfo], + DurableInstrumentationPlugin]``, so a lambda or a plain function was a valid + factory. It is now an object with ``create_plugin``, and no compatibility + path accepts both: a bare callable fails at handler initialization with + guidance naming the replacement. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([bare_callable], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "must be a plugin factory" in str(error.value) + assert "plugins=[MyPluginFactory(exporter)]" in str(error.value) + + +def test_explicit_plugin_class_is_rejected_as_a_factory() -> None: + """A plugin class is not a factory, reversing what this branch documented. - This is the language difference against the TypeScript SDK, which rejects a - class because calling one there throws. Rejecting classes here would break - ``plugins=[MyPlugin]``, which works whenever ``__init__`` takes the info. + Calling a class constructs an instance, so a class satisfied the previous + ``Callable`` registration type and ``plugins=[MyPlugin]`` was accepted. A + class carries no ``create_plugin`` attribute, so it is now rejected at + handler initialization. The replacement is a factory class whose + ``create_plugin`` constructs the plugin, which also keeps setup work out of + the plugin's ``__init__``. """ class _InfoAwarePlugin(DurableInstrumentationPlugin): def __init__(self, info: InvocationStartInfo) -> None: self.info = info - result = load_configured_plugins([_InfoAwarePlugin], environment={}) + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_InfoAwarePlugin], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "a plugin class" in str(error.value) + + +def test_explicit_class_declaring_create_plugin_is_accepted() -> None: + """The requirement is the member, not the kind of object. + + A class that declares ``create_plugin`` as a ``@classmethod`` carries the + attribute, so the class object itself is a valid factory. Nothing in the + contract requires a factory to be an instance. + """ + + class _ClassFactoryPlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + @classmethod + def create_plugin(cls, info: InvocationStartInfo) -> _ClassFactoryPlugin: + return cls(info) + + result = load_configured_plugins([_ClassFactoryPlugin], environment={}) - assert result == [_InfoAwarePlugin] - plugin = result[0](INVOCATION_START_INFO) - assert isinstance(plugin, _InfoAwarePlugin) + assert result == [_ClassFactoryPlugin] + plugin = result[0].create_plugin(INVOCATION_START_INFO) + assert isinstance(plugin, _ClassFactoryPlugin) assert plugin.info is INVOCATION_START_INFO -def test_explicit_plain_function_is_accepted_as_a_factory() -> None: +def test_explicit_factory_object_is_accepted() -> None: result = load_configured_plugins([_plugin_a_factory], environment={}) assert result == [_plugin_a_factory] - assert isinstance(result[0](INVOCATION_START_INFO), _PluginA) + assert isinstance(result[0].create_plugin(INVOCATION_START_INFO), _PluginA) -def test_explicit_callable_object_is_accepted_as_a_factory() -> None: - """A ``__call__`` instance is the documented way to hold handler-lifetime state.""" +def test_explicit_factory_holds_handler_lifetime_state() -> None: + """A factory instance is where state spanning invocations belongs.""" - class _CallableFactory: + class _StatefulFactory: def __init__(self) -> None: self.calls: list[InvocationStartInfo] = [] - def __call__(self, info: InvocationStartInfo) -> _PluginA: + def create_plugin(self, info: InvocationStartInfo) -> _PluginA: self.calls.append(info) return _PluginA() - factory = _CallableFactory() + factory = _StatefulFactory() result = load_configured_plugins([factory], environment={}) assert result == [factory] - assert isinstance(result[0](INVOCATION_START_INFO), _PluginA) + assert isinstance(result[0].create_plugin(INVOCATION_START_INFO), _PluginA) assert factory.calls == [INVOCATION_START_INFO] diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 55095d0e..b2ee6e27 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -22,6 +22,7 @@ ) from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, InvocationEndInfo, InvocationInfo, InvocationStatus, @@ -37,6 +38,7 @@ UserFunctionOutcome, UserFunctionStartInfo, ) +from aws_durable_execution_sdk_python.plugin_discovery import _is_plugin_factory from tests.test_helpers import plugin_factory @@ -555,6 +557,38 @@ def test_subclass_override(self): # endregion DurableInstrumentationPlugin Tests +# region DurableInstrumentationPluginFactory Tests +class TestDurableInstrumentationPluginFactory(unittest.TestCase): + def test_protocol_is_not_runtime_checkable(self): + """``isinstance`` against the protocol must stay unavailable. + + A runtime-checkable protocol requires every declared member, so an + optional second member added later would break any caller's isinstance + check -- the additive extensibility this protocol exists for. The SDK + checks the shape with ``plugin_discovery._is_plugin_factory`` instead, + which also tests that ``create_plugin`` is callable rather than merely + present. + """ + with self.assertRaises(TypeError): + isinstance( # type: ignore[misc] # noqa: B018 + plugin_factory(_NoOpPlugin()), DurableInstrumentationPluginFactory + ) + + def test_factory_shape_check_requires_a_callable_create_plugin(self): + """The shape is one callable member, structurally, without importing it.""" + + class _Data: + create_plugin = "not callable" + + self.assertTrue(_is_plugin_factory(plugin_factory(_NoOpPlugin()))) + self.assertFalse(_is_plugin_factory(_Data())) + self.assertFalse(_is_plugin_factory(lambda info: _NoOpPlugin())) + self.assertFalse(_is_plugin_factory(_NoOpPlugin())) + + +# endregion DurableInstrumentationPluginFactory Tests + + # region PluginExecutor Tests @@ -585,14 +619,15 @@ def test_each_invocation_gets_its_own_instance(self): """Two invocations of one handler never share a plugin instance.""" built: list[_TrackingPlugin] = [] - def build(info: InvocationStartInfo) -> _TrackingPlugin: - plugin = _TrackingPlugin() - built.append(plugin) - return plugin + class _BuildingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: + plugin = _TrackingPlugin() + built.append(plugin) + return plugin # One host for the handler, one executor per invocation -- the shape # durable_execution() uses. - host = PluginHost(plugins=[build]) + host = PluginHost(plugins=[_BuildingFactory()]) for request_id in ("req-1", "req-2"): lambda_context = MagicMock() @@ -659,11 +694,12 @@ class _RecordingPlugin(DurableInstrumentationPlugin): def on_invocation_start(self, info: InvocationStartInfo) -> None: hook_infos.append(info) - def build(info: InvocationStartInfo) -> _RecordingPlugin: - factory_infos.append(info) - return _RecordingPlugin() + class _RecordingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _RecordingPlugin: + factory_infos.append(info) + return _RecordingPlugin() - executor = PluginExecutor(plugins=[build]) + executor = PluginExecutor(plugins=[_RecordingFactory()]) with executor.run(): executor.on_invocation_start( @@ -691,14 +727,15 @@ def __init__(self, label: str) -> None: def on_invocation_start(self, info: InvocationStartInfo) -> None: events.append(f"hook:{self.label}") - def build(label: str): - def factory(info: InvocationStartInfo) -> _OrderedPlugin: - events.append(f"build:{label}") - return _OrderedPlugin(label) + class _OrderedFactory: + def __init__(self, label: str) -> None: + self._label = label - return factory + def create_plugin(self, info: InvocationStartInfo) -> _OrderedPlugin: + events.append(f"build:{self._label}") + return _OrderedPlugin(self._label) - executor = PluginExecutor(plugins=[build("a"), build("b")]) + executor = PluginExecutor(plugins=[_OrderedFactory("a"), _OrderedFactory("b")]) with executor.run(): executor.on_invocation_start( @@ -711,14 +748,11 @@ def factory(info: InvocationStartInfo) -> _OrderedPlugin: self.assertEqual(events, ["build:a", "build:b", "hook:a", "hook:b"]) def test_failing_factory_is_contained(self): - """A raising factory is logged and skipped, like a raising hook.""" + """A factory whose create_plugin raises is logged and skipped.""" surviving = _TrackingPlugin() - def exploding(info: InvocationStartInfo) -> DurableInstrumentationPlugin: - raise RuntimeError("factory boom") - executor = PluginExecutor( - plugins=[exploding, plugin_factory(surviving)], + plugins=[_ExplodingFactory(), plugin_factory(surviving)], ) with self.assertLogs( @@ -736,15 +770,48 @@ def exploding(info: InvocationStartInfo) -> DurableInstrumentationPlugin: # The other factory's plugin still receives its hooks. self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_bare_callable_builds_no_plugin(self): + """The executor calls ``create_plugin``, so a bare callable yields nothing. + + A callable used to be a factory. It is not one now, and an entry that + reaches the executor anyway raises ``AttributeError`` there, which the + executor logs and skips. ``load_configured_plugins`` rejects such an entry + at handler initialization, so this path is only reachable by constructing + an executor directly. + """ + surviving = _TrackingPlugin() + + executor = PluginExecutor( + plugins=[ + (lambda info: _TrackingPlugin()), # type: ignore[list-item] + plugin_factory(surviving), + ], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + self.assertIn("create_plugin", "\n".join(logs.output)) + def test_factory_returning_none_is_contained(self): """A factory that returns nothing is logged and skipped.""" surviving = _TrackingPlugin() - def returns_none(info: InvocationStartInfo): - return None + class _NoneFactory: + def create_plugin(self, info: InvocationStartInfo): + return None executor = PluginExecutor( - plugins=[returns_none, plugin_factory(surviving)], + plugins=[_NoneFactory(), plugin_factory(surviving)], ) with self.assertLogs( @@ -764,11 +831,7 @@ def returns_none(info: InvocationStartInfo): def test_every_failing_factory_leaves_the_executor_usable(self): """All factories failing is not distinguishable from having no plugins.""" - - def exploding(info: InvocationStartInfo) -> DurableInstrumentationPlugin: - raise RuntimeError("factory boom") - - executor = PluginExecutor(plugins=[exploding]) + executor = PluginExecutor(plugins=[_ExplodingFactory()]) with self.assertLogs( "aws_durable_execution_sdk_python.plugin", level=logging.ERROR @@ -2063,6 +2126,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: self.calls.append(f"user_function_end:{info.operation_id}") +class _ExplodingFactory: + """Factory whose ``create_plugin`` raises, for containment tests.""" + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py index de677280..f77a11c2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py +++ b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py @@ -9,6 +9,7 @@ from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, DurableInstrumentationPluginFactory, + InvocationStartInfo, PluginExecutor, ) @@ -27,17 +28,27 @@ def operation_id_sequence(parent_id: str | None = None): yield context._create_step_id() # noqa: SLF001 -def plugin_factory( - plugin: DurableInstrumentationPlugin, -) -> DurableInstrumentationPluginFactory: - """Wrap a plugin instance a test already holds a reference to as a factory. +class _FixedPluginFactory: + """Returns one plugin instance the test already holds, for every invocation. Production factories build a fresh instance per invocation. A test that has to read what the plugin recorded needs the instance it passed in, so it supplies a factory that returns that one. Only valid for a single invocation, which is all these tests run. """ - return lambda info: plugin + + def __init__(self, plugin: DurableInstrumentationPlugin) -> None: + self._plugin = plugin + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + return self._plugin + + +def plugin_factory( + plugin: DurableInstrumentationPlugin, +) -> DurableInstrumentationPluginFactory: + """Wrap a plugin instance a test already holds a reference to as a factory.""" + return _FixedPluginFactory(plugin) @contextlib.contextmanager From 750db4828f4da5201e2b2bc76d10f2ed3669a9e1 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 11:00:42 -0700 Subject: [PATCH 08/28] fix(insight): drop a superseded record; lock wrapper installs Two findings that had been accepted and deferred. Both are real, and the first one left Python as the only SDK without a guard the other two carry, which is backwards for a finding that originated here. Record building runs customer callbacks while the per-execution lock is held: the content transforms and an operation's result transform. A callback can re-enter `on_operation_change` on the same thread, build a newer record and schedule it, after which the outer frame resumes and schedules its older snapshot. Measured with a transform that re-enters while the outer frame holds an empty operation map: with export deferred the exporter saw `[[]]`, so the newer `['s1']` snapshot was coalesced away in the pending slot; with export immediate it saw `[['s1'], []]`, so the stale snapshot arrived after the newer one. Each invocation now counts the record builds it has started. A build takes the next value before it begins and the hand-off queues the record only while that value is still the newest, so an overtaken build's record is dropped. Nothing is lost, because a record is a complete snapshot of one execution and the record that superseded it carries everything the older one carried. That is the same property that makes per-execution coalescing sound. The counter is a plain `int` rather than an atomic, which differs from the Java port deliberately. Every increment and every read happens while `self._lock` is held, because `_emit` is only ever called from inside a `with self._lock` block. So the read-modify-write cannot interleave. Java needs an `AtomicLong` because two of its builds can genuinely run at once; in Python a second thread entering a hook blocks on that lock, so the only builds in flight together are nested frames on one thread that already hold it. The revalidation sits beside the existing closed-gate re-check, inside the same lock hold that queues the record, so no build can start, finish and queue between the check and the queue. The closing record takes no revision and is exempt from the check. One honest note: in Python that exemption is not reachable through a hook, because `on_invocation_end` sets the closed gate before the closing build and all three hooks check the gate before reaching `_emit`, so nothing can advance the counter during that build. The exemption is kept for parity with the JS and Java ports and so the closing record's survival does not depend on that gating holding for every future hook. It is pinned by mutation rather than by a reachable scenario: dropping the exemption fails exactly the test that asserts the closing record still arrives. Second, the OTel wrapper installs were a check-then-set with no lock, and the tracer they reach into is shared. `TracerProvider.get_tracer` caches by instrumentation scope, so two plugins built for two concurrent first invocations receive the same `Tracer` object; `get_tracer(name) is get_tracer(name)` is True. Two installs could therefore each wrap the provider's original, and the loser would keep a wrapper the tracer no longer holds, so its deterministic overrides are ignored and span ids become random, which breaks cross-invocation stitching. The window is a few bytecodes, so 400 trials at the default switch interval produced nothing. With the switch interval lowered to make the interpreter preempt inside that existing window, 448 of 2000 trials produced rival id generators, 191 of 2000 produced rival samplers, and 11 of 200 left a plugin holding an id generator or sampling delegate the tracer no longer held, for each factory class. After the fix all four measurements are zero. Both installs now hold a module-level lock across the read, the check, the construction and the assignment, and each returns whatever the tracer holds, so a plugin that loses the race uses the winner's wrapper and the winner's delegate. This follows the precedent already in this package: the log filter's install is guarded the same way for the same reason. --- .../plugin.py | 82 +++++- .../tests/test_plugin.py | 272 ++++++++++++++++++ .../deterministic_id_generator.py | 39 ++- .../durable_sampling.py | 33 ++- .../plugin_factory.py | 8 +- .../tests/test_plugin_factory.py | 96 +++++++ 6 files changed, 503 insertions(+), 27 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index f80bca97..572e62d7 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -193,8 +193,8 @@ class WorkflowInsightPlugin(DurableInstrumentationPlugin, _ExportState): Two locks, disjoint field sets, and neither is ever taken to reach the other's fields: - * ``_lock`` guards ``_closed``, the ``_operations`` rebind and record - emission. + * ``_lock`` guards ``_closed``, ``_build_revision``, the ``_operations`` + rebind and record emission. * ``_ExportScheduler._condition``'s lock guards the export bookkeeping this instance carries for the scheduler; nothing outside the scheduler reads or writes those fields, and the scheduler never touches the fields above. @@ -236,13 +236,34 @@ def __init__( # operation-change for a checkpoint that completed just before the end) # must emit nothing (mirrors the Java ExecutionState.closed flag). self._closed = False - # Guards `_closed`, the operations rebind and record emission, so a late - # hook can never slip a RUNNING record in after the terminal one. Still - # earns its place with one instance per invocation: the SDK dispatches - # every hook synchronously on the thread that produced the event, so an - # operation-change raised off the checkpointing path runs concurrently - # with the invocation thread's on_invocation_end -- two hooks, one - # instance, genuinely racing. + # Counts the record builds this instance has started. Never decremented. + # + # A record is a complete snapshot of one execution, so the scheduler's + # per-execution slot takes whichever record is handed over last and never + # compares ages. The build that hands its record over last is not the + # build that started last: `_emit` runs customer code (the input/output + # transforms and the operation result overrides) between the snapshot it + # takes and the hand-off, and that code can re-enter a hook on this thread + # and complete a newer build first. Without a comparison of build ages the + # outer frame's older snapshot then replaces the newer one in the slot, or + # is exported after it. Every non-terminal build takes the next value here + # before it starts, and `_emit` hands the record over only while that value + # is still the newest (mirrors the JS `buildRevision` and the Java + # `AtomicLong buildRevision`). + # + # A plain int, not an atomic: every read and every increment happens under + # `_lock`, so the read-modify-write cannot interleave, and the two builds + # this counter distinguishes are nested frames on one thread rather than + # two threads. Java needs an AtomicLong because its two builds really can + # run at once. + self._build_revision = 0 + # Guards `_closed`, `_build_revision`, the operations rebind and record + # emission, so a late hook can never slip a RUNNING record in after the + # terminal one. Still earns its place with one instance per invocation: the + # SDK dispatches every hook synchronously on the thread that produced the + # event, so an operation-change raised off the checkpointing path runs + # concurrently with the invocation thread's on_invocation_end -- two hooks, + # one instance, genuinely racing. # # Reentrant on purpose: `_emit` runs the scheduler's `schedule()` inside # this hold, and `schedule()` releases the record it displaces, which can @@ -421,6 +442,19 @@ def _emit( # cannot change the map mid-build. operations = self._operations + # The revision is taken here, before the build, never after. Customer code + # runs inside the build below and can re-enter a hook on this thread, + # which starts and finishes a newer build. A value read after the build + # would already be that newer build's, so this older record would pass the + # check and replace the newer one. Callers hold self._lock, so the + # increment cannot interleave with another build's. + # + # The closing emit takes no revision; see the hand-off below. + revision = 0 + if not closing: + self._build_revision += 1 + revision = self._build_revision + content = self._shared._content record: dict[str, Any] = { "recordType": "WorkflowInsight", @@ -492,7 +526,35 @@ def _emit( # itself. It is not the same test as "the record is terminal": in # on-change mode a PENDING/RETRY invocation end legitimately emits a # RUNNING record, and that record is the closing one. - if self._closed and not closing: + # + # INVARIANT: the record handed to the scheduler below is the newest build + # this instance has started -- an exporter never stores an older snapshot + # over a newer one for the same execution. + # + # A build that customer code started from inside this one has already + # handed its own, newer record over by the time control returns here, so + # this record is superseded. Dropping it loses nothing: a record is a + # complete snapshot of one execution, so the newer record carries + # everything this one carries. That is the same property that makes the + # scheduler's per-execution coalescing sound. + # + # The check and the hand-off are one critical section. self._lock is held + # across both -- schedule() is called inside this hold -- and every build + # takes that same lock, so no build can start, finish and queue its record + # between this check and the schedule() below. A record that passes + # therefore cannot be queued after the record that supersedes it. (Java + # revalidates inside the scheduler's monitor instead, because there the + # scheduler's monitor is what guards `closed` and the record slot; here + # both facts already belong to self._lock.) + # + # The closing record is exempt from the revision check. Customer code + # inside its build can start a newer non-terminal build, which would make a + # revision taken here stale, and a checked hand-off would then drop the + # closing record and leave a RUNNING snapshot as this execution's last + # exported state. The exemption cannot let a stale record win, because + # `_closed` was set in the same self._lock hold that queues this record and + # every later non-terminal record is rejected above. + if not closing and (self._closed or revision != self._build_revision): return self._shared._scheduler.schedule(self, record) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 8cbf23ea..0e3c9ce2 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -1123,3 +1123,275 @@ def hook() -> None: f"the same execution: {statuses}" ) assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + +# -- a build overtaken by one customer code started from inside it ------------- + + +class GatedCaptureExporter: + """Records every export; optionally blocks inside one execution's export. + + Blocking there pins the single export worker, so records scheduled while it + is held stay in their execution's pending slot and coalesce there. + """ + + max_record_size_bytes = None + + def __init__(self, hold_arn: str | None = None) -> None: + self._hold_arn = hold_arn + self.holding = threading.Event() + self.release = threading.Event() + self._lock = threading.Lock() + self._arrived = threading.Condition(self._lock) + self._records: list[dict[str, Any]] = [] + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + if self._hold_arn is not None and record["executionArn"] == self._hold_arn: + self.holding.set() + self.release.wait(30.0) + with self._arrived: + self._records.append(record) + self._arrived.notify_all() + + def flush(self) -> None: + pass + + def snapshot(self) -> list[dict[str, Any]]: + with self._lock: + return list(self._records) + + def wait_for(self, arn: str, count: int, timeout: float = 10.0) -> bool: + with self._arrived: + return self._arrived.wait_for( + lambda: ( + sum(1 for r in self._records if r["executionArn"] == arn) >= count + ), + timeout, + ) + + +def _records_for(exporter: GatedCaptureExporter, arn: str) -> list[dict[str, Any]]: + return [record for record in exporter.snapshot() if record["executionArn"] == arn] + + +def _force_drain_bounded(factory, timeout: float = 20.0) -> None: + """``_force_drain`` on a bounded thread, so a stalled drain fails the test. + + ``drain()`` parks on a condition with no deadline: that is correct in + production, where the only thing that can release it is the export worker, but + a regression that loses a record leaves it parked for good. Running it on + another thread turns that into an assertion failure instead of a hung suite. + """ + drained = threading.Event() + + def drain() -> None: + _force_drain(factory) + drained.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert drained.wait(timeout), "the drain never returned" + thread.join(5.0) + assert not thread.is_alive() + + +def _exported_operation_names( + exporter: GatedCaptureExporter, arn: str +) -> list[list[str]]: + return [ + [op["name"] for op in record["operations"]] + for record in _records_for(exporter, arn) + ] + + +def _newer_change() -> OperationChangeInfo: + """An operation-change carrying a strictly newer map than the drives below start from.""" + return OperationChangeInfo( + execution_arn=ARN, + updated_operations=_ops(_step("s1")), + operations=_ops(_step("s1")), + ) + + +def test_a_build_overtaken_by_a_nested_one_does_not_coalesce_the_newer_away(): + # _emit runs customer code -- here the content.input transform -- between the + # operation snapshot it takes and the hand-off to the scheduler, and the + # execution's lock is reentrant, so that code can run on_operation_change to + # completion on this same thread. The nested hook adopts a newer operation map + # and schedules its record first. The outer frame then hands over the older + # snapshot it built, and the pending slot takes whichever record arrives last, + # so without a comparison of build ages the newer record is coalesced away and + # the exporter only ever sees the older one. One hook call, one instance, no + # concurrency. + exporter = GatedCaptureExporter(hold_arn=ARN_B) + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + # The primer execution below runs this same transform, so re-enter only + # once the instance under test has been published. + if holder.get("plugin") is None or reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + # Pin the one export worker on another execution's record, so both records + # this execution builds are still in its pending slot when the outer frame + # hands its own over. + primer = factory.create_plugin(_start(arn=ARN_B)) + primer.on_invocation_start(_start(arn=ARN_B)) + assert exporter.holding.wait(10.0), "the export worker never reached the exporter" + + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + # On a bounded thread, so a regression that deadlocks the hook fails here + # instead of hanging the suite. + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() # the nested change hook really did run + exporter.release.set() + _force_drain_bounded(factory) + + assert _exported_operation_names(exporter, ARN) == [["s1"]], ( + "the newer snapshot the nested hook built was coalesced away by the " + "older one the outer frame built: " + f"{_exported_operation_names(exporter, ARN)}" + ) + + +def test_a_build_overtaken_by_a_nested_one_is_not_exported_after_it(): + # Same overtaking, with the newer record already handed to the exporter before + # the outer frame reaches the hand-off. Nothing coalesces, so without a + # comparison of build ages the exporter sees the newer snapshot and then the + # older one, and an exporter that upserts by execution ARN ends up storing the + # older state. + exporter = GatedCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + # Wait for the newer record to reach the exporter, so this frame's older + # record cannot displace it in the pending slot and the two records are + # ordered at the exporter instead. + assert exporter.wait_for(ARN, 1), "the nested record never reached the exporter" + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + _force_drain_bounded(factory) + + assert _exported_operation_names(exporter, ARN) == [["s1"]], ( + "the exporter saw the older snapshot after the newer one for the same " + f"execution: {_exported_operation_names(exporter, ARN)}" + ) + + +def test_the_closing_record_is_exempt_from_the_build_age_check(): + # The closing record must reach the exporters whatever the build ages say. + # Customer code inside its build can start a newer non-terminal build, which + # would leave the closing record's age stale, and a checked hand-off would + # then drop it and leave a RUNNING snapshot as this execution's last exported + # state. Nothing else can rescue it: it is the last record this instance ever + # builds. + # + # The drive raises the build counter above the closing record's own first, by + # overtaking one build with a nested one exactly as the two tests above do, so + # a closing record that is age-checked against a counter it never incremented + # is dropped here. + exporter = GatedCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + assert exporter.wait_for(ARN, 1), "the nested record never reached the exporter" + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + # Push whatever the start hook handed over to the exporter before the end + # hook schedules the closing record, so a record the outer frame scheduled + # is observed here instead of being coalesced away by the closing one. + _force_drain(factory) + # on_invocation_end drains, so every record is at the exporter once this + # returns. + plugin.on_invocation_end(_end(operations=_ops(_step("s1"), _step("s2")))) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hooks never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + + records = _records_for(exporter, ARN) + assert [record["status"] for record in records] == ["RUNNING", "SUCCEEDED"], ( + "the closing record was dropped, or the superseded RUNNING record was " + f"exported: {[record['status'] for record in records]}" + ) + assert [op["name"] for op in records[0]["operations"]] == ["s1"] + assert sorted(op["name"] for op in records[1]["operations"]) == ["s1", "s2"] diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py index c3d7afc1..1386dd2e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py @@ -4,6 +4,7 @@ import contextvars import hashlib +import threading from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass @@ -23,6 +24,17 @@ class _IdOverride: span_id: int | None +# Serializes installation so two invocations binding to one tracer at the same +# time cannot both wrap its original generator. A TracerProvider caches tracers +# by instrumentation scope, so two plugins that ask for the same instrument name +# get the same tracer object. Without this lock both can read the original +# generator, both wrap it, and the second assignment replaces the first: the +# plugin that assigned first then holds a wrapper the tracer no longer uses, its +# deterministic overrides are never consulted, and its workflow and operation +# span IDs are random, which breaks cross-invocation stitching. +_install_lock = threading.Lock() + + def _to_otel_trace_id(execution_arn: str, start_timestamp: datetime) -> int: """Build a deterministic OTel-compatible execution trace ID (128 bits). @@ -123,16 +135,25 @@ def install_on_tracer(cls, tracer: SdkTracer) -> DeterministicIdGenerator: """Return the tracer's deterministic generator, installing one if needed. Installing on the plugin's tracer keeps unrelated instrumentation scopes - on the provider's original generator. Reusing an installed generator also - supports SDK versions that cache tracers by instrumentation scope. + on the provider's original generator. + + The returned generator is always the one the tracer holds. The check and + the install are one critical section, so a caller that arrives while + another is installing waits and then finds the installed generator + instead of wrapping the original a second time. A caller that acted on a + generator the tracer does not hold would set its deterministic overrides + somewhere the tracer never reads. Returning the installed generator also + supports SDK versions that cache tracers by instrumentation scope, which + is what makes two invocations share one tracer in the first place. """ - current_generator = tracer.id_generator - if isinstance(current_generator, cls): - return current_generator - - generator = cls(fallback_id_generator=current_generator) - tracer.id_generator = generator - return generator + with _install_lock: + current_generator = tracer.id_generator + if isinstance(current_generator, cls): + return current_generator + + generator = cls(fallback_id_generator=current_generator) + tracer.id_generator = generator + return generator @contextmanager def use_ids(self, *, trace_id: int | None, span_id: int | None) -> Iterator[None]: diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py index 9d2fa610..9689b341 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py @@ -4,6 +4,7 @@ import functools import inspect +import threading from dataclasses import dataclass from typing import Any @@ -24,6 +25,16 @@ ) +# Serializes installation so two invocations binding to one tracer at the same +# time cannot both wrap its original sampler. A TracerProvider caches tracers by +# instrumentation scope, so two plugins that ask for the same instrument name get +# the same tracer object. Without this lock both can read the original sampler, +# both wrap it, and the second assignment replaces the first: the plugin that +# assigned first then holds a wrapper the tracer no longer uses, and the delegate +# it took from that wrapper is not the delegate the tracer consults. +_install_lock = threading.Lock() + + @dataclass(frozen=True) class DurableSamplingIntent: """Sampling result to apply to each durable span in one invocation.""" @@ -39,12 +50,22 @@ def __init__(self, delegate: Sampler) -> None: @classmethod def install_on_tracer(cls, tracer: SdkTracer) -> "DurableSampler": - current_sampler = tracer.sampler - if isinstance(current_sampler, cls): - return current_sampler - sampler = cls(current_sampler) - tracer.sampler = sampler - return sampler + """Return the tracer's durable sampler, installing one if needed. + + The returned sampler is always the one the tracer holds. The check and the + install are one critical section, so a caller that arrives while another + is installing waits and then finds the installed sampler instead of + wrapping the original a second time. A caller that kept a sampler the + tracer does not hold would also keep a delegate the tracer never + consults. + """ + with _install_lock: + current_sampler = tracer.sampler + if isinstance(current_sampler, cls): + return current_sampler + sampler = cls(current_sampler) + tracer.sampler = sampler + return sampler def should_sample( self, diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py index e5bbfc51..028b9998 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py @@ -15,8 +15,12 @@ itself: the tracer provider (which for the global-provider case may only be installed after the handler module is imported), the tracer, and the deterministic id generator and sampler installed on it. Those installs are -idempotent and scoped to the plugin's own tracer, so building a plugin per -invocation neither stacks wrappers nor disturbs other instrumentation scopes. +scoped to the plugin's own tracer, so building a plugin per invocation does not +disturb other instrumentation scopes. They are also atomic and return whatever +the tracer holds, which matters because a provider caches tracers by +instrumentation scope: two invocations starting at once get one tracer, and a +plugin that lost the install race must use the wrapper on that tracer rather than +one of its own that the tracer would never consult. """ from __future__ import annotations diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py index 32931438..74ca4dea 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py @@ -12,6 +12,8 @@ import importlib import logging +import sys +import threading import tomllib from datetime import UTC, datetime from pathlib import Path @@ -21,7 +23,13 @@ DurableInstrumentationPlugin, InvocationStartInfo, ) +from opentelemetry.sdk.trace import Tracer as SdkTracer +from opentelemetry.sdk.trace import TracerProvider +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import DurableSampler from aws_durable_execution_sdk_python_otel.execution_plugin import ( ExecutionOtelPlugin, ) @@ -178,3 +186,91 @@ def test_declared_entry_points_resolve_to_factories_that_build_plugins() -> None assert callable(create_plugin) assert not isinstance(factory, type) assert isinstance(create_plugin(_invocation_start_info()), expected[name]) + + +# -- concurrent first invocations sharing one cached tracer -------------------- + +# Trials per factory. A single trial reproduces the loss roughly one time in ten +# (measured at 16 threads with the switch interval below), so a handful of trials +# would pass with an unsynchronized install still in place. +_SHARED_TRACER_TRIALS = 200 + +# Concurrent create_plugin calls per trial, standing in for concurrent first +# invocations in one Lambda Managed Instances environment. +_SHARED_TRACER_PLUGINS = 16 + + +def _build_plugins_concurrently( + factory_type: type, provider: TracerProvider, count: int +) -> list[InvocationOtelPlugin | ExecutionOtelPlugin]: + config = OtelPluginConfig(tracer_provider=provider, enrich_logger=False) + factory = factory_type(config) + ready = threading.Barrier(count) + plugins: list[InvocationOtelPlugin | ExecutionOtelPlugin | None] = [None] * count + + def build(index: int) -> None: + ready.wait(10.0) + plugins[index] = factory.create_plugin(_invocation_start_info()) + + workers = [threading.Thread(target=build, args=(index,)) for index in range(count)] + for worker in workers: + worker.start() + for worker in workers: + worker.join(10.0) + assert not any(worker.is_alive() for worker in workers) + built = [plugin for plugin in plugins if plugin is not None] + assert len(built) == count + return built + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_concurrent_plugins_share_the_wrappers_their_tracer_holds( + factory_type: type, +) -> None: + """Every plugin must hold the generator and sampler the tracer actually uses. + + A TracerProvider caches tracers by instrumentation scope, so plugins built for + concurrent first invocations ask for one instrument name and get one tracer + object. Installing the deterministic generator and the durable sampler on that + tracer is a check-then-set: two plugins can each read the original generator, + each wrap it, and the second assignment replaces the first. The plugin that + assigned first then holds a wrapper the tracer no longer consults, so its + deterministic ID overrides are ignored and its workflow and operation span IDs + come out random, which breaks cross-invocation stitching. + + The switch interval is lowered so the interpreter preempts inside that + check-then-set often enough for the loss to appear within the trial count; it + does not create the window, it only makes an existing one likely to be hit. + """ + previous_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-9) + try: + for trial in range(_SHARED_TRACER_TRIALS): + provider = TracerProvider() + plugins = _build_plugins_concurrently( + factory_type, provider, _SHARED_TRACER_PLUGINS + ) + tracer = provider.get_tracer(OtelPluginConfig().instrument_name) + assert isinstance(tracer, SdkTracer) + # One tracer for every plugin: the premise the rest of the assertions + # rest on, and the reason a lost install is not simply harmless. + assert all(plugin._tracer is tracer for plugin in plugins) + assert isinstance(tracer.id_generator, DeterministicIdGenerator) + assert isinstance(tracer.sampler, DurableSampler) + # Exactly one wrapper of each kind exists, and it is the tracer's. + assert {id(plugin._id_generator) for plugin in plugins} == { + id(tracer.id_generator) + }, f"trial {trial}: a plugin holds an id generator the tracer discarded" + assert {id(plugin._sampling_delegate) for plugin in plugins} == { + id(tracer.sampler.delegate) + }, f"trial {trial}: a plugin holds a sampling delegate the tracer discarded" + # The wrapper wraps the provider's original, not another wrapper. + assert not isinstance(tracer.sampler.delegate, DurableSampler) + assert not isinstance( + tracer.id_generator._fallback_id_generator, DeterministicIdGenerator + ) + finally: + sys.setswitchinterval(previous_interval) From f8a5678564c8610036100fd45effcefc57a8e595 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 11:06:24 -0700 Subject: [PATCH 09/28] fix(plugin): reject a factory that returns a non-plugin Two review findings, one behavioural and one documentation. The load-time shape check establishes only that a factory has a callable `create_plugin`. What that call returns is knowable only when it is called, and `_create_plugins` rejected `None` but accepted every other object. So a factory returning `object()` was registered, and it then failed every hook inside `_dispatch_plugin`. The consequence is one logged error per hook per invocation for the life of the function, while that plugin provides no telemetry at all. The return is now checked against `DurableInstrumentationPlugin` and an invalid one is logged once and skipped, which is the same containment the `None` case already had. Every plugin in this repository already subclasses that base class, including all 25 in the conformance handlers, so the check rejects nothing that worked before. Second, the Insight README described a flush cadence the plugin no longer has. It said an invocation that emits no record neither starts nor flushes the worker. Every sampled-in invocation end now drains and flushes, whether or not it emitted a record, which is the cadence the JS and Java plugins have and which the tests already assert. Only a sampled-out execution neither exports nor flushes. The note now says that. --- .../README.md | 9 +++-- .../plugin.py | 14 +++++++ .../tests/plugin_test.py | 38 +++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 5aafb9ac..2d3c9289 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -223,9 +223,12 @@ Behavior is validated cross-SDK by the `insight` conformance suite > **Note (asynchronous export).** Export rendering, truncation, `export()`, and > `flush()` run on one lazy background worker per registered factory. Checkpoint > hooks only replace the latest pending snapshot and wake the worker. Consecutive -> `on-change` snapshots may coalesce while an export is in flight. An invocation -> that emits a record drains the latest snapshot and flushes exporters before it -> returns; invocations that emit nothing do not start or flush the worker. +> `on-change` snapshots may coalesce while an export is in flight. Every +> sampled-in invocation end drains the latest snapshot and flushes exporters +> before it returns, including an end that emitted no record: a buffering +> exporter therefore sees one flush per sampled-in invocation end, which is the +> cadence the JS and Java plugins have. Only a sampled-out execution neither +> exports nor flushes. ## Requirements 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 e7acdd1f..4e0b31e4 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 @@ -627,6 +627,20 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: _factory_name(factory), ) continue + # The load-time shape check can only establish that the factory has + # a callable create_plugin; what that call returns is knowable only + # here. A value that is not a plugin fails every hook inside + # _dispatch_plugin, so registering it would produce one logged error + # per hook per invocation for the life of the function while + # providing no telemetry. Reject it once instead. + if not isinstance(plugin, DurableInstrumentationPlugin): + logger.error( + "Plugin factory %s returned %s, which is not a " + "DurableInstrumentationPlugin; plugin ignored", + _factory_name(factory), + type(plugin).__qualname__, + ) + continue plugins.append(plugin) self._plugins = plugins diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index b2ee6e27..51330fed 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -829,6 +829,44 @@ def create_plugin(self, info: InvocationStartInfo): self.assertIn("returned None", "\n".join(logs.output)) self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_factory_returning_a_non_plugin_is_contained(self): + """A factory whose return is not a plugin is logged once and skipped. + + The load-time shape check establishes only that the factory has a + callable ``create_plugin``; what that call returns is knowable only + here. A value that is not a plugin fails every hook, so registering it + would log one error per hook per invocation and still provide no + telemetry. + """ + surviving = _TrackingPlugin() + + class _WrongTypeFactory: + def create_plugin(self, info: InvocationStartInfo): + return object() + + executor = PluginExecutor( + plugins=[_WrongTypeFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + output = "\n".join(logs.output) + self.assertIn("not a DurableInstrumentationPlugin", output) + self.assertIn("object", output) + # One error, not one per hook: the value never reaches the dispatch. + self.assertEqual(len(logs.output), 1) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_every_failing_factory_leaves_the_executor_usable(self): """All factories failing is not distinguishable from having no plugins.""" executor = PluginExecutor(plugins=[_ExplodingFactory()]) From 98a7044eccd28c60ca905517e3efb340c829ee33 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 12:59:01 -0700 Subject: [PATCH 10/28] style(sdk): format the factory docstring example --- .../src/aws_durable_execution_sdk_python/plugin.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 4e0b31e4..4a100167 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 @@ -480,7 +480,8 @@ def __init__(self, exporter: Exporter) -> None: def create_plugin(self, info: InvocationStartInfo) -> MyPlugin: return MyPlugin(self._exporter) - plugins=[MyPluginFactory(exporter)] + + plugins = [MyPluginFactory(exporter)] A plugin class is not a factory. ``plugins=[MyPlugin]`` used to work because calling a class constructs an instance, and it now fails at handler From c3b8946f5b0fb95f1405b0bd13e0c5d66d0452fd Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 13:20:40 -0700 Subject: [PATCH 11/28] fix(insight): drain outside the hook's lock hold A hook runs customer code while holding the plugin's lock: the input and output transforms, a result override, and __del__ on an object a displaced record carried. The lock is reentrant so that code re-entering a hook on the same thread does not self-deadlock, so a re-entrant on_invocation_end ran its drain inside the outer frame's hold. A drain waits for the export worker, and an exporter that re-enters a hook blocks that worker on the very lock the waiting thread holds, so neither side proceeds and the invocation hangs until Lambda times it out. Hook frames are now counted per thread, shared by every instance because customer code can call a hook on another execution's plugin. A drain is requested rather than performed, and the outermost frame runs it once every lock hold on the thread is released. The unnested case is unchanged. drain() called on the export worker thread is now refused and reported rather than waiting for work only that thread can do, which is the guard Java already had in refuseWaitThatWouldBlockThePump. --- .../_export_scheduler.py | 26 +++ .../plugin.py | 193 ++++++++++++------ .../tests/test_export_scheduler.py | 47 +++++ .../tests/test_plugin.py | 83 ++++++++ 4 files changed, 289 insertions(+), 60 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index f03c5da0..a3c09aa0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -194,7 +194,24 @@ def drain(self, execution: _ExportState) -> None: returns without a flush: failing to start the export worker, and an export worker that has died ``_MAX_CONSECUTIVE_WORKER_FAULTS`` times without completing any work. + + A third path returns immediately: a call made on the export worker thread + itself. Only that worker exports records and completes flushes, so a wait + there would park the one thread able to release it. That happens when an + exporter re-enters a plugin hook and the hook reaches an invocation end. + The call is refused and reported rather than deadlocking the invocation; + the record stays queued and this same worker exports it once it resumes + its loop. (Mirrors the Java + ``ExportScheduler.refuseWaitThatWouldBlockThePump``.) """ + if self._is_export_worker(): + _logger.warning( + "workflow-insight: drain() was called on the export worker " + "thread, the only thread able to serve it, so the call was " + "refused rather than deadlocking the invocation; an exporter " + "re-entered a plugin hook" + ) + return failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: @@ -260,6 +277,15 @@ def drain(self, execution: _ExportState) -> None: # -- internals ------------------------------------------------------------ + def _is_export_worker(self) -> bool: + """Report whether the calling thread is this scheduler's export worker. + + Read under the condition's lock, because ``_worker`` is replaced by the + waiter that starts a replacement and cleared by a worker that exits. + """ + with self._condition: + return self._worker is threading.current_thread() + def _disable_locked(self) -> _Dropped: """Latch asynchronous export off for good and surrender everything queued. diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 572e62d7..c0e478d6 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -42,10 +42,12 @@ from __future__ import annotations +import contextlib import datetime import json import math import threading +from collections.abc import Iterator from typing import Any, Callable from aws_durable_execution_sdk_python.plugin import ( @@ -174,6 +176,26 @@ def _apply_result_override( return None +class _HookFrames(threading.local): + """Nested plugin hook frames on one thread, and the drains they owe. + + Per thread, and shared by every plugin instance on that thread. A hook runs + customer code while holding a plugin's ``_lock``, and that code can call a + hook on *any* live instance, so the frame that must run a deferred drain is + the outermost one on the thread whatever instance it belongs to. + + ``threading.local`` runs ``__init__`` once per thread, so each thread gets its + own counter and its own list. + """ + + def __init__(self) -> None: + self.depth = 0 + self.pending: list[WorkflowInsightPlugin] = [] + + +_hook_frames = _HookFrames() + + class WorkflowInsightPlugin(DurableInstrumentationPlugin, _ExportState): """Everything this environment holds for one invocation of one execution. @@ -284,10 +306,54 @@ def _adopt_operations_locked(self, operations: dict[str, OperationInfo]) -> None # -- hooks ---------------------------------------------------------------- + @contextlib.contextmanager + def _hook_frame(self) -> Iterator[None]: + """Mark a hook frame on this thread and run the drains it owes on exit. + + A hook runs customer code -- the input/output transforms, a result + override, ``__del__`` on an object a displaced record carried -- while + holding ``_lock``, which is reentrant so that such code re-entering a hook + on this thread does not self-deadlock. A re-entrant + ``on_invocation_end`` therefore used to run its drain while the outer + frame still held ``_lock``. A drain waits for the export worker, and an + exporter that re-enters a hook blocks that worker on the very ``_lock`` + the waiting thread holds, so neither side can proceed and the invocation + hangs until Lambda times it out. + + A drain a nested frame asks for is therefore deferred to the outermost + frame, which runs it after every ``_lock`` hold on this thread has been + released. The unnested case is unchanged: the frame is the outermost one, + so its drain runs at the same point it always did. + + The frame state is per thread and shared by every instance, because + customer code inside one execution's build can call a hook on another + execution's instance, and a drain deferred to the outer frame of a + *different* instance is still a drain outside every lock. + """ + state = _hook_frames + state.depth += 1 + try: + yield + finally: + state.depth -= 1 + if state.depth == 0 and state.pending: + owed, state.pending = state.pending, [] + for plugin in owed: + plugin._drain() + + def _request_drain(self) -> None: + """Ask for a drain once the outermost hook frame on this thread unwinds.""" + pending = _hook_frames.pending + if not any(plugin is self for plugin in pending): + pending.append(self) + + def _drain(self) -> None: + self._shared._scheduler.drain(self) + def on_invocation_start(self, info: InvocationStartInfo) -> None: if not self._sampled_in: return - with self._lock: + with self._hook_frame(), self._lock: if self._closed: return # Seed the operation map from the full snapshot. On a cold resume @@ -305,7 +371,7 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: # fabricate it in, and the instance it reaches is its own. if not self._sampled_in: return - with self._lock: + with self._hook_frame(), self._lock: if self._closed: return # Replace state with the full operations snapshot carried by the hook. @@ -322,64 +388,71 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # nothing. return emit_mode = self._shared._emit_mode - with self._lock: - if not self._closed: - # Close the gate before emitting so a concurrent late hook for - # this execution cannot append a RUNNING record after the - # terminal one. - self._closed = True - # Refresh from the fresh end-of-invocation snapshot before - # emitting so the terminal record reflects the final operation - # map. - self._adopt_operations_locked(info.operations) - status = _STATUS_MAP.get(info.status, "RUNNING") - is_terminal = status in ("SUCCEEDED", "FAILED") - is_failure = status == "FAILED" - - if emit_mode == EmitMode.ON_CHANGE: - should_emit = True - elif emit_mode == EmitMode.ON_FAILURE: - should_emit = is_failure - else: # on-complete - should_emit = is_terminal - - if should_emit: - # Only terminal (SUCCEEDED/FAILED) records carry an end time; - # a PENDING/RETRY invocation end maps to RUNNING (still in - # flight) and must omit endTime/durationMs. Passing - # end_time=None makes _emit drop both fields. Output and - # error likewise belong only to a terminal record. - self._emit( - status=status, - end_time=datetime.datetime.now(datetime.UTC) - if is_terminal - else None, - output_raw=info.execution_result if is_terminal else None, - error=info.error if is_terminal else None, - # This is the emit that closed the gate, so it always runs - # with `_closed` already set and must never drop itself. - closing=True, - ) - - # Nothing has to be cleared after an invocation end, including a - # PENDING/RETRY one: this instance IS the state, and the SDK drops it - # when the invocation scope exits. A suspended execution that resumes - # here later gets a fresh instance, seeded from - # InvocationStartInfo.operations. - # - # Drain on EVERY sampled-in invocation end, emitted record or not: JS and - # Java flush once per sampled-in invocation end regardless, and a - # buffering exporter has to see the same rhythm in all three languages - # (an on-failure/on-complete mode that emits nothing for this invocation - # may still be holding records another execution handed it). A sampled-out - # execution returns above, so it neither exports nor flushes. - # - # The drain covers this execution only -- it names this instance, which - # carries its own export bookkeeping: it returns once this execution's - # own record, if any, reached the exporters and a flush that completed - # after this call is done, without waiting on records scheduled after the - # call by other executions. - self._shared._scheduler.drain(self) + with self._hook_frame(): + with self._lock: + if not self._closed: + # Close the gate before emitting so a concurrent late hook for + # this execution cannot append a RUNNING record after the + # terminal one. + self._closed = True + # Refresh from the fresh end-of-invocation snapshot before + # emitting so the terminal record reflects the final operation + # map. + self._adopt_operations_locked(info.operations) + status = _STATUS_MAP.get(info.status, "RUNNING") + is_terminal = status in ("SUCCEEDED", "FAILED") + is_failure = status == "FAILED" + + if emit_mode == EmitMode.ON_CHANGE: + should_emit = True + elif emit_mode == EmitMode.ON_FAILURE: + should_emit = is_failure + else: # on-complete + should_emit = is_terminal + + if should_emit: + # Only terminal (SUCCEEDED/FAILED) records carry an end time; + # a PENDING/RETRY invocation end maps to RUNNING (still in + # flight) and must omit endTime/durationMs. Passing + # end_time=None makes _emit drop both fields. Output and + # error likewise belong only to a terminal record. + self._emit( + status=status, + end_time=datetime.datetime.now(datetime.UTC) + if is_terminal + else None, + output_raw=info.execution_result if is_terminal else None, + error=info.error if is_terminal else None, + # This is the emit that closed the gate, so it always runs + # with `_closed` already set and must never drop itself. + closing=True, + ) + + # Nothing has to be cleared after an invocation end, including a + # PENDING/RETRY one: this instance IS the state, and the SDK drops it + # when the invocation scope exits. A suspended execution that resumes + # here later gets a fresh instance, seeded from + # InvocationStartInfo.operations. + # + # Drain on EVERY sampled-in invocation end, emitted record or not: JS and + # Java flush once per sampled-in invocation end regardless, and a + # buffering exporter has to see the same rhythm in all three languages + # (an on-failure/on-complete mode that emits nothing for this invocation + # may still be holding records another execution handed it). A sampled-out + # execution returns above, so it neither exports nor flushes. + # + # The drain covers this execution only -- it names this instance, which + # carries its own export bookkeeping: it returns once this execution's + # own record, if any, reached the exporters and a flush that completed + # after this call is done, without waiting on records scheduled after the + # call by other executions. + # + # Asked for rather than performed here, so it runs when the outermost + # hook frame on this thread unwinds and every `_lock` hold is released. + # See `_hook_frame`: an invocation end that customer code re-entered + # from inside another hook's build would otherwise wait for the export + # worker while holding the lock that worker may need. + self._request_drain() # -- emission ------------------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index e5ced2e8..e5f5c38a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -4,6 +4,7 @@ from __future__ import annotations +import logging import threading import time from typing import Any @@ -862,3 +863,49 @@ def drain(arn: str) -> None: assert exported < returned_at assert ("flush", None, None) in events[exported:returned_at] assert _wait_until(lambda: _scheduler_is_empty(scheduler)) + + +class ReentrantDrainExporter(CaptureExporter): + """Exporter that drains from inside export(), as a plugin hook would. + + An exporter that re-enters a plugin hook reaches ``on_invocation_end``, and + that hook drains. The re-entry therefore arrives on the export worker thread, + which is the one thread able to serve the wait. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.returned = threading.Event() + + def export(self, record: dict[str, Any]) -> None: + super().export(record) + assert self.scheduler is not None + self.scheduler.drain(ARN_A) + self.returned.set() + + +def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking( + caplog, +) -> None: + """A drain on the worker thread returns instead of parking it. + + Only the export worker exports records and completes flushes. A drain made on + that thread would wait for work only that thread can do, so the wait never + ends and the invocation hangs until Lambda times it out. The call is refused + and reported, and the worker goes back to its loop. + """ + exporter = ReentrantDrainExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + with caplog.at_level(logging.WARNING): + scheduler.schedule(ARN_A, _record("r1")) + + assert exporter.returned.wait(timeout=10), "the refused drain must return" + + assert ("export", "r1") in exporter.calls + assert "refused rather than deadlocking" in caplog.text + # The worker is still serving: a drain from any other thread completes. + scheduler.drain(ARN_B) + assert ("flush", None) in exporter.calls diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 0e3c9ce2..a2fe608c 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -1125,6 +1125,89 @@ def hook() -> None: assert _wait_until(lambda: not factory._scheduler._worker_alive()) +def test_a_reentrant_invocation_end_drains_after_the_lock_is_released(): + # The nested end hook used to drain from inside the outer frame's lock hold. A + # drain waits for the export worker, and an exporter that re-enters a hook + # blocks that worker on the very lock the waiting thread holds, so neither + # side can proceed and the invocation hangs until Lambda times it out. The + # drain a nested frame asks for is therefore deferred to the outermost hook + # frame, which runs it with every lock hold on this thread released. + exporter = ConcurrentCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + events: list[str] = [] + lock_free_at_drain: list[bool] = [] + + def reentering_input(value: Any) -> Any: + if not reentered.is_set(): + reentered.set() + holder["plugin"].on_invocation_end(_end(operations=_ops(_step("s")))) + events.append("nested-end-returned") + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + original_drain = factory._scheduler.drain + + def recording_drain(execution: Any) -> None: + events.append("drain") + lock_free_at_drain.append(_free_for_another_thread(plugin._lock)) + original_drain(execution) + + factory._scheduler.drain = recording_drain # type: ignore[method-assign] + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + + assert events == ["nested-end-returned", "drain"], ( + "the nested end hook drained before its frame unwound, so the drain ran " + f"inside the outer lock hold: {events}" + ) + assert lock_free_at_drain == [True], ( + "the drain ran while this thread still held the execution's lock, which " + "an exporter re-entering a hook turns into a deadlock" + ) + + +def _free_for_another_thread(lock: Any) -> bool: + """Report whether a lock is unheld, as seen from a thread that never took it. + + Asked from another thread on purpose: the lock is reentrant, so the thread + that owns it can always acquire it again and would learn nothing. + """ + acquired: list[bool] = [] + + def probe() -> None: + got = lock.acquire(blocking=False) + acquired.append(got) + if got: + lock.release() + + prober = threading.Thread(target=probe, daemon=True) + prober.start() + prober.join(5.0) + return acquired == [True] + + # -- a build overtaken by one customer code started from inside it ------------- From 3bc7d1e431de060dd47a12134fcab9c48bfb57ec Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 13:20:40 -0700 Subject: [PATCH 12/28] fix(otel): hold the log-filter claim weakly unbind_invocation resets the ContextVar claim only on the thread that ends the invocation, because a ContextVar cannot be reset from another thread. A pool thread that is never used again therefore kept a strong reference to a finished invocation's plugin, its spans and its context tokens for the life of the execution environment, and kept a plugin whose end hook never ran reachable despite _open_invocations being weak. The claim is now a weakref. A collected referent resolves to nothing, which is the answer the liveness check already gives for a finished invocation. While an invocation is open the SDK holds its plugin, which is what keeps the referent alive for every record the filter resolves. --- .../log_filter.py | 34 ++++++++++++--- .../tests/test_log_filter.py | 43 +++++++++++++++++++ 2 files changed, 70 insertions(+), 7 deletions(-) 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 87293990..0a88540a 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 @@ -30,7 +30,8 @@ thread/task for it, through a :class:`contextvars.ContextVar`. A record emitted on a claimed thread resolves to the invocation that claimed it, which is per-thread and per-task and so cannot be overwritten by a - concurrent invocation. + concurrent invocation. The claim is a weak reference, so a pool thread + that is never used again cannot pin a finished invocation's plugin. - ``unbind_invocation`` marks the invocation closed. - The claim reaches the thread running the handler body because the SDK submits that work with a copy of the invocation thread's context, taken @@ -104,9 +105,19 @@ def get_current_span_context(self) -> SpanContext | None: ... # The invocation owning the current thread/task. Set by bind_invocation on every # thread the owning plugin is given control on. -_current_invocation: contextvars.ContextVar[_SpanContextProvider | None] = ( - contextvars.ContextVar("durable_execution_otel_invocation", default=None) -) +# +# A weak reference, because a claim outlives the invocation that made it on every +# thread except the one that ends it: a ContextVar can only be reset by the +# thread that set it, so a pool thread that is never used again keeps whatever +# the claim holds. A strong claim would therefore pin a finished invocation's +# plugin, its spans and its context tokens for the life of the execution +# environment, and would also defeat _open_invocations being weak, since a +# plugin whose end hook never ran would stay reachable through the claim. While +# an invocation is open the SDK holds its plugin, which is what keeps the +# referent alive for every record the filter resolves. +_current_invocation: contextvars.ContextVar[ + weakref.ref[_SpanContextProvider] | None +] = contextvars.ContextVar("durable_execution_otel_invocation", default=None) # Serializes installation so two invocations starting at once cannot both find # a handler filterless and both add a filter to it. @@ -126,7 +137,7 @@ def bind_invocation(provider: _SpanContextProvider) -> None: """ with _registry_lock: _open_invocations.add(provider) - _current_invocation.set(provider) + _current_invocation.set(weakref.ref(provider)) def unbind_invocation(provider: _SpanContextProvider) -> None: @@ -143,7 +154,8 @@ def unbind_invocation(provider: _SpanContextProvider) -> None: """ with _registry_lock: _open_invocations.discard(provider) - if _current_invocation.get() is provider: + claim = _current_invocation.get() + if claim is not None and claim() is provider: _current_invocation.set(None) @@ -161,8 +173,16 @@ def _resolve_provider() -> _SpanContextProvider | None: invocation, and two orderings make that the wrong one: a thread still carrying a finished invocation's claim, and an invocation's own thread that has not reached its invocation-start hook yet. + + A claim whose referent has been collected resolves to nothing as well. The + claim is weak, so a plugin the SDK has released can be gone while the claim + that named it remains on a pool thread. A collected referent means the + invocation is over, which is the same answer the liveness check gives. """ - claimed = _current_invocation.get() + claim = _current_invocation.get() + if claim is None: + return None + claimed = claim() if claimed is None: return None with _registry_lock: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index 583bcdcf..770a0f53 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -3,8 +3,10 @@ from __future__ import annotations import contextvars +import gc import logging import threading +import weakref from datetime import UTC, datetime import opentelemetry.context as otel_context @@ -346,6 +348,47 @@ def emit() -> None: worker.join(timeout=10) +def test_a_claimed_worker_does_not_pin_the_finished_invocation(): + """A claim left on a pool thread does not keep the plugin alive. + + ``unbind_invocation`` resets the claim only on the thread that ends the + invocation, so a pool thread that is never used again keeps the claim it was + given. The claim is a weak reference, so what it keeps is nothing: once the + invocation ends and the SDK releases the plugin, the plugin is collectable + even while the claimed thread is still running. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix="only")) + + # A worker running in a copy of the claiming thread's context, held alive for + # the length of the test, as a pooled worker would be. + claimed_context = contextvars.copy_context() + release_worker = threading.Event() + resolved_while_open: list[bool] = [] + + def hold_the_claim() -> None: + resolved_while_open.append(log_filter_module._resolve_provider() is not None) + assert release_worker.wait(timeout=10) + + worker = threading.Thread( + target=claimed_context.run, args=(hold_the_claim,), name="claimed-worker" + ) + worker.start() + try: + collected = threading.Event() + weakref.finalize(plugin, collected.set) + plugin.on_invocation_end(_invocation_end_info(suffix="only")) + + del plugin + gc.collect() + + assert resolved_while_open == [True], "the claim must resolve while open" + assert collected.is_set(), "the claimed worker must not pin the plugin" + finally: + release_worker.set() + worker.join(timeout=10) + + def test_a_record_emitted_before_its_invocation_binds_is_not_given_the_open_one(): """A record emitted before its own invocation binds is not correlated. From 5d874f12075549520f8886a2d11b23895890623f Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 13:20:40 -0700 Subject: [PATCH 13/28] fix(sdk): reject a factory class at registration plugins=[MyFactory] -- the factory class rather than an instance of it -- passed the registration check, because create_plugin read off the class is a plain function and therefore callable. The per-invocation call supplies only the info, Python binds it to self, and the resulting TypeError is contained like any other factory failure: instrumentation is silently absent for the lifetime of the function. Registration now binds one positional argument to the member's signature, which no factory code runs. A bound method, a classmethod, a staticmethod and a __call__ on an instance all bind; an instance method read off the class does not. A callable whose signature cannot be read is accepted on the member alone, because a missing description of a factory is not evidence of a broken one. --- .../plugin_discovery.py | 78 ++++++++++++---- .../tests/plugin_discovery_test.py | 91 +++++++++++++++++++ 2 files changed, 153 insertions(+), 16 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index c0c83209..259a769a 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -1,5 +1,6 @@ from __future__ import annotations +import inspect import logging import os from collections.abc import Mapping, Sequence @@ -17,6 +18,10 @@ PLUGIN_ENTRY_POINT_GROUP = "aws_durable_execution.plugins" PLUGIN_ENVIRONMENT_VARIABLE = "DURABLE_EXECUTION_PLUGINS" +# Stands in for the InvocationStartInfo when a factory's signature is checked at +# registration. Only the bind is performed, so nothing reads it. +_ARGUMENT_PROBE = object() + def _parse_configured_plugin_names(environment: Mapping[str, str]) -> list[str]: configured_plugins = environment.get(PLUGIN_ENVIRONMENT_VARIABLE) @@ -58,21 +63,57 @@ def _is_plugin_factory(value: object) -> bool: """Report whether a value has the shape of a plugin factory. :class:`DurableInstrumentationPluginFactory` declares one method, so the - shape is one member: a callable ``create_plugin``. The attribute is fetched - and tested for callability rather than merely for presence, because an - object carrying a non-callable ``create_plugin`` would otherwise pass here - and fail at invocation time. + shape is one member: a ``create_plugin`` that can be called with one + positional argument. The attribute is fetched and tested rather than merely + checked for presence, because an object carrying a non-callable + ``create_plugin`` would otherwise pass here and fail at invocation time. + + Callability alone is not enough. ``plugins=[MyFactory]`` -- the factory + *class* rather than an instance of it -- resolves ``create_plugin`` to a + plain function whose first parameter is ``self``, which is callable. The + per-invocation call supplies only the info, Python binds it to ``self``, and + the resulting :exc:`TypeError` is contained like any other factory failure: + telemetry is silently absent for the lifetime of the function. Binding one + positional argument to the signature rejects that at registration instead. + The bind is a signature operation, so no factory code runs. + + A callable with no introspectable signature -- a C-implemented callable, for + example -- is accepted on the member alone. ``inspect.signature`` raises for + it, and refusing a factory because its signature could not be read would + reject a usable factory over a missing description of it. Structural rather than nominal, so a factory need not import the SDK protocol to satisfy it. The protocol is deliberately not ``@runtime_checkable``; see its docstring. - The check cannot go further than one member. Whether ``create_plugin`` - accepts the info, and whether it returns a plugin, is only knowable by - calling it, and calling it at load time is what the per-invocation factory - design avoids: there is no invocation yet. + The check stops there. Whether ``create_plugin`` returns a plugin is only + knowable by calling it, and calling it at load time is what the + per-invocation factory design avoids: there is no invocation yet. That case + is checked per invocation by :meth:`PluginExecutor._create_plugins`. """ - return callable(getattr(value, "create_plugin", None)) + create_plugin = getattr(value, "create_plugin", None) + if not callable(create_plugin): + return False + return _accepts_one_positional_argument(create_plugin) + + +def _accepts_one_positional_argument(create_plugin: object) -> bool: + """Report whether one positional argument can be bound to a callable. + + A bound method, a ``@classmethod`` or ``@staticmethod`` read off a class, and + a ``__call__`` on an instance all present the signature the SDK calls, so all + three bind. An instance method read off the class does not: its first + parameter is ``self``, so one argument leaves the info unbound. + """ + try: + signature = inspect.signature(create_plugin) # type: ignore[arg-type] + except (TypeError, ValueError): + return True + try: + signature.bind(_ARGUMENT_PROBE) + except TypeError: + return False + return True def _load_factory( @@ -103,7 +144,7 @@ def _load_factory( "create_plugin(info) method returning a " "DurableInstrumentationPlugin -- but resolved to " f"{_qualified_type_name(factory)}. Name the factory instance, not a " - "plugin and not a plugin class." + "plugin, not a plugin class, and not the factory class." ) return cast(DurableInstrumentationPluginFactory, factory) @@ -126,11 +167,15 @@ def _validate_explicit_factories( A plugin *class* is rejected, and so is any bare callable. Both were accepted while the registration type was ``Callable``: a lambda satisfied it directly, and a class satisfied it because calling a class constructs an - instance. Neither carries ``create_plugin``, so ``plugins=[MyPlugin]`` and - ``plugins=[lambda info: MyPlugin()]`` now fail here. The replacement is a - small factory class, which is also where setup work that can fail belongs. A - class that declares ``create_plugin`` as a ``@classmethod`` is accepted, - because the requirement is the member and not the kind of object. + instance. Neither carries a ``create_plugin`` the SDK can call, so + ``plugins=[MyPlugin]`` and ``plugins=[lambda info: MyPlugin()]`` now fail + here. A *factory* class passed instead of an instance of it fails here too: + ``MyFactory.create_plugin`` is callable, but its first parameter is ``self``, + so the per-invocation call binds the info to ``self``. The replacement is a + small factory class, instantiated, which is also where setup work that can + fail belongs. A class that declares ``create_plugin`` as a ``@classmethod`` + or a ``@staticmethod`` is accepted, because that member presents the + signature the SDK calls. """ factories = list(explicit_plugins or []) for index, factory in enumerate(factories): @@ -140,7 +185,8 @@ class that declares ``create_plugin`` as a ``@classmethod`` is accepted, "plugin factory -- an object with a create_plugin(info) method " "returning a DurableInstrumentationPlugin -- but is " f"{_qualified_type_name(factory)}. Pass a factory rather than a " - "plugin, a plugin class, or a plain callable, for example " + "plugin, a plugin class, or a plain callable, and pass a factory " + "instance rather than the factory class, for example " "plugins=[MyPluginFactory(exporter)]." ) return factories diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index 54b0f5be..b79498fb 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -1,7 +1,9 @@ from __future__ import annotations +import inspect import logging import os +import time from unittest.mock import Mock, patch import pytest @@ -577,6 +579,95 @@ def create_plugin(cls, info: InvocationStartInfo) -> _ClassFactoryPlugin: assert plugin.info is INVOCATION_START_INFO +def test_explicit_factory_class_with_an_instance_method_is_rejected() -> None: + """The factory class is not the factory, and callability does not reveal it. + + ``MyFactory.create_plugin`` read off the class is a plain function whose + first parameter is ``self``, so it is callable and used to pass. The + per-invocation call supplies only the info, Python binds it to ``self``, and + the resulting ``TypeError`` is contained like any other factory failure: + instrumentation is silently absent for the lifetime of the function. The + signature is bound at registration so the mistake fails here instead. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_PluginAFactory], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "a factory instance rather than the factory class" in str(error.value) + + +def test_discovery_rejects_a_factory_class_at_the_entry_point() -> None: + """The entry-point path applies the same signature check.""" + entry_point = _FakeEntryPoint("a", _PluginAFactory) + + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ), + pytest.raises(PluginLoadError) as error, + ): + load_configured_plugins( + None, + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert "must resolve to a plugin factory" in str(error.value) + assert "not the factory class" in str(error.value) + + +def test_explicit_class_declaring_a_static_create_plugin_is_accepted() -> None: + """A ``@staticmethod`` presents the signature the SDK calls, so it binds.""" + + class _StaticFactoryPlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + @staticmethod + def create_plugin(info: InvocationStartInfo) -> _StaticFactoryPlugin: + return _StaticFactoryPlugin(info) + + result = load_configured_plugins([_StaticFactoryPlugin], environment={}) + + assert result == [_StaticFactoryPlugin] + assert isinstance( + result[0].create_plugin(INVOCATION_START_INFO), _StaticFactoryPlugin + ) + + +def test_explicit_factory_without_an_introspectable_signature_is_accepted() -> None: + """A signature that cannot be read is not evidence of a broken factory. + + ``inspect.signature`` raises ``ValueError`` for some C-implemented callables, + ``time.strftime`` among them. Rejecting such a factory would refuse a usable + one over a missing description of it, so the member alone decides. + """ + + class _UnreadableSignatureFactory: + create_plugin = staticmethod(time.strftime) + + factory = _UnreadableSignatureFactory() + + with pytest.raises(ValueError, match="no signature"): + inspect.signature(time.strftime) + + assert load_configured_plugins([factory], environment={}) == [factory] # type: ignore[list-item, comparison-overlap] + + +def test_explicit_factory_taking_no_argument_is_rejected() -> None: + """A ``create_plugin`` that takes nothing cannot receive the info.""" + + class _NoArgumentFactory: + def create_plugin(self) -> _PluginA: + return _PluginA() + + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_NoArgumentFactory()], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "create_plugin(info) method" in str(error.value) + + def test_explicit_factory_object_is_accepted() -> None: result = load_configured_plugins([_plugin_a_factory], environment={}) From 9b8834343f756d14c5d6653c6e6d5876c77fedd8 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 14:15:28 -0700 Subject: [PATCH 14/28] fix(plugin): fire the end hook on every exit The invocation wrapper fired on_invocation_end on the success path and on except Exception, so an exit by BaseException skipped it. A handler that surfaces an asyncio.CancelledError -- user code awaiting a cancelled task -- leaves that way, and so does a KeyboardInterrupt or SystemExit. The invocation then ended without the one hook a plugin has to finish on: Insight never drained, so the records it held for that execution were dropped, and OTel never ended the spans it had opened. Teardown still ran, so nothing failed visibly. Every exit now fires the hook and re-raises the exception unchanged. Plugin-code containment widens with it. The factory and hook boundaries caught Exception, so a factory or hook raising CancelledError failed an execution it was only observing and stopped the remaining plugins from running. Both now contain every BaseException except KeyboardInterrupt, SystemExit and GeneratorExit, which are instructions to the calling thread rather than reports of a plugin defect. CancelledError is not among them: nothing cancels the invocation thread or the single-worker plugin pool, so one arriving from plugin code came from the plugin's own asyncio use. ErrorObject.from_exception is annotated BaseException, which is what its own helper already accepted. --- .../lambda_service.py | 2 +- .../plugin.py | 59 ++++++++- .../tests/plugin_test.py | 123 ++++++++++++++++++ 3 files changed, 179 insertions(+), 5 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index eb5ee78b..797ae255 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -235,7 +235,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ErrorObject: ) @classmethod - def from_exception(cls, exception: Exception) -> ErrorObject: + def from_exception(cls, exception: BaseException) -> ErrorObject: # SerDesError and subclasses pin to the base discriminator so replay # always reconstructs them as SerDesError. if isinstance(exception, SerDesError): 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 4a100167..f9eb8240 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 @@ -530,6 +530,21 @@ def _factory_name(factory: object) -> str: return getattr(factory, "__qualname__", None) or type(factory).__name__ +# Raised out of plugin code, these three are not reports of a plugin defect but +# instructions to the thread that is running: stop. Containing one would drop the +# instruction and return a thread that was told to unwind to the work after the +# plugin. They are re-raised; every other BaseException is contained. +# +# asyncio.CancelledError is deliberately NOT here. It derives from BaseException +# and it does mean "stop" for the task that was cancelled, but the task here is +# the SDK's, not the plugin's: nothing cancels the invocation thread or the +# single-worker plugin pool. A CancelledError arriving from plugin code therefore +# came from the plugin's own asyncio use -- an awaited task it let be cancelled -- +# which is a plugin defect and belongs on the contained side, or the plugin's +# failure would fail an execution it was only observing. +_PLUGIN_THREAD_CONTROL_EXCEPTIONS = (KeyboardInterrupt, SystemExit, GeneratorExit) + + class PluginExecutor: """One invocation's plugin instances, metadata and dispatch. @@ -611,12 +626,21 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: :func:`plugin_discovery.load_configured_plugins` rejects such an entry while the handler is being initialized, so the silent case is not reachable through ``durable_execution()``. + + Containment covers every ``BaseException`` except the three that instruct + the calling thread to stop; see + :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS`. Narrowing it to ``Exception`` + left the contract conditional on a factory never raising outside that + hierarchy, and a factory that awaits a cancelled task raises + ``asyncio.CancelledError``, which is outside it. """ plugins: list[DurableInstrumentationPlugin] = [] for factory in self._plugin_factories: try: plugin = factory.create_plugin(info) - except Exception: + except _PLUGIN_THREAD_CONTROL_EXCEPTIONS: + raise + except BaseException: # noqa: BLE001 - a factory must not fail the execution # log and ignore the exception logger.exception( "Plugin factory %s exception ignored", _factory_name(factory) @@ -647,7 +671,14 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: @staticmethod def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: - """Invoke the appropriate plugin callback. Runs inside the thread pool.""" + """Invoke the appropriate plugin callback. Runs inside the thread pool. + + Contains every ``BaseException`` except the three that instruct the + calling thread to stop, the same rule the factory boundary uses. The + thread here is the executor's own single worker, which nothing outside + this class cancels or interrupts, so an exception outside the ``Exception`` + hierarchy arriving here was raised by the plugin. + """ try: match info: case InvocationStartInfo(): @@ -666,7 +697,9 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: plugin.on_user_function_end(info) case _: raise RuntimeError(f"Unknown info type: {type(info)}") - except Exception: + except _PLUGIN_THREAD_CONTROL_EXCEPTIONS: + raise + except BaseException: # noqa: BLE001 - a hook must not fail the execution # log and ignore the exception logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) @@ -1068,7 +1101,25 @@ def wrapper(event: Any, context: LambdaContext): output=DurableExecutionInvocationOutput.from_dict(output), ) return output - except Exception as e: + except BaseException as e: + # Every exit fires the end hook, not only the ones that + # derive from Exception. A handler that surfaces an + # asyncio.CancelledError -- user code that awaited a + # cancelled task, most simply -- leaves the invocation by + # a BaseException, and an invocation that ends without + # its end hook costs the plugins the only point at which + # they can finish: Insight never drains, so the records it + # holds for this execution are dropped, and OTel never + # ends the spans it opened, so they are never exported. + # The teardown below still runs either way, which is why + # the gap was silent rather than a leak. + # + # KeyboardInterrupt and SystemExit reach here too, and + # they also fire the hook. The hook is what a plugin needs + # to flush, and a process being torn down is when flushing + # matters; the cost is the same bounded work any + # invocation end does. The exception itself is re-raised + # unchanged, so what the caller sees is untouched. plugin_executor.on_invocation_end( output=DurableExecutionInvocationOutput.create_retry( ErrorObject.from_exception(e) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 51330fed..efafb6fc 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,3 +1,4 @@ +import asyncio import contextlib import datetime import logging @@ -646,6 +647,41 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_a_body_raising_cancellation_still_fires_the_end_hook(self): + """Every exit fires the end hook, not only the ones deriving from Exception. + + A handler that surfaces an ``asyncio.CancelledError`` left the invocation + without its end hook, so Insight never drained the records it held for the + execution and OTel never ended the spans it had opened. Nothing failed + visibly, which is why the gap was silent. + """ + for raised in ( + asyncio.CancelledError("cancelled"), + KeyboardInterrupt(), + SystemExit(), + ): + with self.subTest(raised=type(raised).__name__): + plugin = _TrackingPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise raised + + with self.assertRaises(type(raised)): + handler({}, LAMBDA_CTX) + + self.assertEqual( + plugin.calls, + ["invocation_start:req-1", "invocation_end:req-1"], + ) + def test_host_hands_out_a_new_executor_per_invocation(self): """The host itself holds no per-invocation state to overwrite.""" host = PluginHost(plugins=[plugin_factory(_TrackingPlugin())]) @@ -891,6 +927,69 @@ def test_every_failing_factory_leaves_the_executor_usable(self): ), ) + def test_a_factory_raising_cancellation_is_contained(self): + """Containment is not limited to ``Exception``. + + ``asyncio.CancelledError`` derives from ``BaseException``, so a factory + that awaits a cancelled task used to abort the invocation it was only + instrumenting and stop the remaining factories from running. + """ + surviving = _TrackingPlugin() + + executor = PluginExecutor( + plugins=[_CancellingFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("factory cancelled", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_a_factory_raising_thread_control_still_propagates(self): + """The three that tell the thread to stop are not contained. + + Containing one would drop the instruction and hand the thread back to the + work that follows the plugin. + """ + for control in (KeyboardInterrupt, SystemExit, GeneratorExit): + with self.subTest(control=control.__name__): + executor = PluginExecutor( + plugins=[_ThreadControlFactory(control)], + ) + + with executor.run(), self.assertRaises(control): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + def test_a_hook_raising_cancellation_is_contained(self): + """The hook boundary uses the same rule as the factory boundary.""" + tracking = _TrackingPlugin() + executor = PluginExecutor( + plugins=[plugin_factory(_CancellingPlugin()), plugin_factory(tracking)] + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with _invocation(executor, tracking): + executor.execute_plugins(OPERATION_START_INFO, sync=True) + + self.assertIn("hook cancelled", "\n".join(logs.output)) + self.assertIn("operation_start:op-2", tracking.calls) + class TestPluginExecutor(unittest.TestCase): def test_no_thread_pool_when_plugins_is_none(self): @@ -2171,6 +2270,30 @@ def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlug raise RuntimeError("factory boom") +class _CancellingFactory: + """Factory whose ``create_plugin`` raises outside the ``Exception`` hierarchy.""" + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise asyncio.CancelledError("factory cancelled") + + +class _ThreadControlFactory: + """Factory that raises one of the three exceptions that must propagate.""" + + def __init__(self, control: type[BaseException]) -> None: + self._control = control + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise self._control + + +class _CancellingPlugin(DurableInstrumentationPlugin): + """Plugin whose hook raises outside the ``Exception`` hierarchy.""" + + def on_operation_start(self, info): + raise asyncio.CancelledError("hook cancelled") + + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" From ac786f60dca1081bb5e6b4ca845921221cfa68fc Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 15:12:22 -0700 Subject: [PATCH 15/28] fix(plugin): send one end notification per invocation The success dispatch of on_invocation_end sat inside the try that reports a handler failure, so an end hook that raised was caught as though the handler had failed and the hook ran a second time with a RETRY outcome. Later plugins were told the wrong thing about an invocation that succeeded, and an exporter exported twice. A hook can raise: _dispatch_plugin re-raises the three exceptions that instruct the calling thread to stop, and from_dict can reject an output the handler built. The output is parsed inside the try and the success hook is dispatched after it, so each invocation produces exactly one end notification. Registration also rejects more factory-class shapes. A signature bind accepts create_plugin(self, info=None) and create_plugin(self, *args), because the probe binds to self and what remains is satisfied, so plugins=[MyFactory] still passed and then failed on every invocation where the error is swallowed. For a class the kind of the member now decides: a classmethod carries __self__, a staticmethod is identified through its descriptor, and an attribute holding a callable object takes no implicit first argument. Anything else read off a class takes self and is rejected. inspect.getattr_static walks the MRO without running a descriptor, so this still runs no factory code. --- .../plugin.py | 17 ++-- .../plugin_discovery.py | 45 +++++++++- .../tests/plugin_discovery_test.py | 85 ++++++++++++++++--- .../tests/plugin_test.py | 41 +++++++++ 4 files changed, 167 insertions(+), 21 deletions(-) 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 f9eb8240..e4eebc7d 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 @@ -1094,13 +1094,18 @@ def decorator( @functools.wraps(func) def wrapper(event: Any, context: LambdaContext): with self.invocation() as plugin_executor: + # The end hook is dispatched exactly once per invocation, so + # the success dispatch sits outside the try. Inside it, an + # end hook that raised would be caught as though the handler + # had failed, and the hook would run a second time with a + # RETRY outcome -- telling later plugins the wrong thing about + # an invocation that succeeded, and letting an exporter export + # twice. A hook can raise: _dispatch_plugin re-raises the three + # exceptions that instruct the calling thread to stop, and + # from_dict below can reject an output the handler built. try: output = func(event, context, plugin_executor) - - plugin_executor.on_invocation_end( - output=DurableExecutionInvocationOutput.from_dict(output), - ) - return output + completed = DurableExecutionInvocationOutput.from_dict(output) except BaseException as e: # Every exit fires the end hook, not only the ones that # derive from Exception. A handler that surfaces an @@ -1126,6 +1131,8 @@ def wrapper(event: Any, context: LambdaContext): ), ) raise + plugin_executor.on_invocation_end(output=completed) + return output return wrapper diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index 259a769a..9ff7290d 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -73,9 +73,12 @@ def _is_plugin_factory(value: object) -> bool: plain function whose first parameter is ``self``, which is callable. The per-invocation call supplies only the info, Python binds it to ``self``, and the resulting :exc:`TypeError` is contained like any other factory failure: - telemetry is silently absent for the lifetime of the function. Binding one - positional argument to the signature rejects that at registration instead. - The bind is a signature operation, so no factory code runs. + telemetry is silently absent for the lifetime of the function. Two checks + reject that at registration instead. For a class, the *kind* of the member + decides, because a signature bind cannot tell ``create_plugin(self, info)`` + from ``create_plugin(info)``: see :func:`_is_unbound_instance_method`. For + everything else, binding one positional argument to the signature rejects a + member that cannot receive the info. Neither check runs factory code. A callable with no introspectable signature -- a C-implemented callable, for example -- is accepted on the member alone. ``inspect.signature`` raises for @@ -94,9 +97,45 @@ def _is_plugin_factory(value: object) -> bool: create_plugin = getattr(value, "create_plugin", None) if not callable(create_plugin): return False + if isinstance(value, type) and _is_unbound_instance_method(value, create_plugin): + return False return _accepts_one_positional_argument(create_plugin) +def _is_unbound_instance_method(cls: type, create_plugin: object) -> bool: + """Report whether a class's ``create_plugin`` is an instance method. + + Read off the class, an instance method is a plain function whose first + parameter is ``self``, so the per-invocation call binds the info to ``self`` + and the factory never sees it. The signature bind below cannot catch every + such shape: ``create_plugin(self, info=None)`` and + ``create_plugin(self, *args)`` both bind one argument to ``self`` and leave + the rest satisfied. The kind of the member decides it instead. + + Three shapes read off a class are usable and none of them is a plain + function. A ``@classmethod`` is already bound to the class, so it carries + ``__self__``. A ``@staticmethod`` is a plain function, but its descriptor says + it takes no implicit first argument. And an attribute holding a callable + object -- ``create_plugin = SomeCallable()`` -- is not a function at all and + takes no implicit first argument either. Anything else read off a class takes + ``self`` and cannot serve. + + :func:`inspect.getattr_static` is what distinguishes the ``@staticmethod``, + because it returns the descriptor rather than what reading the attribute + produces. It walks the MRO without running any descriptor, so no factory code + runs here. + """ + if getattr(create_plugin, "__self__", None) is not None: + return False + if not inspect.isfunction(create_plugin): + return False + try: + declared = inspect.getattr_static(cls, "create_plugin") + except AttributeError: + return False + return not isinstance(declared, staticmethod) + + def _accepts_one_positional_argument(create_plugin: object) -> bool: """Report whether one positional argument can be bound to a callable. diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index b79498fb..2aac4ba5 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -45,6 +45,20 @@ def create_plugin(self, info: InvocationStartInfo) -> _PluginB: return _PluginB() +class _DefaultedArgumentFactory: + """Instance method whose info parameter has a default, so one argument binds.""" + + def create_plugin(self, info: InvocationStartInfo | None = None) -> _PluginA: + return _PluginA() + + +class _VariadicFactory: + """Instance method taking ``*args``, so any argument count binds.""" + + def create_plugin(self, *args: object) -> _PluginA: + return _PluginA() + + _plugin_a_factory = _PluginAFactory() _plugin_b_factory = _PluginBFactory() @@ -579,26 +593,49 @@ def create_plugin(cls, info: InvocationStartInfo) -> _ClassFactoryPlugin: assert plugin.info is INVOCATION_START_INFO -def test_explicit_factory_class_with_an_instance_method_is_rejected() -> None: - """The factory class is not the factory, and callability does not reveal it. - - ``MyFactory.create_plugin`` read off the class is a plain function whose - first parameter is ``self``, so it is callable and used to pass. The - per-invocation call supplies only the info, Python binds it to ``self``, and - the resulting ``TypeError`` is contained like any other factory failure: - instrumentation is silently absent for the lifetime of the function. The - signature is bound at registration so the mistake fails here instead. +@pytest.mark.parametrize( + "factory_class", + [ + _PluginAFactory, + _DefaultedArgumentFactory, + _VariadicFactory, + ], + ids=["plain", "defaulted", "variadic"], +) +def test_explicit_factory_class_with_an_instance_method_is_rejected( + factory_class: type, +) -> None: + """The factory class is not the factory, and no signature shape rescues it. + + ``MyFactory.create_plugin`` read off the class is a plain function whose first + parameter is ``self``, so the per-invocation call binds the info to ``self`` + and the factory never sees it. A signature bind alone does not catch every + such shape: ``create_plugin(self, info=None)`` and + ``create_plugin(self, *args)`` both bind one argument to ``self`` and leave + the rest satisfied, so they used to pass and then fail on every invocation + where the error is swallowed. The kind of the member decides instead. """ with pytest.raises(PluginLoadError) as error: - load_configured_plugins([_PluginAFactory], environment={}) # type: ignore[list-item] + load_configured_plugins([factory_class], environment={}) # type: ignore[list-item] assert "plugins[0]" in str(error.value) assert "a factory instance rather than the factory class" in str(error.value) -def test_discovery_rejects_a_factory_class_at_the_entry_point() -> None: - """The entry-point path applies the same signature check.""" - entry_point = _FakeEntryPoint("a", _PluginAFactory) +@pytest.mark.parametrize( + "factory_class", + [ + _PluginAFactory, + _DefaultedArgumentFactory, + _VariadicFactory, + ], + ids=["plain", "defaulted", "variadic"], +) +def test_discovery_rejects_a_factory_class_at_the_entry_point( + factory_class: type, +) -> None: + """The entry-point path applies the same rule.""" + entry_point = _FakeEntryPoint("a", factory_class) with ( patch( @@ -616,6 +653,28 @@ def test_discovery_rejects_a_factory_class_at_the_entry_point() -> None: assert "not the factory class" in str(error.value) +def test_explicit_class_holding_a_callable_create_plugin_is_accepted() -> None: + """A class attribute holding a callable takes no implicit first argument. + + Reading it off the class produces the callable itself, so the info reaches it. + """ + + class _CallableMember: + def __call__(self, info: InvocationStartInfo) -> _PluginA: + return _PluginA() + + class _MemberFactory: + create_plugin = _CallableMember() + + result = load_configured_plugins([_MemberFactory], environment={}) # type: ignore[list-item] + + assert result == [_MemberFactory] + assert isinstance( + _MemberFactory.create_plugin(INVOCATION_START_INFO), + _PluginA, + ) + + def test_explicit_class_declaring_a_static_create_plugin_is_accepted() -> None: """A ``@staticmethod`` presents the signature the SDK calls, so it binds.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index efafb6fc..604977d7 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -647,6 +647,36 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_an_end_hook_that_stops_the_thread_is_not_reported_twice(self): + """Exactly one end notification per invocation, whatever the hook does. + + ``_dispatch_plugin`` re-raises the three exceptions that instruct the + calling thread to stop, so a hook can raise. With the success dispatch + inside the try, that raise was caught as a handler failure and the hook + ran a second time with a RETRY outcome -- the wrong outcome for an + invocation that succeeded, and a second export for an exporter. + """ + plugin = _ControlOnEndPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + return { + "Status": ServiceInvocationStatus.SUCCEEDED.value, + "Result": None, + } + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual(plugin.end_statuses, [InvocationStatus.SUCCEEDED]) + def test_a_body_raising_cancellation_still_fires_the_end_hook(self): """Every exit fires the end hook, not only the ones deriving from Exception. @@ -2294,6 +2324,17 @@ def on_operation_start(self, info): raise asyncio.CancelledError("hook cancelled") +class _ControlOnEndPlugin(DurableInstrumentationPlugin): + """Plugin whose end hook records the outcome and then stops the thread.""" + + def __init__(self) -> None: + self.end_statuses: list[InvocationStatus] = [] + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + self.end_statuses.append(info.status) + raise KeyboardInterrupt + + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" From 3cbb20256b68cb1d5b31e29cc2c003b130fdb0a3 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 15:51:35 -0700 Subject: [PATCH 16/28] fix(plugin): pair the end hook with the start hook A start hook raising one of the three exceptions that instruct the calling thread to stop propagates out of the dispatch loop, so plugins later in the list never receive their start hook. The invocation-end hook that the propagating exception then triggers reached them anyway, leaving a plugin to tear down state it had never been told to build. Hooks after the start hook are now dispatched to the plugins that received the start hook, which makes the pairing an invariant rather than a coincidence. A plugin counts as started before its hook is dispatched, because a hook that begins and then fails may already have allocated what its end hook releases. A refused drain also requests a flush now. An exporter that re-enters a plugin hook can queue a record from the export worker, and drain() refuses the wait that hook makes. The worker exits its loop once nothing is pending and no flush is requested, so the record reached the exporters and the worker stopped, leaving a buffering exporter holding an execution's terminal telemetry when Lambda froze the environment. The request is made without waiting, so the worker stays unblocked, and its barrier covers the record just queued. --- .../_export_scheduler.py | 26 ++++++++++++-- .../tests/test_export_scheduler.py | 10 ++++-- .../plugin.py | 25 ++++++++++++- .../tests/plugin_test.py | 36 +++++++++++++++++++ 4 files changed, 92 insertions(+), 5 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index a3c09aa0..40e42941 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -201,8 +201,8 @@ def drain(self, execution: _ExportState) -> None: exporter re-enters a plugin hook and the hook reaches an invocation end. The call is refused and reported rather than deadlocking the invocation; the record stays queued and this same worker exports it once it resumes - its loop. (Mirrors the Java - ``ExportScheduler.refuseWaitThatWouldBlockThePump``.) + its loop, and a flush covering it is requested on the way out. (Mirrors + the Java ``ExportScheduler.refuseWaitThatWouldBlockThePump``.) """ if self._is_export_worker(): _logger.warning( @@ -211,6 +211,14 @@ def drain(self, execution: _ExportState) -> None: "refused rather than deadlocking the invocation; an exporter " "re-entered a plugin hook" ) + # The refused call still leaves a flush behind. The hook that + # re-entered may have queued a record, and this worker exits its loop + # once nothing is pending and no flush is requested -- so without the + # request the record would be handed to the exporters and the worker + # would stop, leaving a buffering exporter holding an execution's + # terminal telemetry when Lambda freezes the environment. Requesting + # it rather than waiting for it is what keeps the worker unblocked. + self._request_flush() return failed_pending: _Dropped | None = None start_error: Exception | None = None @@ -286,6 +294,20 @@ def _is_export_worker(self) -> bool: with self._condition: return self._worker is threading.current_thread() + def _request_flush(self) -> None: + """Ask the worker for a flush covering everything scheduled so far. + + Returns without waiting, so it is safe to call from the worker itself. The + barrier is raised to the current schedule counter, which is what makes the + flush cover a record queued moments ago rather than running before it. + """ + with self._condition: + if self._disabled: + return + self._flush_requested = True + self._flush_barrier = max(self._flush_barrier, self._seq) + self._condition.notify_all() + def _disable_locked(self) -> _Dropped: """Latch asynchronous export off for good and surrender everything queued. diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index e5f5c38a..a10fa448 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -903,9 +903,15 @@ def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking( scheduler.schedule(ARN_A, _record("r1")) assert exporter.returned.wait(timeout=10), "the refused drain must return" + assert _wait_until(lambda: "refused rather than deadlocking" in caplog.text) assert ("export", "r1") in exporter.calls - assert "refused rather than deadlocking" in caplog.text + # The refused drain still leaves a flush behind, with no external drain to ask + # for one. The re-entering hook may have queued a record, and the worker exits + # once nothing is pending and no flush is requested, so without the request a + # buffering exporter would be holding that record when the environment froze. + assert _wait_until(lambda: ("flush", None) in exporter.calls), ( + "a refused drain must request a flush on its way out" + ) # The worker is still serving: a drain from any other thread completes. scheduler.drain(ARN_B) - assert ("flush", None) in exporter.calls 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 e4eebc7d..336f4403 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 @@ -569,6 +569,10 @@ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): # on_invocation_start and emptied when the invocation scope exits. self._plugin_factories = list(plugins or []) self._plugins: list[DurableInstrumentationPlugin] = [] + # The subset of _plugins whose invocation-start hook has been dispatched. + # Every later hook is dispatched to this list, so a plugin that never + # received its start hook never receives its end hook. + self._started: list[DurableInstrumentationPlugin] = [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None @@ -611,6 +615,7 @@ class exists to prevent; failing loudly here keeps the bug from # drained, so no queued dispatch still holds one: nothing outlives # the invocation. self._plugins = [] + self._started = [] def _create_plugins(self, info: InvocationStartInfo) -> None: """Build this invocation's plugin instances from its start info. @@ -704,9 +709,27 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) def execute_plugins(self, info, sync): + """Dispatch one hook to this invocation's plugins. + + A plugin receives a hook only once it has received the invocation-start + hook, which makes the pairing an invariant rather than a coincidence. + Without it one dispatch order breaks the pairing: a start hook that raises + one of the three exceptions :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS` + names propagates out of this loop, so plugins later in the list never + receive their start hook -- and the invocation-end hook that the + propagating exception then triggers used to reach them anyway, leaving a + plugin to tear down state it had never been told to build. + + A plugin is counted as started before its start hook is dispatched rather + than after, because a hook that begins and then fails may already have + allocated what its end hook releases. + """ if not self._executor: return - for plugin in self._plugins: + starting = isinstance(info, InvocationStartInfo) + for plugin in self._plugins if starting else self._started: + if starting: + self._started.append(plugin) if sync: # this is called synchronously, so plugins will be able to manipulate thread local objects self._dispatch_plugin(plugin, info) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 604977d7..4d4212ce 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -647,6 +647,35 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_a_start_hook_that_stops_the_thread_leaves_later_plugins_unpaired(self): + """A plugin that never received the start hook never receives the end hook. + + A start hook raising one of the three control exceptions propagates out of + the dispatch loop, so plugins after it in the list never receive their + start hook. The invocation-end hook that the propagating exception then + triggers used to reach them anyway, leaving a plugin to tear down state it + had never been told to build. + """ + later = _TrackingPlugin() + host = PluginHost( + plugins=[plugin_factory(_ControlOnStartPlugin()), plugin_factory(later)] + ) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise AssertionError("unreachable: the start hook stops the thread") + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual(later.calls, []) + def test_an_end_hook_that_stops_the_thread_is_not_reported_twice(self): """Exactly one end notification per invocation, whatever the hook does. @@ -2335,6 +2364,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: raise KeyboardInterrupt +class _ControlOnStartPlugin(DurableInstrumentationPlugin): + """Plugin whose start hook stops the thread, so later plugins never start.""" + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + raise KeyboardInterrupt + + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" From 84159aa222a10087bba29cb0f627875921d057de Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 16:12:39 -0700 Subject: [PATCH 17/28] fix(plugin): finish the end dispatch, flush only new work Two follow-ups to the previous commit, both from its own review. An end hook raising one of the three control exceptions aborted the dispatch loop, so plugins after it lost the hook that is their only chance to finish: Insight would not drain and OTel would leave spans unended. Every plugin the loop reaches has already started, so the first such exception is now held and re-raised once every plugin has been called. No other hook defers, because stopping a start-hook loop early leaves later plugins with nothing to clean up: the pairing rule withholds their end hook too. The flush a refused drain requests is now conditional on a pending record. An exporter whose flush() re-enters a plugin hook arrives at the refused drain from inside a flush, and an unconditional request asked for the next one, which re-entered and requested again for as long as the environment lived. A pending record is what distinguishes the record the re-entering hook queued from that loop. --- .../_export_scheduler.py | 36 ++++++--- .../tests/test_export_scheduler.py | 73 ++++++++++++++++--- .../plugin.py | 29 +++++++- .../tests/plugin_test.py | 34 +++++++++ 4 files changed, 147 insertions(+), 25 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 40e42941..dd1a52a8 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -211,14 +211,20 @@ def drain(self, execution: _ExportState) -> None: "refused rather than deadlocking the invocation; an exporter " "re-entered a plugin hook" ) - # The refused call still leaves a flush behind. The hook that - # re-entered may have queued a record, and this worker exits its loop - # once nothing is pending and no flush is requested -- so without the - # request the record would be handed to the exporters and the worker - # would stop, leaving a buffering exporter holding an execution's - # terminal telemetry when Lambda freezes the environment. Requesting - # it rather than waiting for it is what keeps the worker unblocked. - self._request_flush() + # The refused call still leaves a flush behind, but only when there is + # something for it to cover. The hook that re-entered may have queued a + # record, and this worker exits its loop once nothing is pending and no + # flush is requested -- so without the request the record would be + # handed to the exporters and the worker would stop, leaving a + # buffering exporter holding an execution's terminal telemetry when + # Lambda freezes the environment. + # + # Requesting one unconditionally would livelock instead: an exporter + # whose flush() re-enters a hook arrives here from inside a flush, and + # an unconditional request would ask for the next one, which re-enters + # again, for as long as the environment lives. A pending record is what + # distinguishes new work from that loop. + self._request_flush_for_pending_records() return failed_pending: _Dropped | None = None start_error: Exception | None = None @@ -294,15 +300,23 @@ def _is_export_worker(self) -> bool: with self._condition: return self._worker is threading.current_thread() - def _request_flush(self) -> None: - """Ask the worker for a flush covering everything scheduled so far. + def _request_flush_for_pending_records(self) -> None: + """Ask the worker for a flush, but only if a record is waiting for one. Returns without waiting, so it is safe to call from the worker itself. The barrier is raised to the current schedule counter, which is what makes the flush cover a record queued moments ago rather than running before it. + + A pending record is the condition, not a formality. This is called from a + drain refused on the worker thread, and one way to reach that is an + exporter whose ``flush()`` re-enters a plugin hook: the call then arrives + from inside a flush, and requesting the next one unconditionally would + produce a flush that re-enters, requests, and flushes again for as long as + the environment lives -- after the invocation has returned. Nothing is + pending in that case, so nothing is requested. """ with self._condition: - if self._disabled: + if self._disabled or not self._pending: return self._flush_requested = True self._flush_barrier = max(self._flush_barrier, self._seq) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index a10fa448..a09a236f 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -866,22 +866,27 @@ def drain(arn: str) -> None: class ReentrantDrainExporter(CaptureExporter): - """Exporter that drains from inside export(), as a plugin hook would. + """Exporter that queues a record and drains from inside export(). - An exporter that re-enters a plugin hook reaches ``on_invocation_end``, and - that hook drains. The re-entry therefore arrives on the export worker thread, - which is the one thread able to serve the wait. + This is what an exporter re-entering a plugin hook produces: the hook emits + its record and then, at an invocation end, drains. Both arrive on the export + worker, which is the one thread able to serve the wait. """ def __init__(self) -> None: super().__init__() self.scheduler: _ArnScheduler | None = None self.returned = threading.Event() + self._reentered = False def export(self, record: dict[str, Any]) -> None: super().export(record) assert self.scheduler is not None - self.scheduler.drain(ARN_A) + if self._reentered: + return + self._reentered = True + self.scheduler.schedule(ARN_B, _record("r2")) + self.scheduler.drain(ARN_B) self.returned.set() @@ -907,11 +912,59 @@ def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking( assert ("export", "r1") in exporter.calls # The refused drain still leaves a flush behind, with no external drain to ask - # for one. The re-entering hook may have queued a record, and the worker exits - # once nothing is pending and no flush is requested, so without the request a + # for one. The re-entering hook queued a record, and the worker exits once + # nothing is pending and no flush is requested, so without the request a # buffering exporter would be holding that record when the environment froze. - assert _wait_until(lambda: ("flush", None) in exporter.calls), ( - "a refused drain must request a flush on its way out" - ) + assert _wait_until(lambda: ("export", "r2") in exporter.calls) + assert _wait_until( + lambda: exporter.calls.index(("flush", None)) + > exporter.calls.index(("export", "r2")) + ), "a refused drain must request a flush that covers the record it queued" # The worker is still serving: a drain from any other thread completes. scheduler.drain(ARN_B) + + +class ReentrantFlushExporter(CaptureExporter): + """Exporter whose flush() drains, as a hook re-entered from a flush would. + + An exporter that re-enters a plugin hook from ``flush()`` reaches + ``on_invocation_end``, which drains. The drain arrives on the export worker + from inside a flush, with nothing pending. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.flushes = 0 + + def flush(self) -> None: + super().flush() + self.flushes += 1 + assert self.scheduler is not None + self.scheduler.drain(ARN_A) + + +def test_a_drain_refused_from_inside_a_flush_does_not_re_arm_it() -> None: + """A refused drain with nothing pending asks for no further flush. + + Requesting one unconditionally would keep the worker flushing for as long as + the environment lived: the flush re-enters the hook, the hook drains, the + refused drain asks for the next flush. A pending record is what distinguishes + new work from that loop, so a drain refused from inside a flush leaves no + request behind. + """ + exporter = ReentrantFlushExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + scheduler.schedule(ARN_A, _record("r1")) + scheduler.drain(ARN_A) + + flushes_after_drain = exporter.flushes + assert flushes_after_drain >= 1, "the drain must have flushed" + + # Give a re-armed flush time to appear. The worker retires when nothing is + # pending and no flush is requested, so a bounded count here is the whole + # assertion: an unconditional request never settles. + assert not _wait_until(lambda: exporter.flushes > flushes_after_drain, timeout=1.0) + assert _wait_until(lambda: not scheduler._worker_alive()) 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 336f4403..bc1daa99 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 @@ -723,19 +723,40 @@ def execute_plugins(self, info, sync): A plugin is counted as started before its start hook is dispatched rather than after, because a hook that begins and then fails may already have allocated what its end hook releases. + + The invocation-end hook is the one hook that finishes dispatching even + when a plugin raises one of those three. Every plugin it reaches has + already started, so cutting the loop short costs a plugin its only chance + to finish: Insight 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. """ if not self._executor: return starting = isinstance(info, InvocationStartInfo) + ending = isinstance(info, InvocationEndInfo) + deferred_control: BaseException | None = None for plugin in self._plugins if starting else self._started: if starting: self._started.append(plugin) - if sync: - # this is called synchronously, so plugins will be able to manipulate thread local objects - self._dispatch_plugin(plugin, info) - else: + if not sync: # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) + continue + # this is called synchronously, so plugins will be able to manipulate thread local objects + if not ending: + self._dispatch_plugin(plugin, info) + continue + try: + self._dispatch_plugin(plugin, info) + except _PLUGIN_THREAD_CONTROL_EXCEPTIONS as control: + if deferred_control is None: + deferred_control = control + if deferred_control is not None: + raise deferred_control def _snapshot_operation_infos( self, diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 4d4212ce..7e2c0d79 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -647,6 +647,40 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self): + """Every started plugin receives the end hook, then the thread stops. + + The end hook is a plugin's only chance to finish -- Insight drains there + and OTel ends its spans -- so a plugin raising one of the three control + exceptions must not cost the plugins after it in the list their own end + hook. The exception is held and re-raised once every plugin has been + called. + """ + later = _TrackingPlugin() + host = PluginHost( + plugins=[plugin_factory(_ControlOnEndPlugin()), plugin_factory(later)] + ) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + return { + "Status": ServiceInvocationStatus.SUCCEEDED.value, + "Result": None, + } + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual( + later.calls, ["invocation_start:req-1", "invocation_end:req-1"] + ) + def test_a_start_hook_that_stops_the_thread_leaves_later_plugins_unpaired(self): """A plugin that never received the start hook never receives the end hook. From 2e84c0dc7be2f039b078fc356b57558c3ba10423 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 17:45:53 -0700 Subject: [PATCH 18/28] fix(plugin): split exception groups, drain before the build A BaseExceptionGroup is neither case the containment boundary tested. A group carrying a KeyboardInterrupt is not an instance of one, so the tuple handler naming the three did not match it and the broad handler contained the interrupt inside it. Plugin code produces such a group without asking: an asyncio.TaskGroup whose task is interrupted raises one. Both boundaries now partition what plugin code raised. The control leaves are returned to the caller and re-raised; what remains is logged as a contained plugin failure. A group of ordinary failures is therefore contained whole, and a mixed group is logged and then propagates only its control part. The end-hook dispatch catches BaseException rather than naming the three, because everything reaching it has already been partitioned. Insight registers its invocation-end drain on entering the hook rather than after building the record. _emit runs customer code, and a content transform raising a BaseException escapes _apply_data_content, which contains only Exception. That failure left the hook with the drain unrequested, so records this execution had already scheduled stayed in a buffering exporter when the environment froze. --- .../plugin.py | 23 +++-- .../tests/test_plugin.py | 54 +++++++++++- .../plugin.py | 87 +++++++++++++------ .../tests/plugin_test.py | 68 +++++++++++++++ 4 files changed, 200 insertions(+), 32 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index c0e478d6..17529770 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -389,6 +389,16 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: return emit_mode = self._shared._emit_mode with self._hook_frame(): + # The drain is registered before anything can fail, not after the + # record is built. `_emit` runs customer code -- the content and + # result transforms, and `__del__` on an object a displaced record + # carried -- and a failure there leaves this hook by way of the SDK's + # containment. Registering afterwards meant such a failure skipped + # the drain, so records this execution had already scheduled stayed + # in a buffering exporter when the environment froze. Registering + # here costs nothing when the hook succeeds: the frame runs the drain + # once, on the way out, either way. + self._request_drain() with self._lock: if not self._closed: # Close the gate before emitting so a concurrent late hook for @@ -447,12 +457,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # after this call is done, without waiting on records scheduled after the # call by other executions. # - # Asked for rather than performed here, so it runs when the outermost - # hook frame on this thread unwinds and every `_lock` hold is released. - # See `_hook_frame`: an invocation end that customer code re-entered - # from inside another hook's build would otherwise wait for the export - # worker while holding the lock that worker may need. - self._request_drain() + # Asked for rather than performed, so it runs when the outermost hook + # frame on this thread unwinds and every `_lock` hold is released. See + # `_hook_frame`: an invocation end that customer code re-entered from + # inside another hook's build would otherwise wait for the export + # worker while holding the lock that worker may need. The request + # itself is made at the top of this hook, so a failure in the build + # below cannot skip it. # -- emission ------------------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index a2fe608c..7a3c000b 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -19,12 +19,15 @@ from __future__ import annotations +import asyncio import datetime import itertools import threading import time from typing import Any +import pytest + from aws_durable_execution_sdk_python.lambda_service import ( ErrorObject, OperationStatus, @@ -680,6 +683,7 @@ class ConcurrentCaptureExporter: def __init__(self) -> None: self.records: list[dict[str, Any]] = [] + self.flushes = 0 self._lock = threading.Lock() def render(self, record: dict[str, Any]) -> Any: @@ -690,7 +694,8 @@ def export(self, record: dict[str, Any]) -> None: self.records.append(record) def flush(self) -> None: - pass + with self._lock: + self.flushes += 1 def snapshot(self) -> list[dict[str, Any]]: with self._lock: @@ -1188,6 +1193,53 @@ def hook() -> None: ) +def test_an_end_transform_that_fails_still_drains_what_was_scheduled(): + # `_emit` runs customer code between the snapshot and the hand-off: the content + # transforms, a result override, and __del__ on an object a displaced record + # carried. A failure there leaves the hook through the SDK's containment, and + # the drain used to be requested after the build, so that failure skipped it -- + # leaving records this execution had already scheduled in a buffering exporter + # when the environment froze. The drain is now requested on entry. + exporter = ConcurrentCaptureExporter() + calls: list[str] = [] + + def failing_on_the_second_call(value: Any) -> Any: + calls.append("input") + if len(calls) > 1: + # CancelledError rather than an ordinary exception on purpose: + # _apply_data_content contains Exception so that a failing redactor + # cannot leak the raw value, and a BaseException is what escapes the + # build and leaves the hook. + raise asyncio.CancelledError("transform cancelled") + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=failing_on_the_second_call), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + + # The first emit succeeds and schedules a RUNNING record. + plugin.on_invocation_start(start) + + # The terminal emit fails inside the transform, so this hook raises. The SDK + # contains that; here it is raised directly, which is the same code path. + with pytest.raises(asyncio.CancelledError): + plugin.on_invocation_end(_end(operations=_ops(_step("s")))) + + statuses = [record["status"] for record in exporter.snapshot()] + assert statuses == ["RUNNING"], ( + "the record scheduled before the failing transform must have been drained " + f"to the exporters, not left pending: {statuses}" + ) + assert exporter.flushes >= 1, "the drain must have flushed" + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + def _free_for_another_thread(lock: Any) -> bool: """Report whether a lock is unheld, as seen from a thread that never took it. 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 bc1daa99..76db81d9 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 @@ -545,6 +545,36 @@ def _factory_name(factory: object) -> str: _PLUGIN_THREAD_CONTROL_EXCEPTIONS = (KeyboardInterrupt, SystemExit, GeneratorExit) +def _contain_plugin_failure( + error: BaseException, message: str, *message_args: object +) -> BaseException | None: + """Log what plugin code raised, and return the part that must not be contained. + + Returns ``None`` when the whole failure was contained, or the part that + instructs the calling thread to stop, which the caller re-raises. + + A :class:`BaseExceptionGroup` is split rather than tested, because it is + neither of the two cases a plain ``isinstance`` chain covers: a group carrying + a :class:`KeyboardInterrupt` is not an instance of one, so a tuple handler + naming the three does not match it and a broad handler would contain the + interrupt inside it. Plugin code produces such a group without asking for it + -- an ``asyncio.TaskGroup`` whose task is interrupted raises one -- so the + group is partitioned: the control leaves are returned to be re-raised, and + what remains is logged like any other contained plugin failure. + """ + control: BaseException | None + contained: BaseException | None + if isinstance(error, BaseExceptionGroup): + control, contained = error.split(_PLUGIN_THREAD_CONTROL_EXCEPTIONS) + elif isinstance(error, _PLUGIN_THREAD_CONTROL_EXCEPTIONS): + control, contained = error, None + else: + control, contained = None, error + if contained is not None: + logger.error(message, *message_args, exc_info=contained) + return control + + class PluginExecutor: """One invocation's plugin instances, metadata and dispatch. @@ -632,24 +662,24 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: while the handler is being initialized, so the silent case is not reachable through ``durable_execution()``. - Containment covers every ``BaseException`` except the three that instruct + Containment covers every ``BaseException`` except the parts that instruct the calling thread to stop; see - :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS`. Narrowing it to ``Exception`` - left the contract conditional on a factory never raising outside that - hierarchy, and a factory that awaits a cancelled task raises + :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS` and + :func:`_contain_plugin_failure`. Narrowing it to ``Exception`` left the + contract conditional on a factory never raising outside that hierarchy, + and a factory that awaits a cancelled task raises ``asyncio.CancelledError``, which is outside it. """ plugins: list[DurableInstrumentationPlugin] = [] for factory in self._plugin_factories: try: plugin = factory.create_plugin(info) - except _PLUGIN_THREAD_CONTROL_EXCEPTIONS: - raise - except BaseException: # noqa: BLE001 - a factory must not fail the execution - # log and ignore the exception - logger.exception( - "Plugin factory %s exception ignored", _factory_name(factory) + except BaseException as error: # noqa: BLE001 - a factory must not fail the execution + control = _contain_plugin_failure( + error, "Plugin factory %s exception ignored", _factory_name(factory) ) + if control is not None: + raise control from None continue if plugin is None: logger.error( @@ -678,11 +708,11 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: """Invoke the appropriate plugin callback. Runs inside the thread pool. - Contains every ``BaseException`` except the three that instruct the - calling thread to stop, the same rule the factory boundary uses. The - thread here is the executor's own single worker, which nothing outside - this class cancels or interrupts, so an exception outside the ``Exception`` - hierarchy arriving here was raised by the plugin. + Contains every ``BaseException`` except the parts that instruct the calling + thread to stop, the same rule the factory boundary uses. The thread here + is the executor's own single worker, which nothing outside this class + cancels or interrupts, so an exception outside the ``Exception`` hierarchy + arriving here was raised by the plugin. """ try: match info: @@ -702,11 +732,12 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: plugin.on_user_function_end(info) case _: raise RuntimeError(f"Unknown info type: {type(info)}") - except _PLUGIN_THREAD_CONTROL_EXCEPTIONS: - raise - except BaseException: # noqa: BLE001 - a hook must not fail the execution - # log and ignore the exception - logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) + except BaseException as error: # noqa: BLE001 - a hook must not fail the execution + control = _contain_plugin_failure( + error, "Plugin %s exception ignored", plugin.__class__.__name__ + ) + if control is not None: + raise control from None def execute_plugins(self, info, sync): """Dispatch one hook to this invocation's plugins. @@ -725,14 +756,20 @@ def execute_plugins(self, info, sync): allocated what its end hook releases. The invocation-end hook is the one hook that finishes dispatching even - when a plugin raises one of those three. Every plugin it reaches has - already started, so cutting the loop short costs a plugin its only chance - to finish: Insight would not drain, and OTel would leave spans unended. - The first such exception is held and re-raised once every plugin has been + when a plugin raises one of those. Every plugin it reaches has already + started, so cutting the loop short costs a plugin its only chance to + finish: Insight 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. + + Anything :meth:`_dispatch_plugin` raises is already a thread-control + failure -- it contains everything else -- so the end path catches + ``BaseException`` rather than naming the three again. Naming them would + miss a :class:`BaseExceptionGroup` carrying one, which is what + :func:`_contain_plugin_failure` hands back. """ if not self._executor: return @@ -752,7 +789,7 @@ def execute_plugins(self, info, sync): continue try: self._dispatch_plugin(plugin, info) - except _PLUGIN_THREAD_CONTROL_EXCEPTIONS as control: + 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: diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 7e2c0d79..335fc457 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -647,6 +647,64 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_a_factory_raising_a_group_with_a_control_exception_propagates(self): + """A group carrying a control exception is not contained. + + ``BaseExceptionGroup`` is neither of the two cases an ``isinstance`` chain + covers: a group holding a ``KeyboardInterrupt`` is not an instance of one, + so naming the three in a handler does not match it and a broad handler + would swallow the interrupt inside it. Plugin code produces such a group + without asking for one -- an ``asyncio.TaskGroup`` whose task is + interrupted raises it. + """ + group = BaseExceptionGroup( + "plugin group", [ValueError("contained"), KeyboardInterrupt()] + ) + executor = PluginExecutor(plugins=[_GroupRaisingFactory(group)]) + + with ( + self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs, + executor.run(), + self.assertRaises(BaseExceptionGroup) as raised, + ): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + # Only the control leaf propagates; the rest was logged as a contained + # plugin failure. + self.assertEqual( + [type(leaf) for leaf in raised.exception.exceptions], [KeyboardInterrupt] + ) + self.assertIn("contained", "\n".join(logs.output)) + + def test_a_factory_raising_a_group_without_a_control_exception_is_contained(self): + """A group of ordinary failures is contained like any other.""" + surviving = _TrackingPlugin() + group = BaseExceptionGroup("plugin group", [ValueError("boom")]) + executor = PluginExecutor( + plugins=[_GroupRaisingFactory(group), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("boom", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self): """Every started plugin receives the end hook, then the thread stops. @@ -2380,6 +2438,16 @@ def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlug raise self._control +class _GroupRaisingFactory: + """Factory that raises a ``BaseExceptionGroup``, as an ``asyncio.TaskGroup`` does.""" + + def __init__(self, group: BaseExceptionGroup) -> None: + self._group = group + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise self._group + + class _CancellingPlugin(DurableInstrumentationPlugin): """Plugin whose hook raises outside the ``Exception`` hierarchy.""" From 0559917f94638a04ffa8885524f7ad67be7a3b79 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 20:19:42 -0700 Subject: [PATCH 19/28] fix(plugin): dispatch on one thread, keep the real failure Two findings from a fresh review pass. The failure exit fired the end hook inside its except block, so an exception out of the hook left that block before the handler's own failure was re-raised. The caller saw the plugin's exception and the real failure survived only as __context__. The end-hook dispatch does raise: it re-raises the control exceptions it holds through the fan-out. Instrumentation does not decide what an execution failed with, so the hook's exception is now contained and logged there and the original is re-raised unchanged. execute_plugins also had a sync parameter no caller ever passed. Its asynchronous branch submitted the hook to a per-invocation thread pool, which skipped the end-hook fan-out rule and swallowed a control exception in a Future nobody reads -- a way to opt out of the invariants this PR establishes, on a path with no production caller. The parameter and the pool are removed: hooks are dispatched on the calling thread, which is what lets a plugin set thread-affine state the SDK's own logging reads, and the gate that used to test the pool now tests the factory list. The pool created no thread either way, because ThreadPoolExecutor spawns lazily on first submit. --- .../plugin.py | 82 +++++++------- .../tests/plugin_test.py | 101 ++++++++++++++---- 2 files changed, 122 insertions(+), 61 deletions(-) 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 76db81d9..a8d665c1 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,6 @@ import functools import logging from collections.abc import Iterator, Mapping, Sequence -from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, MutableMapping, Protocol, cast @@ -603,7 +602,6 @@ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): # Every later hook is dispatched to this list, so a plugin that never # received its start hook never receives its end hook. self._started: list[DurableInstrumentationPlugin] = [] - self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None self._run_entered = False @@ -625,25 +623,14 @@ class exists to prevent; failing loudly here keeps the bug from ) raise RuntimeError(msg) self._run_entered = True - if self._plugin_factories: - self._executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="plugin-executor", - ) try: yield finally: self._invocation_status = None self._operations_provider = None - # Shut down the thread pool, waiting for pending tasks to complete. - # The pool belongs to this invocation, so this drains only this - # invocation's queued dispatches and cannot cut short a concurrent - # invocation's. - if self._executor: - self._executor.shutdown(wait=True) - # Drop this invocation's plugin instances. After the pool has - # drained, so no queued dispatch still holds one: nothing outlives - # the invocation. + # Drop this invocation's plugin instances: nothing outlives the + # invocation. Every dispatch is synchronous, so there is no queued + # work still holding one. self._plugins = [] self._started = [] @@ -739,7 +726,7 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: if control is not None: raise control from None - def execute_plugins(self, info, sync): + def execute_plugins(self, info): """Dispatch one hook to this invocation's plugins. A plugin receives a hook only once it has received the invocation-start @@ -770,8 +757,18 @@ def execute_plugins(self, info, sync): ``BaseException`` rather than naming the three again. Naming them would miss a :class:`BaseExceptionGroup` carrying one, which is what :func:`_contain_plugin_failure` hands back. + + Every hook is dispatched on the calling thread. That is what lets a + plugin set a ``ThreadLocal`` or an MDC key the SDK's own logging then + reads, and it is what makes the pairing and re-raise rules above + enforceable: a hook dispatched to a pool would land in a + :class:`~concurrent.futures.Future` nobody reads, so a control exception + 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. """ - if not self._executor: + if not self._plugin_factories: return starting = isinstance(info, InvocationStartInfo) ending = isinstance(info, InvocationEndInfo) @@ -779,11 +776,6 @@ def execute_plugins(self, info, sync): for plugin in self._plugins if starting else self._started: if starting: self._started.append(plugin) - if not sync: - # this is called asynchronously, so plugins cannot manipulate thread local objects - self._executor.submit(self._dispatch_plugin, plugin, info) - continue - # this is called synchronously, so plugins will be able to manipulate thread local objects if not ending: self._dispatch_plugin(plugin, info) continue @@ -870,7 +862,7 @@ def on_invocation_start( # Build this invocation's plugin instances from the very info their first # hook receives, and before that hook is dispatched. self._create_plugins(self._invocation_status) - self.execute_plugins(self._invocation_status, sync=True) + self.execute_plugins(self._invocation_status) def _snapshot_execution_input(self, execution_input: Any) -> Any: """Deep-copy the execution input so the plugin view is isolated. @@ -918,7 +910,7 @@ def on_invocation_end( operations=self._snapshot_operation_infos(self._operations_provider), ) ) - self.execute_plugins(invocation_end_info, sync=True) + self.execute_plugins(invocation_end_info) def on_user_function_start( self, @@ -939,7 +931,7 @@ def on_user_function_start( is_replay_children=is_replay_children, attempt=attempt, ) - self.execute_plugins(start_info, sync=True) + self.execute_plugins(start_info) return start_info def on_user_function_end( @@ -952,7 +944,6 @@ def on_user_function_end( """Execute plugins when a user function returns, fails, or is incomplete.""" self.execute_plugins( UserFunctionEndInfo.from_start_info(start_info, error, outcome=outcome), - sync=True, ) def on_operation_action( @@ -982,7 +973,6 @@ def on_operation_action( is_replayed=previous_operation is not None, status=OperationStatus.STARTED, ), - sync=True, ) def on_operation_replay(self, operation: Operation) -> None: @@ -1000,7 +990,7 @@ def on_operation_replay(self, operation: Operation) -> None: is_replayed=True, status=operation.status, ) - self.execute_plugins(start_info, sync=True) + self.execute_plugins(start_info) def on_child_context_end( self, @@ -1025,7 +1015,6 @@ def on_child_context_end( error=error, is_replayed=is_replayed, ), - sync=True, ) def on_operation_update( @@ -1075,7 +1064,6 @@ def on_operation_update( ), is_replayed=False, ), - sync=True, ) if ( @@ -1103,7 +1091,6 @@ def on_operation_update( }, operations=_to_operation_info_map(operations), ), - sync=True, ) @staticmethod @@ -1204,13 +1191,30 @@ def wrapper(event: Any, context: LambdaContext): # they also fire the hook. The hook is what a plugin needs # to flush, and a process being torn down is when flushing # matters; the cost is the same bounded work any - # invocation end does. The exception itself is re-raised - # unchanged, so what the caller sees is untouched. - plugin_executor.on_invocation_end( - output=DurableExecutionInvocationOutput.create_retry( - ErrorObject.from_exception(e) - ), - ) + # invocation end does. + # + # The handler's exception is what the caller sees, + # whatever the hook does. The end-hook dispatch can raise + # -- it re-raises the control exceptions it holds through + # the fan-out -- and letting that replace the handler's + # failure would report an instrumentation problem as the + # execution's outcome and leave the real failure reachable + # only as __context__. Instrumentation does not decide + # what an execution failed with, so the hook's exception + # is contained here and the original is re-raised + # unchanged. + try: + plugin_executor.on_invocation_end( + output=DurableExecutionInvocationOutput.create_retry( + ErrorObject.from_exception(e) + ), + ) + except BaseException: # noqa: BLE001 - the handler's failure wins + logger.exception( + "Plugin invocation-end hook failed while the " + "invocation was already failing; the original " + "failure is raised" + ) raise plugin_executor.on_invocation_end(output=completed) return output diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 335fc457..200af119 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -3,6 +3,7 @@ import datetime import logging import pickle +import threading import unittest from collections.abc import Iterator from copy import deepcopy @@ -705,6 +706,38 @@ def test_a_factory_raising_a_group_without_a_control_exception_is_contained(self self.assertIn("boom", "\n".join(logs.output)) self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_a_failing_end_hook_does_not_replace_the_handlers_failure(self): + """Instrumentation does not decide what an execution failed with. + + The end-hook dispatch can raise: it re-raises the control exceptions it + holds through the fan-out. On the failure exit that raise used to leave + the ``except`` block before the handler's own exception was re-raised, so + the caller saw the plugin's exception and the real failure survived only + as ``__context__``. + """ + plugin = _ControlOnEndPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise ValueError("the real failure") + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ): + with self.assertRaises(ValueError) as raised: + handler({}, LAMBDA_CTX) + + self.assertEqual(str(raised.exception), "the real failure") + # The hook still ran, and still saw the failing outcome. + self.assertEqual(plugin.end_statuses, [InvocationStatus.RETRY]) + def test_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self): """Every started plugin receives the end hook, then the thread stops. @@ -1136,26 +1169,50 @@ def test_a_hook_raising_cancellation_is_contained(self): "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ) as logs: with _invocation(executor, tracking): - executor.execute_plugins(OPERATION_START_INFO, sync=True) + executor.execute_plugins(OPERATION_START_INFO) self.assertIn("hook cancelled", "\n".join(logs.output)) self.assertIn("operation_start:op-2", tracking.calls) class TestPluginExecutor(unittest.TestCase): - def test_no_thread_pool_when_plugins_is_none(self): - """Tests that PluginExecutor does not create a thread pool when plugins is empty.""" - executor = PluginExecutor(plugins=None) - self.assertIsNone(executor._executor) - - def test_no_thread_pool_when_plugins_is_empty_list(self): - executor = PluginExecutor(plugins=[]) - self.assertIsNone(executor._executor) + def test_dispatch_is_a_no_op_when_no_factory_is_registered(self): + """An executor with no factories dispatches nothing and needs no thread. + + Hooks are dispatched on the calling thread, which is what lets a plugin + set thread-affine state the SDK's own logging then reads, and what makes + the end-hook pairing and re-raise rules enforceable. There is therefore no + pool to create, and an executor with nothing registered returns before it + touches anything. + """ + for plugins in (None, []): + with self.subTest(plugins=plugins): + executor = PluginExecutor(plugins=plugins) + self.assertEqual(executor._plugin_factories, []) + with executor.run(): + # Nothing registered, so no hook reaches a plugin and no + # dispatch raises on the way through. + executor.execute_plugins(INVOCATION_START_INFO) + executor.execute_plugins(OPERATION_START_INFO) + + def test_hooks_are_dispatched_on_the_calling_thread(self): + """Thread affinity is the contract, so it is asserted rather than assumed.""" + seen: list[str] = [] + + class _ThreadRecordingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + seen.append(threading.current_thread().name) - def test_thread_pool_created_when_plugins_provided(self): - executor = PluginExecutor(plugins=[plugin_factory(_NoOpPlugin())]) + executor = PluginExecutor(plugins=[plugin_factory(_ThreadRecordingPlugin())]) with executor.run(): - self.assertIsNotNone(executor._executor) + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertEqual(seen, [threading.current_thread().name]) def test_start_is_noop_when_empty(self): executor = PluginExecutor(plugins=[]) @@ -1234,37 +1291,37 @@ def setUp(self): def test_dispatch_invocation_start_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(INVOCATION_START_INFO, sync=True) + self.executor.execute_plugins(INVOCATION_START_INFO) self.assertIn("invocation_start:req-1", self.plugin.calls) def test_dispatch_invocation_end_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(INVOCATION_END_INFO, sync=True) + self.executor.execute_plugins(INVOCATION_END_INFO) self.assertIn("invocation_end:req-1", self.plugin.calls) def test_dispatch_operation_end_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(OPERATION_END_INFO, sync=False) + self.executor.execute_plugins(OPERATION_END_INFO) self.assertIn("operation_end:op-1", self.plugin.calls) def test_dispatch_operation_start_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(OPERATION_START_INFO, sync=False) + self.executor.execute_plugins(OPERATION_START_INFO) self.assertIn("operation_start:op-2", self.plugin.calls) def test_dispatch_operation_change_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(OPERATION_CHANGE_INFO, sync=False) + self.executor.execute_plugins(OPERATION_CHANGE_INFO) self.assertIn("operation_change:op-1", self.plugin.calls) def test_dispatch_user_function_start_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(USER_FUNCTION_START_INFO, sync=True) + self.executor.execute_plugins(USER_FUNCTION_START_INFO) self.assertIn("user_function_start:op-1", self.plugin.calls) def test_dispatch_user_function_end_info(self): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins(USER_FUNCTION_END_INFO, sync=True) + self.executor.execute_plugins(USER_FUNCTION_END_INFO) self.assertIn("user_function_end:op-1", self.plugin.calls) def test_dispatch_unknown_type_logs_exception(self): @@ -1273,7 +1330,7 @@ def test_dispatch_unknown_type_logs_exception(self): "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): with _invocation(self.executor, self.plugin): - self.executor.execute_plugins("not a valid info type", sync=True) + self.executor.execute_plugins("not a valid info type") def test_plugin_exception_is_swallowed(self): """If a plugin raises, the exception is logged and execution continues.""" @@ -1287,7 +1344,7 @@ def test_plugin_exception_is_swallowed(self): "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): with _invocation(executor, tracking_plugin): - executor.execute_plugins(OPERATION_START_INFO, sync=True) + executor.execute_plugins(OPERATION_START_INFO) # The second plugin should still have been called self.assertIn("operation_start:op-2", tracking_plugin.calls) @@ -1298,7 +1355,7 @@ def test_multiple_plugins_all_called(self): executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)]) with _invocation(executor, p1, p2): - executor.execute_plugins(OPERATION_START_INFO, sync=True) + executor.execute_plugins(OPERATION_START_INFO) self.assertIn("operation_start:op-2", p1.calls) self.assertIn("operation_start:op-2", p2.calls) From aeea6ff932c898576943f339baa479915e3fd134 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 20:19:42 -0700 Subject: [PATCH 20/28] docs(insight): correct two rationales, note drain latency The log filter's bind_invocation took a process-global lock on every operation-start and user-function-start hook although the work is idempotent after the first call on a thread. It now returns early when the thread already carries this provider's claim. Two comments described an ordering the SDK prevents. Insight's lock said a checkpoint-path operation-change "genuinely races" the invocation-end hook; the SDK joins the checkpoint thread and the branch pools before that hook is dispatched. The lock stays -- it is what makes reentrancy from customer code inside a build safe, and a guard resting on the SDK's join ordering is one refactor from being wrong -- but the comment now says which of the two it is. The Insight README documents what the invocation-end drain costs under concurrency: it waits for every record any execution had pending, and one worker serializes every export and flush, so ends are released together at the slowest. Measured with a 30 ms exporter: ~72 ms at one execution, ~1.8 s each at 48. --- .../README.md | 13 +++++++++++++ .../plugin.py | 15 ++++++++++----- .../log_filter.py | 8 ++++++++ 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 2d3c9289..f3c2847b 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -230,6 +230,19 @@ Behavior is validated cross-SDK by the `insight` conformance suite > cadence the JS and Java plugins have. Only a sampled-out execution neither > exports nor flushes. +> **Note (invocation-end latency under concurrency).** The drain an invocation +> end performs waits for every record any execution had pending when it was +> called, and one worker serializes all exports and all flushes, so every +> concurrently ending invocation is released together at the slowest one. The wait +> therefore grows with the number of executions the environment is running, not +> just with this execution's own work: measured with a 30 ms exporter, one +> execution ended in ~72 ms and 48 concurrent executions in ~1.8 s each. That is +> the deliberate trade against the alternative — releasing an end before its +> record reached the exporters, which is what silently lost terminal records +> before. It matters for an exporter that makes a network call per record: budget +> invocation-end time against the environment's concurrency, not against one +> execution. + ## Requirements - `aws-durable-execution-sdk-python` with the plugin invocation hooks that diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 17529770..4cf0e0c3 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -281,11 +281,16 @@ def __init__( self._build_revision = 0 # Guards `_closed`, `_build_revision`, the operations rebind and record # emission, so a late hook can never slip a RUNNING record in after the - # terminal one. Still earns its place with one instance per invocation: the - # SDK dispatches every hook synchronously on the thread that produced the - # event, so an operation-change raised off the checkpointing path runs - # concurrently with the invocation thread's on_invocation_end -- two hooks, - # one instance, genuinely racing. + # terminal one. + # + # What it protects against is reentrancy, not two threads. The SDK + # dispatches every hook synchronously on the thread that produced the + # event, and it joins the checkpoint thread and the branch pools before + # the invocation-end hook is dispatched, so a checkpoint-path + # operation-change cannot overlap `on_invocation_end` -- an earlier + # version of this comment claimed it could. The lock still earns its place + # for the reason below, and it stays because a guard whose correctness + # rests on the SDK's join ordering is one refactor away from being wrong. # # Reentrant on purpose: `_emit` runs the scheduler's `schedule()` inside # this hold, and `schedule()` releases the record it displaces, which can 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 0a88540a..afc07060 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 @@ -132,9 +132,17 @@ def bind_invocation(provider: _SpanContextProvider) -> None: executing user code. Idempotent, so a plugin can call it from every such hook without tracking which threads it has already claimed. + A thread that already carries this provider's claim returns before taking the + registry lock. Every operation-start and user-function-start hook calls this, + and after the first call on a thread there is nothing to add: membership in + the open-invocation set is idempotent, and the claim is already in place. + Args: provider: The plugin serving the invocation that owns this thread. """ + claim = _current_invocation.get() + if claim is not None and claim() is provider: + return with _registry_lock: _open_invocations.add(provider) _current_invocation.set(weakref.ref(provider)) From 9a7206db800d85c616cf06fe42bbdb9c4d230eb1 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 20:19:42 -0700 Subject: [PATCH 21/28] fix(deps): bound the core dependency above Both plugin packages required aws-durable-execution-sdk-python>=3.0.0 with no ceiling. The lower bound exists because their entry points resolve to factories, which core 2.x cannot call: pip accepts the resolution and the handler fails at initialization. Without a ceiling the next core major that changes the plugin contract reproduces that exactly. Both are now >=3.0.0,<4, and each package's metadata test asserts the specifier rejects the next core major as well as accepting the current one. A second case asserts it still admits the next core patch, because a <= ceiling looks equivalent and excludes it. --- .../pyproject.toml | 7 ++- .../tests/test_package_metadata.py | 46 ++++++++++++++++- .../pyproject.toml | 7 ++- .../tests/test_package_metadata.py | 50 ++++++++++++++++++- pyproject.toml | 2 +- 5 files changed, 104 insertions(+), 8 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml index 0b25d4e1..af270c89 100644 --- a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml @@ -21,12 +21,15 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", ] dependencies = [ - # >=3.0.0: the first release whose `plugins` argument takes factories. + # >=3.0.0,<4: 3.0.0 is the first release whose `plugins` argument takes + # factories, and the ceiling is the same reasoning applied forwards -- the next + # core major that changes the plugin contract would install and then fail at + # handler initialization exactly as 2.x does below. # `workflow_insight()` returns a factory, which core 2.x cannot call, so an # install resolved against 2.x fails at handler initialization. 3.0.0 also # carries the invocation-hook fields this plugin reads # (InvocationInfo.execution_input / InvocationEndInfo.execution_result). - "aws-durable-execution-sdk-python>=3.0.0", + "aws-durable-execution-sdk-python>=3.0.0,<4", ] [project.optional-dependencies] diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py index e9130b3b..8134746f 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py @@ -16,6 +16,7 @@ import tomllib from pathlib import Path +from packaging.specifiers import SpecifierSet from packaging.version import Version @@ -48,7 +49,7 @@ def _core_dependency_lower_bound() -> str: dependencies = tomllib.load(pyproject)["project"]["dependencies"] bounds = [ - dependency.removeprefix(CORE_DISTRIBUTION + ">=") + dependency.removeprefix(CORE_DISTRIBUTION + ">=").split(",", 1)[0] for dependency in dependencies if dependency.startswith(CORE_DISTRIBUTION + ">=") ] @@ -76,3 +77,46 @@ def test_core_dependency_bound_matches_the_core_major_in_this_repository() -> No assert _major(lower_bound) == _major(core_version) assert Version(lower_bound) <= Version(core_version) + + +def test_core_dependency_excludes_the_next_core_major() -> None: + """A lower bound alone is the same defect one major later. + + The lower bound exists because this package's entry points resolve to plugin + factories, which the core major below cannot call: pip accepts the resolution + and the handler fails at initialization. Without a ceiling the next core major + that changes the plugin contract reproduces exactly that, so the specifier has + to reject it rather than only reject what came before. + """ + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + specifiers = [ + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ] + assert len(specifiers) == 1 + + core_major = _major(_core_version()) + assert specifiers[0].contains(_core_version(), prereleases=True) + assert not specifiers[0].contains(f"{core_major + 1}.0.0", prereleases=True) + + +def test_core_dependency_admits_a_later_core_patch() -> None: + """The ceiling belongs on the major, not on the version built here. + + A ``<=`` ceiling looks equivalent and is not: it excludes the next core patch, + so the first core patch release puts this claim out of date for a change that + cannot have touched the plugin contract. + """ + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + specifier = next( + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ) + + core = Version(_core_version()) + next_patch = f"{core.major}.{core.minor}.{core.micro + 1}" + assert specifier.contains(next_patch, prereleases=True) diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index de1d940d..574ff6fe 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,11 +22,14 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - # >=3.0.0: the first release whose `plugins` argument takes factories. + # >=3.0.0,<4: 3.0.0 is the first release whose `plugins` argument takes + # factories, and the ceiling is the same reasoning applied forwards -- the next + # core major that changes the plugin contract would install and then fail at + # handler initialization exactly as 2.x does below. # DurableInstrumentationPluginProvider was removed in 3.0.0, and this package's # entry points resolve to factories, so core 2.x accepts the install and then # fails at handler initialization. - "aws-durable-execution-sdk-python>=3.0.0", + "aws-durable-execution-sdk-python>=3.0.0,<4", ] [project.entry-points."aws_durable_execution.plugins"] diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 59ef0a13..5a59d9bf 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -2,13 +2,14 @@ import tomllib from pathlib import Path +from packaging.specifiers import SpecifierSet from packaging.version import Version PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] CORE_DISTRIBUTION = "aws-durable-execution-sdk-python" -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=3.0.0" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=3.0.0,<4" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -112,7 +113,7 @@ def _core_dependency_lower_bound() -> str: "dependencies" ] bounds = [ - dependency.removeprefix(CORE_DISTRIBUTION + ">=") + dependency.removeprefix(CORE_DISTRIBUTION + ">=").split(",", 1)[0] for dependency in dependencies if dependency.startswith(CORE_DISTRIBUTION + ">=") ] @@ -157,3 +158,48 @@ def test_layer_sdk_pin_matches_the_core_version_in_this_repository() -> None: pinned_version = tomllib.load(metadata_file)["layer"]["sdk-version"] assert pinned_version == _core_version() + + +def test_core_dependency_excludes_the_next_core_major() -> None: + """A lower bound alone is the same defect one major later. + + The lower bound exists because this package's entry points resolve to plugin + factories, which the core major below cannot call: pip accepts the resolution + and the handler fails at initialization. Without a ceiling the next core major + that changes the plugin contract reproduces exactly that, so the specifier has + to reject it rather than only reject what came before. + """ + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + specifiers = [ + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ] + assert len(specifiers) == 1 + + core_major = _major(_core_version()) + assert specifiers[0].contains(_core_version(), prereleases=True) + assert not specifiers[0].contains(f"{core_major + 1}.0.0", prereleases=True) + + +def test_core_dependency_admits_a_later_core_patch() -> None: + """The ceiling belongs on the major, not on the version built here. + + A ``<=`` ceiling looks equivalent and is not: it excludes the next core patch, + so the first core patch release puts this claim out of date for a change that + cannot have touched the plugin contract. + """ + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + specifier = next( + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ) + + core = Version(_core_version()) + next_patch = f"{core.major}.{core.minor}.{core.micro + 1}" + assert specifier.contains(next_patch, prereleases=True) diff --git a/pyproject.toml b/pyproject.toml index 9c2e4972..ab0f0b85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" [tool.hatch.envs.test-pypi-otel] dependencies = [ - "aws-durable-execution-sdk-python>=3.0.0", + "aws-durable-execution-sdk-python>=3.0.0,<4", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", From 869f6b95659e069b6f78039d3bcaeace615b590e Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 20:46:16 -0700 Subject: [PATCH 22/28] fix(insight): flush a record already taken for export The condition on the refused drain's flush request was pending-only, and the worker takes a record out of the pending map before it calls the exporter. A hook re-entered from inside export() that reached an invocation end emitting no record therefore found nothing pending, asked for no flush, and left the snapshot it had just been handed in a buffering exporter until the environment froze -- the loss the request exists to prevent, arriving by the other door. The condition now separates the two re-entry paths, because they need opposite answers: a flush already in flight with nothing newly pending is the one case that asks for nothing, which is what stops an exporter whose flush() re-enters from flushing for as long as the environment lives. The plugin contract also said a plugin may hold "per-execution" state in its attributes. Per invocation is narrower, and the difference is the replay model: an execution spans as many invocations as it waits, retries or resumes, and the instance is dropped when each returns. The README and the factory protocol now say per-invocation and name the two homes for anything that has to outlive it -- the operation-map snapshot the invocation hooks carry, or the factory. --- .../_export_scheduler.py | 23 +++++---- .../tests/test_export_scheduler.py | 47 +++++++++++++++++++ .../README.md | 13 ++++- .../plugin.py | 11 ++++- 4 files changed, 83 insertions(+), 11 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index dd1a52a8..66ffb205 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -301,22 +301,27 @@ def _is_export_worker(self) -> bool: return self._worker is threading.current_thread() def _request_flush_for_pending_records(self) -> None: - """Ask the worker for a flush, but only if a record is waiting for one. + """Ask the worker for a flush, unless a flush already in flight covers it. Returns without waiting, so it is safe to call from the worker itself. The barrier is raised to the current schedule counter, which is what makes the flush cover a record queued moments ago rather than running before it. - A pending record is the condition, not a formality. This is called from a - drain refused on the worker thread, and one way to reach that is an - exporter whose ``flush()`` re-enters a plugin hook: the call then arrives - from inside a flush, and requesting the next one unconditionally would - produce a flush that re-enters, requests, and flushes again for as long as - the environment lives -- after the invocation has returned. Nothing is - pending in that case, so nothing is requested. + The condition separates the two ways a refused drain is reached, because + they need opposite answers. Re-entered from an exporter's ``export()``, + the record has already left ``_pending`` -- the worker takes it before it + calls the exporter -- so a pending-only test would skip the request and + leave that snapshot buffered until the environment froze. Re-entered from + an exporter's ``flush()``, a flush is in flight and covers what was + exported before it, so requesting another would produce a flush that + re-enters, requests, and flushes again for as long as the environment + lived -- after the invocation returned. A flush in flight with nothing + newly pending is therefore the one case that asks for nothing. """ with self._condition: - if self._disabled or not self._pending: + if self._disabled: + return + if not self._pending and self._flush_in_flight is not None: return self._flush_requested = True self._flush_barrier = max(self._flush_barrier, self._seq) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index a09a236f..7f9f454d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -968,3 +968,50 @@ def test_a_drain_refused_from_inside_a_flush_does_not_re_arm_it() -> None: # assertion: an unconditional request never settles. assert not _wait_until(lambda: exporter.flushes > flushes_after_drain, timeout=1.0) assert _wait_until(lambda: not scheduler._worker_alive()) + + +class ReentrantDrainDuringExportExporter(CaptureExporter): + """Exporter whose export() drains without queueing anything new. + + The worker takes a record out of the pending map before it calls the + exporter, so a hook re-entered from inside ``export()`` and reaching an + invocation end that emits no record finds nothing pending -- while the record + it was just handed is still only in the exporter's buffer. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.returned = threading.Event() + self._reentered = False + + def export(self, record: dict[str, Any]) -> None: + super().export(record) + assert self.scheduler is not None + if self._reentered: + return + self._reentered = True + self.scheduler.drain(ARN_B) + self.returned.set() + + +def test_a_drain_refused_from_inside_an_export_still_flushes() -> None: + """A refused drain from inside an export asks for a flush. + + The record it must cover has already left the pending map, so a + pending-only condition would leave that snapshot in a buffering exporter + when the environment froze -- which is the loss the refusal path exists to + prevent, arriving by the other door. + """ + exporter = ReentrantDrainDuringExportExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + scheduler.schedule(ARN_A, _record("r1")) + + assert exporter.returned.wait(timeout=10), "the refused drain must return" + assert _wait_until(lambda: ("flush", None) in exporter.calls), ( + "the exported record must be flushed even though nothing was pending" + ) + assert ("export", "r1") in exporter.calls + assert _wait_until(lambda: not scheduler._worker_alive()) diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 214f1422..d2df0d1f 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -47,7 +47,18 @@ A plugin is registered as a *factory*, not as an instance. A factory is an objec with a `create_plugin(info)` method taking the invocation's `InvocationStartInfo` and returning a `DurableInstrumentationPlugin`; the SDK calls that method once per invocation, so the instance it returns serves that one invocation only and can -hold per-execution state in ordinary attributes. +hold **per-invocation** state in ordinary attributes. + +Per-invocation is narrower than per-execution, and the difference matters. A +durable execution spans as many invocations as it waits, retries or resumes, and +the instance is dropped when each of those returns — so anything a plugin keeps in +its attributes is gone by the next invocation of the same execution. State that +has to survive that has two honest homes: rebuild it from the operation map the +invocation hooks carry (`InvocationStartInfo.operations` is a full snapshot, +which is how the bundled Insight plugin reports operations that completed in an +earlier invocation), or put it on the factory, which outlives every invocation — +keyed by execution ARN, and pruned by the owner, because the SDK will not tell the +factory when an execution ends for good. A factory is an object with a method rather than a plain callable so the registration type can grow a second, optional member later -- a process-level 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 a8d665c1..874d27d5 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 @@ -505,9 +505,18 @@ def create_plugin( :class:`InvocationStartInfo` -- the same object the returned instance's ``on_invocation_start`` then receives -- and before any hook fires. The instance serves only that invocation and is dropped when it returns, so a - plugin can hold per-execution state in ordinary instance attributes + plugin can hold that invocation's state in ordinary instance attributes without keying it by execution ARN. + Per invocation is narrower than per execution. A durable execution spans + as many invocations as it waits, retries or resumes, so state a plugin + leaves in its attributes is gone by the next invocation of the same + execution. Anything that has to survive that is rebuilt from the operation + map the invocation hooks carry -- ``InvocationStartInfo.operations`` is a + full snapshot, including operations that completed in an earlier + invocation -- or kept on the factory, which outlives every invocation and + is therefore the caller's to key and to prune. + ``info`` is positional-only, so an implementation may name the parameter whatever reads best; a named protocol parameter would pin that name for every implementation. From f5174ab3184872df0f3bf83a44fbe8482ae1264f Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 21:14:06 -0700 Subject: [PATCH 23/28] fix(insight): release displaced records outside every lock `_emit` hands a record to the scheduler while holding this instance's lock, and that hand-off displaces the record already pending for the execution. Releasing the displaced one can run a customer `__del__`, because the record is a snapshot of customer data, and a finalizer that reaches another execution's plugin blocks on that instance's lock. Two threads doing that to each other at once hang both invocations, and the invocation end is awaited before the response. The reentrant lock does not help: it covers re-entry on the same instance, which is why customer code inside a build can call this execution's hooks safely, and says nothing about a second instance. `schedule()` now returns what it displaced and the hook frame releases it once the outermost hook returns, with no lock held. Releases run before the deferred drains, so a finalizer that schedules a record is covered by the drain that follows it. --- .../_export_scheduler.py | 23 ++++++-- .../plugin.py | 27 +++++++-- .../tests/test_plugin.py | 56 +++++++++++++++++++ 3 files changed, 96 insertions(+), 10 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 66ffb205..4de1f930 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -144,14 +144,28 @@ def __init__(self, exporters: list[InsightExporter]) -> None: # a worker that keeps making progress never approaches the bound. self._worker_faults = 0 - def schedule(self, execution: _ExportState, record: dict[str, Any]) -> None: - """Replace this execution's pending snapshot; never runs exporters inline.""" + def schedule(self, execution: _ExportState, record: dict[str, Any]) -> list[Any]: + """Replace this execution's pending snapshot; never runs exporters inline. + + Returns the records this call displaced, for the caller to release once it + holds no lock. They are handed back rather than dropped here because + releasing one can run a customer ``__del__``, and this method is called + from inside the calling plugin's own lock. A finalizer that re-entered a + *different* execution's plugin would then block on that instance's lock + while holding this one, which deadlocks both invocations if the mirror + image happens on another thread at the same time. The plugin's hook frame + drops them when the outermost hook returns; see + ``WorkflowInsightPlugin._hook_frame``. + + A caller with no frame to defer to may drop the returned list + immediately: doing so is only unsafe while a plugin lock is held. + """ displaced: dict[str, Any] | None = None failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: if self._disabled: - return + return [] self._seq += 1 execution.scheduled_seq = self._seq displaced = execution.pending_record @@ -161,14 +175,13 @@ def schedule(self, execution: _ExportState, record: dict[str, Any]) -> None: self._pending[execution] = None failed_pending, start_error = self._ensure_worker_locked() self._condition.notify_all() - # Releasing either record may run custom finalizers, so do it unlocked. - del displaced, failed_pending if start_error is not None: _logger.warning( "workflow-insight: could not start export worker; disabling " "asynchronous export: %s", start_error, ) + return [item for item in (displaced, failed_pending) if item is not None] def drain(self, execution: _ExportState) -> None: """Wait until this execution's latest record is exported and exporters flush. diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 4cf0e0c3..a46c947d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -191,6 +191,11 @@ class _HookFrames(threading.local): def __init__(self) -> None: self.depth = 0 self.pending: list[WorkflowInsightPlugin] = [] + # Records displaced from the scheduler's pending slots, held until every + # plugin lock on this thread is released. Releasing one can run a customer + # finalizer, and a finalizer that reaches another execution's plugin must + # not do so while this thread holds a plugin lock. + self.releases: list[Any] = [] _hook_frames = _HookFrames() @@ -341,10 +346,16 @@ def _hook_frame(self) -> Iterator[None]: yield finally: state.depth -= 1 - if state.depth == 0 and state.pending: - owed, state.pending = state.pending, [] - for plugin in owed: - plugin._drain() + if state.depth == 0: + # Released before the drains, and with no lock held: a finalizer + # that schedules a record is then covered by the drain that + # follows it. + if state.releases: + state.releases.clear() + if state.pending: + owed, state.pending = state.pending, [] + for plugin in owed: + plugin._drain() def _request_drain(self) -> None: """Ask for a drain once the outermost hook frame on this thread unwinds.""" @@ -645,7 +656,13 @@ def _emit( # every later non-terminal record is rejected above. if not closing and (self._closed or revision != self._build_revision): return - self._shared._scheduler.schedule(self, record) + # The records this hand-off displaces are released by the hook frame, not + # here: this runs inside `_lock`, and a displaced record can carry a + # customer object whose `__del__` re-enters a hook. Re-entering *this* + # execution's hook is safe because `_lock` is reentrant; re-entering + # another execution's is not, and two threads doing it to each other at + # once would hang both invocations. See `_hook_frame`. + _hook_frames.releases.extend(self._shared._scheduler.schedule(self, record)) class _WorkflowInsightFactory: diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 7a3c000b..2b40dc17 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -1240,6 +1240,62 @@ def failing_on_the_second_call(value: Any) -> Any: assert _wait_until(lambda: not factory._scheduler._worker_alive()) +def test_a_displaced_records_finalizer_runs_with_no_plugin_lock_held(): + # `_emit` hands the record to the scheduler while holding this instance's + # `_lock`, and that hand-off displaces the record already pending for this + # execution. Releasing the displaced one can run a customer `__del__` -- it is + # a snapshot of customer data -- and a finalizer that reaches ANOTHER + # execution's plugin blocks on that instance's lock. Two threads doing that to + # each other at once hang both invocations, which the reentrant lock does not + # help with: it only covers re-entry on the same instance. The hook frame + # therefore releases displaced records after every lock is dropped. + exporter = ConcurrentCaptureExporter() + observed: list[bool] = [] + holder: dict[str, Any] = {} + + class _FinalizerProbe: + def __del__(self) -> None: + plugin = holder.get("plugin") + if plugin is None: + return + # `_is_owned()` answers "does the calling thread hold this lock", + # which is the question here. Acquiring it would not: an RLock lets + # its owner acquire it again. + observed.append(not plugin._lock._is_owned()) + + # A transform returning a fresh object per emit is what makes the record the + # only reference to it, so releasing the displaced record is what collects it. + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=lambda value: _FinalizerProbe()), + ) + ) + plugin = factory.create_plugin(_start(operations={})) + holder["plugin"] = plugin + + # The first emit queues a record holding a probe; the next displaces it while + # the worker is still parked, so the release happens on this thread. + plugin.on_invocation_start(_start(operations={})) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, + updated_operations=_ops(_step("s")), + operations=_ops(_step("s")), + ) + ) + plugin.on_invocation_end(_end(operations=_ops(_step("s")))) + + assert observed, "the displaced record's finalizer must have run" + assert all(observed), ( + "a displaced record was released while this thread still held the " + "plugin's lock, which is what deadlocks two invocations whose finalizers " + "reach each other" + ) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + def _free_for_another_thread(lock: Any) -> bool: """Report whether a lock is unheld, as seen from a thread that never took it. From 3014e8ad1aea186eb7ce29c1397bd345b83b4d66 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 21:35:48 -0700 Subject: [PATCH 24/28] test(insight): make the finalizer probe deterministic The new test used the capturing exporter, which retains every record it is given. If the export worker exported the first record before it was displaced, the probe stayed reachable from the exporter and no finalizer ran -- so the test decided on a race it did not control, and lost it in CI while winning locally. It now uses an exporter that keeps nothing, and collects before asserting. --- .../tests/test_plugin.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 2b40dc17..6d9fa0f1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -21,6 +21,7 @@ import asyncio import datetime +import gc import itertools import threading import time @@ -1249,7 +1250,6 @@ def test_a_displaced_records_finalizer_runs_with_no_plugin_lock_held(): # each other at once hang both invocations, which the reentrant lock does not # help with: it only covers re-entry on the same instance. The hook frame # therefore releases displaced records after every lock is dropped. - exporter = ConcurrentCaptureExporter() observed: list[bool] = [] holder: dict[str, Any] = {} @@ -1263,11 +1263,31 @@ def __del__(self) -> None: # its owner acquire it again. observed.append(not plugin._lock._is_owned()) + class _DiscardingExporter: + """Keeps no record, so the displaced one is the last reference to a probe. + + An exporter that retained records would decide this test by a race: if the + worker exported the first record before it was displaced, the probe stays + reachable from the exporter and no finalizer runs. + """ + + max_record_size_bytes = None + exports = 0 + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + type(self).exports += 1 + + def flush(self) -> None: + pass + # A transform returning a fresh object per emit is what makes the record the # only reference to it, so releasing the displaced record is what collects it. factory = workflow_insight( WorkflowInsightConfig( - exporters=[exporter], + exporters=[_DiscardingExporter()], emit_mode="on-change", content=ContentConfig(input=lambda value: _FinalizerProbe()), ) @@ -1275,8 +1295,8 @@ def __del__(self) -> None: plugin = factory.create_plugin(_start(operations={})) holder["plugin"] = plugin - # The first emit queues a record holding a probe; the next displaces it while - # the worker is still parked, so the release happens on this thread. + # The first emit queues a record holding a probe; the next displaces it, and + # the hook frame releases it on this thread. plugin.on_invocation_start(_start(operations={})) plugin.on_operation_change( OperationChangeInfo( @@ -1286,6 +1306,7 @@ def __del__(self) -> None: ) ) plugin.on_invocation_end(_end(operations=_ops(_step("s")))) + gc.collect() assert observed, "the displaced record's finalizer must have run" assert all(observed), ( From fb4c010bf5e8435f07e5c07eed7234783ea16cb9 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 21:44:48 -0700 Subject: [PATCH 25/28] fix(plugin): keep factory diagnostics out of customer code Naming a failing factory for the log read the factory itself, so a factory whose __getattr__ or __getattribute__ raises made the containment raise: a create_plugin failure that should have been logged and skipped failed the invocation instead. The name now comes from the factory's type, which no instance attribute hook can intercept, and the whole lookup is wrapped, because a name is never worth failing a hook for. The scheduler's disabled path also handed the newly built record straight to the caller's local, so it was released when _emit returned -- still inside the plugin's lock, which is the finalizer hazard the previous commit closed for displaced records. It is now returned for the hook frame to release. --- .../_export_scheduler.py | 5 ++- .../plugin.py | 21 +++++++++- .../tests/plugin_test.py | 42 +++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 4de1f930..46fa6a1a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -165,7 +165,10 @@ def schedule(self, execution: _ExportState, record: dict[str, Any]) -> list[Any] start_error: Exception | None = None with self._condition: if self._disabled: - return [] + # The record is handed back rather than dropped here for the same + # reason a displaced one is: this runs inside the calling plugin's + # lock, and releasing the record can run a customer finalizer. + return [record] self._seq += 1 execution.scheduled_seq = self._seq displaced = execution.pending_record 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 874d27d5..420002e0 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 @@ -534,8 +534,25 @@ def create_plugin( def _factory_name(factory: object) -> str: - """Best available name for a factory, for log messages.""" - return getattr(factory, "__qualname__", None) or type(factory).__name__ + """Best available name for a factory, for log messages. + + Read from the factory's *type*, not from the factory. This runs while a + factory failure is being contained, and a ``getattr`` on the instance would + call a custom ``__getattr__`` or ``__getattribute__`` -- so a factory whose + attribute hook raises would make the containment itself raise, turning a + contained plugin failure into a failed execution. A class registered directly + as a factory is read through the class object, which carries its own + ``__qualname__``. + + Wrapped as well, because diagnostics must not be the thing that fails: a name + that cannot be produced is reported as unavailable rather than raised. + """ + try: + if isinstance(factory, type): + return factory.__qualname__ + return type(factory).__qualname__ + except BaseException: # noqa: BLE001 - a name is never worth failing a hook for + return "" # Raised out of plugin code, these three are not reports of a plugin defect but diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 200af119..517eded2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -648,6 +648,32 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: self.assertEqual(built[0].calls, ["invocation_start:req-1"]) self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + def test_a_hostile_factory_attribute_hook_does_not_escape_containment(self): + """Naming a failing factory must not be what fails the invocation. + + ``_factory_name`` runs while a factory failure is being contained, so it + reads the factory's type rather than the factory: an instance + ``__getattr__`` belongs to customer code and can raise. + """ + surviving = _TrackingPlugin() + executor = PluginExecutor( + plugins=[_HostileAttributeFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("_HostileAttributeFactory", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + def test_a_factory_raising_a_group_with_a_control_exception_propagates(self): """A group carrying a control exception is not contained. @@ -2505,6 +2531,22 @@ def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlug raise self._group +class _HostileAttributeFactory: + """Factory that fails, and whose attribute hook fails too. + + The second failure is the point: naming a failing factory for the log must not + reach customer code that raises, or containment raises instead of containing. + """ + + def __getattr__(self, name: str) -> object: + msg = f"attribute hook refuses {name}" + raise RuntimeError(msg) + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + msg = "factory boom" + raise RuntimeError(msg) + + class _CancellingPlugin(DurableInstrumentationPlugin): """Plugin whose hook raises outside the ``Exception`` hierarchy.""" From 3824d1b9ec6d003131b511c7de96ac90098f6020 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 18 Sep 2026 22:05:49 -0700 Subject: [PATCH 26/28] fix(plugin): name the invalid return type safely too The invalid-return message read type(plugin).__qualname__ directly, which a metaclass attribute hook can intercept, in the same containment path the factory name was just moved out of. Both now go through one helper that reads the type and cannot raise. --- .../aws_durable_execution_sdk_python/plugin.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) 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 420002e0..1a1f4692 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 @@ -533,6 +533,20 @@ def create_plugin( ... +def _type_name(value: object) -> str: + """Best available name for a value's type, for log messages. + + Reads the type rather than the value, so no instance ``__getattr__`` or + ``__getattribute__`` runs, and wraps the lookup, because these names are built + while a plugin failure is being contained: a name that raises would turn a + contained failure into a failed execution. + """ + try: + return type(value).__qualname__ + except BaseException: # noqa: BLE001 - a name is never worth failing a hook for + return "" + + def _factory_name(factory: object) -> str: """Best available name for a factory, for log messages. @@ -550,9 +564,9 @@ def _factory_name(factory: object) -> str: try: if isinstance(factory, type): return factory.__qualname__ - return type(factory).__qualname__ except BaseException: # noqa: BLE001 - a name is never worth failing a hook for return "" + return _type_name(factory) # Raised out of plugin code, these three are not reports of a plugin defect but @@ -711,7 +725,7 @@ def _create_plugins(self, info: InvocationStartInfo) -> None: "Plugin factory %s returned %s, which is not a " "DurableInstrumentationPlugin; plugin ignored", _factory_name(factory), - type(plugin).__qualname__, + _type_name(plugin), ) continue plugins.append(plugin) From 698693fbf64684615593d580771d4e57a724e67c Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 19 Sep 2026 07:46:09 -0700 Subject: [PATCH 27/28] fix(plugin): name the rejected class, not its metaclass The registration diagnostic named type(value), which is the metaclass when the entry is a class. Both likely migration mistakes are classes: plugins=[MyPlugin] is the previous major's shape, and plugins=[MyPluginFactory] is this major's shape with the parentheses left off. Both reported builtins.type, which identifies neither. A class and a function are now named by their own qualified name, and the kind is stated -- "the class", "the function", "an instance of" -- because a factory class and an instance of it share one name and passing the class is itself a rejected shape. --- .../plugin_discovery.py | 48 +++++++++++-- .../tests/plugin_discovery_test.py | 70 +++++++++++++++++-- 2 files changed, 108 insertions(+), 10 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index 9ff7290d..281e7275 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -54,9 +54,42 @@ def _distribution_name(entry_point: metadata.EntryPoint) -> str: return distribution.metadata.get("Name", "unknown distribution") -def _qualified_type_name(value: object) -> str: - value_type = type(value) - return f"{value_type.__module__}.{value_type.__qualname__}" +def _module_qualified_name(named: object) -> str: + """Join a class's or function's module and qualified name. + + Both attributes are read with a default, because a rejected entry can be any + object at all. A missing name must not replace the configuration error with an + ``AttributeError``. + """ + module = getattr(named, "__module__", None) or "unknown module" + qualified = getattr(named, "__qualname__", None) or getattr(named, "__name__", None) + return f"{module}.{qualified or ''}" + + +def _describe_value(value: object) -> str: + """Describe a rejected registration entry so the caller can identify it. + + Two kinds of value are named by themselves rather than by their type. A + class's type is its metaclass, which is ``builtins.type`` for an ordinary + class. A function's type is ``builtins.function``. Neither of those two names + says which value was passed. + + Both of the likeliest migration mistakes are classes: ``plugins=[MyPlugin]`` + is the shape the previous major accepted, and ``plugins=[MyPluginFactory]`` is + this major's shape with the parentheses left off. Naming the type would report + ``builtins.type`` for both. So a class and a function are named directly, and + every other value is named by its type. + + The kind is named alongside the name -- "the class", "the function", "an + instance of" -- because a factory class and an instance of that factory class + share one qualified name, and passing the class where an instance is required + is itself one of the rejected shapes. + """ + if isinstance(value, type): + return f"the class {_module_qualified_name(value)}" + if inspect.isroutine(value): + return f"the function {_module_qualified_name(value)}" + return f"an instance of {_module_qualified_name(type(value))}" def _is_plugin_factory(value: object) -> bool: @@ -182,7 +215,7 @@ def _load_factory( "resolve to a plugin factory -- an object with a " "create_plugin(info) method returning a " "DurableInstrumentationPlugin -- but resolved to " - f"{_qualified_type_name(factory)}. Name the factory instance, not a " + f"{_describe_value(factory)}. Name the factory instance, not a " "plugin, not a plugin class, and not the factory class." ) @@ -201,7 +234,10 @@ def _validate_explicit_factories( for the lifetime of the function, and nothing fails. Raising here converts that into one configuration failure while the handler is being initialized. The position is named because a caller passing several entries cannot - otherwise tell which one is wrong. + otherwise tell which one is wrong. The entry itself is named too, by + :func:`_describe_value`, which names a class and a function directly rather + than by type: the type of a class is ``builtins.type``, and that would + identify no particular class. A plugin *class* is rejected, and so is any bare callable. Both were accepted while the registration type was ``Callable``: a lambda satisfied it @@ -223,7 +259,7 @@ def _validate_explicit_factories( f"Durable instrumentation plugin at plugins[{index}] must be a " "plugin factory -- an object with a create_plugin(info) method " "returning a DurableInstrumentationPlugin -- but is " - f"{_qualified_type_name(factory)}. Pass a factory rather than a " + f"{_describe_value(factory)}. Pass a factory rather than a " "plugin, a plugin class, or a plain callable, and pass a factory " "instance rather than the factory class, for example " "plugins=[MyPluginFactory(exporter)]." diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index 2aac4ba5..83aac8ec 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -403,18 +403,21 @@ def test_discovery_names_unknown_distribution_in_load_failure() -> None: @pytest.mark.parametrize( - ("resolved_value", "expected_type_name"), + ("resolved_value", "expected_description"), [ (_PluginA(), "_PluginA"), (object(), "builtins.object"), ("not-a-factory", "builtins.str"), (None, "builtins.NoneType"), - (lambda info: _PluginA(), "builtins.function"), + # A function is named by its own qualified name, not by its type. Its type + # is ``builtins.function`` for every function ever written, so naming the + # type would identify no particular one. + (lambda info: _PluginA(), ""), ], ) def test_discovery_rejects_entry_point_without_create_plugin( resolved_value: object, - expected_type_name: str, + expected_description: str, ) -> None: """A plugin *instance* at the entry point is now the common mistake. @@ -438,7 +441,7 @@ def test_discovery_rejects_entry_point_without_create_plugin( assert "must resolve to a plugin factory" in str(error.value) assert "create_plugin(info) method" in str(error.value) - assert expected_type_name in str(error.value) + assert expected_description in str(error.value) def test_discovery_rejects_a_plugin_class_at_the_entry_point() -> None: @@ -483,6 +486,9 @@ def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: assert "plugins[1]" in str(error.value) assert "must be a plugin factory" in str(error.value) assert "_PluginB" in str(error.value) + # An instance and the class it was built from share one qualified name, so the + # message states which of the two was passed. + assert "an instance of" in str(error.value) @pytest.mark.parametrize( @@ -653,6 +659,62 @@ def test_discovery_rejects_a_factory_class_at_the_entry_point( assert "not the factory class" in str(error.value) +@pytest.mark.parametrize( + "rejected_class", + [_PluginA, _PluginAFactory], + ids=["plugin-class", "factory-class"], +) +def test_a_rejected_class_is_named_by_itself_not_by_its_metaclass( + rejected_class: type, +) -> None: + """The message has to name the class that was passed. + + The type of a class is its metaclass, which is ``builtins.type`` for both + classes here. Both are rejected shapes a caller reaches by accident: + ``plugins=[MyPlugin]`` is what the previous major accepted, and + ``plugins=[MyPluginFactory]`` is this major's shape with the parentheses left + off. Naming the type would report ``builtins.type`` for either one and + distinguish neither. So the class is named directly, and the message says it + was a class rather than an instance. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([rejected_class], environment={}) # type: ignore[list-item] + + message = str(error.value) + qualified = f"{rejected_class.__module__}.{rejected_class.__qualname__}" + assert f"the class {qualified}" in message + assert "builtins.type" not in message + + +@pytest.mark.parametrize( + "rejected_class", + [_PluginA, _PluginAFactory], + ids=["plugin-class", "factory-class"], +) +def test_a_rejected_class_at_the_entry_point_is_named_by_itself( + rejected_class: type, +) -> None: + """The entry-point path applies the same naming rule.""" + entry_point = _FakeEntryPoint("a", rejected_class) + + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ), + pytest.raises(PluginLoadError) as error, + ): + load_configured_plugins( + None, + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + message = str(error.value) + qualified = f"{rejected_class.__module__}.{rejected_class.__qualname__}" + assert f"the class {qualified}" in message + assert "builtins.type" not in message + + def test_explicit_class_holding_a_callable_create_plugin_is_accepted() -> None: """A class attribute holding a callable takes no implicit first argument. From 4face578875c5058d1c57bfa21491450667f8249 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Sat, 19 Sep 2026 07:46:16 -0700 Subject: [PATCH 28/28] fix(insight): make the plugin factory type public workflow_insight() declared _WorkflowInsightFactory as its return type. That name is private and absent from __all__, so a py.typed consumer could not annotate the value it holds without importing a private symbol. The sibling OTel package, changed in the same release, exports InvocationOtelPluginFactory and ExecutionOtelPluginFactory. The class is now WorkflowInsightPluginFactory and is exported from the package root. The _ExportState mixin also set five non-underscored fields, and the mixin lands on WorkflowInsightPlugin, which the package exports. Those fields belong to the export scheduler and are guarded by its lock, so a public name advertised scheduler bookkeeping as part of the plugin's API. They are now underscore-prefixed; _ExportScheduler is declared in the same module, so it still reads them directly. --- .../README.md | 3 +- .../__init__.py | 2 + .../_export_scheduler.py | 59 +++++++++++-------- .../plugin.py | 12 ++-- .../tests/test_export_scheduler.py | 6 +- .../tests/test_plugin.py | 41 +++++++++++++ 6 files changed, 90 insertions(+), 33 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index f3c2847b..f57a8bbb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -46,7 +46,8 @@ def handler(event, context): `plugins` argument takes: the SDK calls its `create_plugin` once per invocation to build that invocation's plugin instance. The factory holds the resolved configuration and the exporters, so configuration is per handler while record -state is per invocation. +state is per invocation. Its type is `WorkflowInsightPluginFactory`, exported from +the package root for annotating a value you hold. With no exporter configured, records are written to the function's own CloudWatch log group as single JSON lines (the `LambdaLogExporter` default), diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py index 488c2ada..3cf30c9a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py @@ -33,6 +33,7 @@ ) from aws_durable_execution_sdk_python_insight.plugin import ( WorkflowInsightPlugin, + WorkflowInsightPluginFactory, workflow_insight, ) from aws_durable_execution_sdk_python_insight.truncation import truncate_record @@ -77,6 +78,7 @@ "SQSExporter", "WorkflowInsightConfig", "WorkflowInsightPlugin", + "WorkflowInsightPluginFactory", "apply_operations_format", "build_operations_by_name", "truncate_record", diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 46fa6a1a..bea39fcb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -62,27 +62,34 @@ class _ExportState: Every field here is guarded by ``_ExportScheduler._condition``. They belong to the scheduler: nothing outside it reads or writes them, and it never touches the hook-facing state on the same object. + + Every field is underscore-prefixed, because the mixin's fields land on + :class:`WorkflowInsightPlugin`, which the package exports. A public name there + would advertise scheduler bookkeeping as part of the plugin's API. The + scheduler reads these names directly, which is why they are not name-mangled: + :class:`_ExportScheduler` is declared in this module, so this class's private + fields are within its own module's reach. """ def __init__(self) -> None: # Newest sequence number the scheduler assigned to this execution. - self.scheduled_seq = 0 + self._scheduled_seq = 0 # This execution's latest record, waiting for the export worker, or None # when nothing of its own is outstanding. A repeat emission replaces it, # which is what per-execution coalescing means; the scheduler's queue # holds this object exactly while this field is set. - self.pending_record: dict[str, Any] | None = None + self._pending_record: dict[str, Any] | None = None # Newest sequence number already handed to every exporter. - self.exported_seq = 0 + self._exported_seq = 0 # Value of the scheduler's export counter when that export finished, so # a waiter can tell whether a completed flush covered its own record. - self.exported_at = 0 + self._exported_at = 0 # drain() calls currently parked on this execution. Nothing depends on # it: a waiter holds this object directly, so the bookkeeping it waits # on can no longer be reclaimed from under it. It is kept because it is # the only way to observe that a drain really parked rather than raced # past. - self.waiters = 0 + self._waiters = 0 # What a caller has to release once it is back outside the lock: the records the @@ -101,7 +108,7 @@ def __init__(self, exporters: list[InsightExporter]) -> None: # set, keyed by the execution object itself. A repeat schedule for an # execution replaces the record the object carries and keeps the object's # position, so coalescing never lets one execution jump the queue. An - # execution is in here exactly while its `pending_record` is set. + # execution is in here exactly while its `_pending_record` is set. self._pending: dict[_ExportState, None] = {} self._seq = 0 self._export_count = 0 @@ -170,9 +177,9 @@ def schedule(self, execution: _ExportState, record: dict[str, Any]) -> list[Any] # lock, and releasing the record can run a customer finalizer. return [record] self._seq += 1 - execution.scheduled_seq = self._seq - displaced = execution.pending_record - execution.pending_record = record + execution._scheduled_seq = self._seq + displaced = execution._pending_record + execution._pending_record = record # Re-queuing an execution that is already queued is a no-op that # keeps its arrival position. self._pending[execution] = None @@ -247,9 +254,9 @@ def drain(self, execution: _ExportState) -> None: with self._condition: if self._disabled: return - execution.waiters += 1 + execution._waiters += 1 try: - want_seq = execution.scheduled_seq + want_seq = execution._scheduled_seq # A drain always flushes, so require a flush that covers every # export completed before this call as well as our own. want_flush = self._export_count @@ -261,11 +268,11 @@ def drain(self, execution: _ExportState) -> None: # Export counter value a flush has to cover to release us: # our own record's export plus everything already exported # when this call started. Recomputed every pass, because - # execution.exported_at only becomes ours once our record is + # execution._exported_at only becomes ours once our record is # out. - need = max(execution.exported_at, want_flush) + need = max(execution._exported_at, want_flush) if ( - execution.exported_seq >= want_seq + execution._exported_seq >= want_seq and self._flushed_through >= need and self._flushes_completed > want_flushes ): @@ -296,7 +303,7 @@ def drain(self, execution: _ExportState) -> None: self._condition.notify_all() self._condition.wait() finally: - execution.waiters -= 1 + execution._waiters -= 1 del failed_pending if start_error is not None: _logger.warning( @@ -361,9 +368,11 @@ def _disable_locked(self) -> _Dropped: """ self._disabled = True self._worker = None - dropped = [(execution, execution.pending_record) for execution in self._pending] + dropped = [ + (execution, execution._pending_record) for execution in self._pending + ] for execution, _ in dropped: - execution.pending_record = None + execution._pending_record = None self._pending = {} self._flush_requested = False self._flush_barrier = 0 @@ -390,7 +399,7 @@ def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: def _blocking_pending_locked(self) -> bool: """True while a record scheduled at or before the flush barrier is pending.""" barrier = self._flush_barrier - return any(execution.scheduled_seq <= barrier for execution in self._pending) + return any(execution._scheduled_seq <= barrier for execution in self._pending) def _run(self) -> None: # The worker slot must be empty whenever no worker is running, or @@ -470,9 +479,9 @@ def _run_loop(self) -> None: if self._pending: execution = next(iter(self._pending)) del self._pending[execution] - record = execution.pending_record - execution.pending_record = None - seq = execution.scheduled_seq + record = execution._pending_record + execution._pending_record = None + seq = execution._scheduled_seq break self._condition.wait() @@ -481,7 +490,7 @@ def _run_loop(self) -> None: # Taking the record consumed this execution's pending slot, so # nothing will ever export that snapshot again. The bookkeeping # must therefore advance whatever export() did: skip it and - # execution.exported_seq never reaches a waiter's want_seq, so a + # execution._exported_seq never reaches a waiter's want_seq, so a # drain parked on this execution is never released. Count the # attempt in a finally so that holds even if _export() raises. # @@ -503,9 +512,9 @@ def _run_loop(self) -> None: del record with self._condition: self._export_count += 1 - if seq > execution.exported_seq: - execution.exported_seq = seq - execution.exported_at = self._export_count + if seq > execution._exported_seq: + execution._exported_seq = seq + execution._exported_at = self._export_count if exported: # This worker completed work, so any earlier worker # death was not the start of a fault the work diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index a46c947d..d56bb5d4 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -231,7 +231,7 @@ class WorkflowInsightPlugin(DurableInstrumentationPlugin, _ExportState): """ def __init__( - self, shared: _WorkflowInsightFactory, info: InvocationStartInfo + self, shared: WorkflowInsightPluginFactory, info: InvocationStartInfo ) -> None: _ExportState.__init__(self) self._shared = shared @@ -665,9 +665,13 @@ def _emit( _hook_frames.releases.extend(self._shared._scheduler.schedule(self, record)) -class _WorkflowInsightFactory: +class WorkflowInsightPluginFactory: """The handler-lifetime half of the plugin: what is NOT per-execution. + Built by :func:`workflow_insight`, which is the supported way to obtain one. + The class is public because it is the declared return type of that function, + and a ``py.typed`` consumer must be able to name the type it holds. + Satisfies the SDK's ``DurableInstrumentationPluginFactory`` -- its ``create_plugin`` is called with an ``InvocationStartInfo`` and returns the plugin instance for that invocation. Everything it holds is either immutable @@ -722,11 +726,11 @@ def create_plugin(self, info: InvocationStartInfo) -> WorkflowInsightPlugin: return WorkflowInsightPlugin(self, info) -def workflow_insight(config: WorkflowInsightConfig) -> _WorkflowInsightFactory: +def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPluginFactory: """Create a Workflow Insight plugin factory. Mirrors the JS ``workflowInsight()``. Pass the result straight to ``@durable_execution(plugins=[...])``: the SDK calls its ``create_plugin`` once per invocation to build that invocation's plugin instance. """ - return _WorkflowInsightFactory(config) + return WorkflowInsightPluginFactory(config) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 7f9f454d..7cd56b9d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -507,7 +507,7 @@ def _scheduler_is_empty(scheduler: _ExportScheduler) -> bool: def _drain_waiters(scheduler: _ArnScheduler, arn: str) -> int: """How many drain() calls are currently parked on this execution.""" with scheduler._condition: - return scheduler._execution(arn).waiters + return scheduler._execution(arn)._waiters def test_drain_stays_parked_until_a_flush_covering_its_record_completes() -> None: @@ -672,7 +672,7 @@ def both_parked() -> bool: if second is None: return False with scheduler._condition: - return first.waiters == 1 and second.waiters == 1 + return first._waiters == 1 and second._waiters == 1 assert _wait_until(both_parked), "a drain raced past the flush it needs" with returned_lock: @@ -740,7 +740,7 @@ def fail_start(self) -> None: # noqa: ARG001 # bookkeeping that used to need clearing in a second map is on these # objects now, so the queue and the records are one thing to release. assert all( - execution.pending_record is None + execution._pending_record is None for execution in scheduler.executions.values() ) assert scheduler._flush_requested is False diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 6d9fa0f1..b55c59cc 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -50,6 +50,7 @@ LambdaLogExporter, OperationOverride, WorkflowInsightConfig, + WorkflowInsightPluginFactory, workflow_insight, ) from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportState @@ -186,6 +187,46 @@ def _run( return plugin +# -- public surface ---------------------------------------------------------- + + +def test_the_factory_type_is_public_and_re_exported(): + """A consumer must be able to name the type ``workflow_insight()`` returns. + + The package ships ``py.typed``, so a consumer annotating the value it holds + needs the name. A private name would force an import from a private module. + The sibling OTel package exports ``InvocationOtelPluginFactory`` and + ``ExecutionOtelPluginFactory`` for the same reason, so this keeps the two + plugin packages consistent. + """ + import aws_durable_execution_sdk_python_insight as pkg + + factory = workflow_insight(WorkflowInsightConfig(exporters=[CaptureExporter()])) + + assert type(factory) is WorkflowInsightPluginFactory + assert not WorkflowInsightPluginFactory.__name__.startswith("_") + assert "WorkflowInsightPluginFactory" in pkg.__all__ + assert pkg.WorkflowInsightPluginFactory is WorkflowInsightPluginFactory + + +def test_the_plugin_exposes_no_public_attributes(): + """Export bookkeeping must not become part of the plugin's public surface. + + ``WorkflowInsightPlugin`` is exported from the package root, and it mixes in + :class:`_ExportState`, so every field that mixin sets lands on the exported + class. Those fields belong to the export scheduler, which owns them under its + own lock. A public name among them would advertise scheduler bookkeeping as + something a consumer may read or set. So every field on the instance is + underscore-prefixed. + """ + exporter = CaptureExporter() + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin = _invocation(factory, _start()) + + public = sorted(name for name in vars(plugin) if not name.startswith("_")) + assert public == [] + + # -- existing record-building coverage ---------------------------------------