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())