From fb8af739a92f7fd8070d928029150fe4e5c9cc61 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Wed, 16 Sep 2026 17:14:49 -0700 Subject: [PATCH] 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())