diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index ed47c053..1c3b5f24 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -55,10 +55,21 @@ and `top-level` vs `full-tree` operation detail all mirror the JS plugin. Behavior is validated cross-SDK by the `insight` conformance suite (`aws-durable-execution-conformance-tests-insight`). -> **Note (`on-change` emission).** In `on-change` mode, exporter calls currently -> run synchronously on the SDK checkpoint path, so a slow exporter can delay -> workflow progress. Asynchronous scheduling/coalescing is deferred and tracked -> in [issue #687](https://github.com/aws/aws-durable-execution-sdk-python/issues/687). +> **Note (asynchronous export).** Exporter work — per-exporter copy, rendering, +> truncation, `export()` and `flush()` — runs on a background daemon worker per +> exporter, never on the SDK checkpoint path, so a slow exporter does not delay +> workflow progress. Because each configured exporter is driven by its own +> single background worker, each exporter object may belong to only one live +> `WorkflowInsightPlugin`: listing it twice or sharing it across plugin instances +> raises `ValueError`. Separate instances of the same exporter class (e.g. two +> `S3Exporter`s for different buckets) are fine. A blocked lane retains at most +> one record in flight and one latest pending snapshot. Rapid cumulative +> snapshots are coalesced, so a lane may skip intermediate `on-change` records; +> the terminal record is always delivered under normal completion. At invocation +> end the plugin drains and flushes the touched exporters under a single shared +> deadline +> (`WorkflowInsightConfig.export_timeout_seconds`, default `5.0`); on timeout the +> workflow response is returned and record delivery degrades to best-effort. ## Requirements 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 new file mode 100644 index 00000000..f5a94d5d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Asynchronous export scheduler for the Workflow Insight plugin. + +Exporter work runs outside the SDK checkpoint thread. Each exporter has one +lazy daemon worker that serializes copying, rendering, truncation, export, and +flush calls. + +Each exporter lane: + +* Keeps at most one record in flight and one latest pending snapshot. A newer + pending snapshot replaces the older one. +* Uses a lane-wide flush barrier at invocation end. All barriers share one + timeout; timed-out barriers are removed without replacing a blocked worker. +* Stops its worker after an invocation has drained and the lane becomes idle. +""" + +from __future__ import annotations + +import copy +import logging +import threading +import time +from collections import deque +from typing import Any + +from aws_durable_execution_sdk_python_insight.truncation import truncate_record +from aws_durable_execution_sdk_python_insight.types import InsightExporter + + +_logger = logging.getLogger("aws_durable_execution_sdk_python_insight") + + +def _copy_record_containers(record: dict[str, Any]) -> dict[str, Any]: + """Copy built-in containers while treating custom values as opaque leaves.""" + memo: dict[int, Any] = {} + seen: set[int] = set() + stack: list[Any] = [record] + while stack: + item = stack.pop() + identity = id(item) + if identity in seen: + continue + seen.add(identity) + if type(item) is dict: + stack.extend(item.keys()) + stack.extend(item.values()) + elif type(item) in {list, tuple, set, frozenset, deque}: + stack.extend(item) + else: + memo[identity] = item + return copy.deepcopy(record, memo) + + +# Queue entry kinds. +_RECORD = "record" +_FLUSH = "flush" + + +class _FlushBarrier: + """A one-shot flush marker the invocation-end thread waits on. + + The worker completes the barrier after it has flushed (or skipped a cancelled + barrier). ``canceled`` is set by the waiter when the shared timeout elapses so + a later, still-blocked worker skips the now-pointless flush. + """ + + __slots__ = ("_event", "canceled", "failed") + + def __init__(self) -> None: + self._event = threading.Event() + self.canceled = False + self.failed = False + + def complete(self) -> None: + self._event.set() + + def wait(self, timeout: float) -> bool: + return self._event.wait(timeout if timeout > 0 else 0) + + def is_done(self) -> bool: + return self._event.is_set() + + +class _ExporterLane: + """A single exporter's serial worker lane. + + All mutable state is guarded by ``_cond``. The worker is the only consumer of + the queue; scheduling threads are producers that wake it via ``notify``. + """ + + def __init__(self, exporter: InsightExporter) -> None: + self._exporter = exporter + # Explicit non-reentrant Lock rather than Condition()'s default RLock: + # the lane never re-acquires ``_cond`` while already holding it (worker + # I/O -- export/flush -- runs outside the lock and no locked helper + # re-enters), so recursion support is unnecessary. A plain Lock also + # makes any accidental recursive acquisition fail loudly instead of + # silently succeeding. + self._cond = threading.Condition(threading.Lock()) + # Ordered work list: entries are (_RECORD, None) or (_FLUSH, barrier). + self._queue: deque[tuple[str, Any]] = deque() + # At most one record waits behind the in-flight export. Replacing this + # snapshot moves its queue token to the newest scheduling position. + self._pending: dict[str, Any] | None = None + self._stop_when_idle = False + self._worker: threading.Thread | None = None + self._disabled = False + + # -- producer API (checkpoint / invocation-end threads) ------------------- + + def schedule(self, record: dict[str, Any]) -> None: + with self._cond: + if self._disabled: + return + self._stop_when_idle = False + if self._pending is None: + self._queue.append((_RECORD, None)) + else: + self._move_record_token_to_back() + self._pending = record + self._ensure_worker_locked() + self._cond.notify() + + def enqueue_flush(self) -> _FlushBarrier: + barrier = _FlushBarrier() + with self._cond: + if self._disabled: + barrier.canceled = True + barrier.failed = True + barrier.complete() + return barrier + self._queue.append((_FLUSH, barrier)) + self._ensure_worker_locked() + self._cond.notify() + return barrier + + def request_stop_when_idle(self) -> None: + with self._cond: + self._stop_when_idle = True + self._cond.notify() + + def cancel_flush(self, barrier: _FlushBarrier) -> None: + """Stop waiting while retaining one flush at the latest covered point.""" + with self._cond: + barrier.canceled = True + marker_indexes = { + index + for index, (kind, payload) in enumerate(self._queue) + if kind == _FLUSH and (payload is barrier or payload is None) + } + if not any( + kind == _FLUSH and payload is barrier for kind, payload in self._queue + ): + # The worker already owns this barrier. Do not erase a detached + # flush installed by a later invocation while this one was in flight. + return + # A flush covers every record before its queue position. Keep the + # rightmost canceled/detached marker so coalescing never narrows the + # set of records that will eventually be published. + rightmost = max(marker_indexes) + coalesced: deque[tuple[str, Any]] = deque() + for index, item in enumerate(self._queue): + if index == rightmost: + coalesced.append((_FLUSH, None)) + elif index not in marker_indexes: + coalesced.append(item) + self._queue = coalesced + barrier.complete() + + # -- queue bookkeeping (must hold ``_cond``) ------------------------------ + + def _move_record_token_to_back(self) -> None: + for index, (kind, _) in enumerate(self._queue): + if kind == _RECORD: + del self._queue[index] + self._queue.append((_RECORD, None)) + return + + def _disable_locked(self, exc: Exception) -> None: + self._disabled = True + self._worker = None + self._pending = None + for kind, payload in self._queue: + if kind == _FLUSH and payload is not None: + barrier: _FlushBarrier = payload + barrier.canceled = True + barrier.failed = True + barrier.complete() + self._queue.clear() + _logger.warning( + "workflow-insight: could not start worker for exporter %s; " + "disabling this lane: %s", + type(self._exporter).__name__, + exc, + ) + + def _ensure_worker_locked(self) -> None: + # Never create a replacement while a prior worker is alive (a blocked + # worker keeps ``_worker`` non-None). A worker that exits cleanly nulls + # ``_worker`` under the lock before returning, so this check is a + # race-free "start iff there is no live worker". + if self._disabled: + return + if self._worker is None or not self._worker.is_alive(): + worker = threading.Thread( + target=self._run_worker, + name=f"workflow-insight-export-{id(self)}", + daemon=True, + ) + self._worker = worker + try: + worker.start() + except Exception as exc: # noqa: BLE001 - instrumentation must not break hooks + self._disable_locked(exc) + + # -- worker (single daemon thread) --------------------------------------- + + def _run_worker(self) -> None: + while True: + with self._cond: + while not self._queue and not self._stop_when_idle: + self._cond.wait() + if not self._queue and self._stop_when_idle: + # Idle stop: null ``_worker`` under the lock so a concurrent + # scheduler starts a fresh worker rather than assuming this + # one will pick the work up. + self._worker = None + return + kind, payload = self._queue.popleft() + record: dict[str, Any] | None = None + if kind == _RECORD: + record = self._pending + self._pending = None + if record is None: + continue + + if kind == _RECORD and record is not None: + self._export_one(record) + else: # _FLUSH + barrier: _FlushBarrier | None = payload + self._flush() + if barrier is not None: + barrier.complete() + + def _export_one(self, record: dict[str, Any]) -> None: + exporter = self._exporter + # Copy the record's built-in containers for lane isolation, but preserve + # custom values as opaque leaves for exporter-specific rendering. This + # keeps one lane's render/truncation mutations out of other lanes without + # requiring custom-renderable values to implement ``deepcopy``. + try: + local = _copy_record_containers(record) + except Exception as exc: # noqa: BLE001 - export remains best-effort + _logger.warning( + "workflow-insight: record container copy failed for exporter %s; " + "using the original record without lane isolation: %s", + type(exporter).__name__, + exc, + ) + local = record + try: + shaped = truncate_record( + local, exporter.max_record_size_bytes, exporter.render + ) + except Exception as exc: # noqa: BLE001 - render/truncation is best-effort + _logger.warning( + "workflow-insight: render/truncation failed for exporter %s: %s", + type(exporter).__name__, + exc, + ) + return + try: + exporter.export(shaped) + except Exception as exc: # noqa: BLE001 - one export must not break the lane + _logger.warning( + "workflow-insight: exporter %s export failed: %s", + type(exporter).__name__, + exc, + ) + + def _flush(self) -> None: + try: + self._exporter.flush() + except Exception as exc: # noqa: BLE001 - a failing flush completes the barrier + _logger.warning( + "workflow-insight: exporter %s flush failed: %s", + type(self._exporter).__name__, + exc, + ) + + # -- test / introspection helpers ---------------------------------------- + + def _worker_alive(self) -> bool: + with self._cond: + return self._worker is not None and self._worker.is_alive() + + def _pending_count(self) -> int: + with self._cond: + return int(self._pending is not None) + + def _queue_len(self) -> int: + with self._cond: + return len(self._queue) + + def _queued_flush_count(self) -> int: + with self._cond: + return sum(1 for kind, _ in self._queue if kind == _FLUSH) + + +class _ExportScheduler: + """Owns one :class:`_ExporterLane` per exporter and fans records out to them.""" + + def __init__(self, exporters: list[InsightExporter]) -> None: + self._lanes = [_ExporterLane(exporter) for exporter in exporters] + + def schedule(self, _execution_arn: str, record: dict[str, Any]) -> None: + """Fan a canonical record out to every lane. Returns immediately.""" + for lane in self._lanes: + lane.schedule(record) + + def end_invocation(self, timeout_seconds: float) -> bool: + """Drain and flush every touched lane under one shared timeout. + + Enqueues a flush barrier per lane (after that lane's latest record), + waits for all barriers against a single deadline, then asks every worker + to stop once idle. Returns ``True`` if every barrier completed within the + deadline, ``False`` if delivery degraded to best-effort on timeout. + """ + barriers = [(lane, lane.enqueue_flush()) for lane in self._lanes] + deadline = time.monotonic() + timeout_seconds + degraded = False + for lane, barrier in barriers: + remaining = deadline - time.monotonic() + if not barrier.wait(remaining): + # Timed out: cancel this lane's barrier and pull its still-queued + # _FLUSH marker out now, so a stale barrier per invocation cannot + # accumulate behind a blocked worker. + lane.cancel_flush(barrier) + degraded = True + elif barrier.failed: + degraded = True + for lane in self._lanes: + lane.request_stop_when_idle() + if degraded: + _logger.warning( + "workflow-insight: export drain/flush exceeded %.3fs; " + "record delivery is best-effort for this invocation", + timeout_seconds, + ) + return not degraded 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 13f16a06..c49f1bdd 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 @@ -33,8 +33,8 @@ import datetime import json import math -import sys import threading +import weakref from typing import Any, Callable from aws_durable_execution_sdk_python.plugin import ( @@ -47,10 +47,10 @@ OperationType, ) +from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportScheduler from aws_durable_execution_sdk_python_insight.exporters.lambda_log_exporter import ( LambdaLogExporter, ) -from aws_durable_execution_sdk_python_insight.truncation import truncate_record from aws_durable_execution_sdk_python_insight.types import ( ContentConfig, EmitMode, @@ -58,6 +58,7 @@ OperationDetail, OperationOverride, WorkflowInsightConfig, + _validate_exporter_instances, ) @@ -71,6 +72,34 @@ InvocationStatus.RETRY: "RUNNING", } +_exporter_owner_lock = threading.Lock() +_exporter_owners: list[weakref.ReferenceType[Any]] = [] + + +def _claim_exporter_lanes(lanes: list[Any]) -> None: + """Give each exporter object to at most one live scheduler lane.""" + + live_owners: list[Any] = [] + with _exporter_owner_lock: + for lane_ref in _exporter_owners: + owner = lane_ref() + if owner is not None: + live_owners.append(owner) + _exporter_owners[:] = [weakref.ref(owner) for owner in live_owners] + for lane in lanes: + exporter = lane._exporter + for owner in live_owners: + if owner._exporter is exporter: + raise ValueError( + "the same exporter instance cannot be shared across " + "Workflow Insight plugin instances" + ) + # Callback-free weakrefs avoid lock re-entry during synchronous finalization. + # Dead entries are pruned at the start of every subsequent claim. + _exporter_owners.extend(weakref.ref(lane) for lane in lanes) + # Keep dereferenced lanes alive until the ownership lock has been released. + # Releasing the last lane reference can finalize a custom exporter synchronously. + def _parse_execution_arn(execution_arn: str) -> dict[str, str]: # arn::lambda:::function::/durable-execution// @@ -161,7 +190,7 @@ def _apply_result_override( class _ExecutionState: - __slots__ = ("start_time", "parsed_arn", "cached_input", "operations") + __slots__ = ("start_time", "parsed_arn", "cached_input", "operations", "scheduled") def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: self.start_time = start_time @@ -170,6 +199,10 @@ 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] = {} + # True once at least one record was scheduled for the current + # invocation; gates the invocation-end drain/flush so a no-op invocation + # (e.g. on-complete + non-terminal end) never touches a lane. + self.scheduled: bool = False class WorkflowInsightPlugin(DurableInstrumentationPlugin): @@ -199,14 +232,24 @@ def __init__(self, config: WorkflowInsightConfig) -> None: if ops is not None: for override in ops.overrides: self._overrides_by_name[override.operation_name] = override - # Default-exporter parity with the JS plugin: an omitted OR an explicitly - # empty exporter list falls back to the Lambda log exporter, so the - # plugin is never a silent no-op. A non-empty list is used verbatim. + # Snapshot and revalidate at use time. ``WorkflowInsightConfig`` is + # frozen, but its caller-owned list can still be mutated after + # construction; the exact snapshot used to create lanes must preserve + # the one-worker-per-distinct-instance invariant. + configured_exporters = list(config.exporters) + _validate_exporter_instances(configured_exporters) self._exporters: list[InsightExporter] = ( - list(config.exporters) if config.exporters else [LambdaLogExporter()] + configured_exporters if configured_exporters else [LambdaLogExporter()] ) + # One shared deadline (seconds) for the invocation-end drain + flush. + self._export_timeout = float(config.export_timeout_seconds) + # Off-thread export scheduler: one lazy daemon worker per exporter. All + # copy/render/truncation/export/flush runs there, never on the checkpoint + # hook thread. + self._scheduler = _ExportScheduler(self._exporters) self._state: dict[str, _ExecutionState] = {} self._lock = threading.Lock() + _claim_exporter_lanes(self._scheduler._lanes) # -- sampling / state ----------------------------------------------------- @@ -239,6 +282,21 @@ def _adopt_operations( with self._lock: state.operations = dict(operations) + def _mark_scheduled(self, state: _ExecutionState) -> None: + # Mutate ``scheduled`` under the same lock that guards every other field + # of the shared ``_ExecutionState``. The lock is released before any + # scheduler/exporter work runs (see ``_schedule_record``), so it never + # covers I/O and introduces no new lock ordering. + with self._lock: + state.scheduled = True + + def _was_scheduled(self, state: _ExecutionState) -> bool: + # Read ``scheduled`` under the lock, then act on the returned snapshot + # outside it -- the invocation-end drain must not hold the plugin lock + # while calling the scheduler. + with self._lock: + return state.scheduled + # -- hooks ---------------------------------------------------------------- def on_invocation_start(self, info: InvocationStartInfo) -> None: @@ -257,7 +315,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: # plugin instance never saw via per-operation hooks. self._adopt_operations(state, info.operations) if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( + self._schedule_record( arn, state, status="RUNNING", @@ -267,23 +325,30 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) def on_operation_change(self, info: OperationChangeInfo) -> None: + # Non-on-change modes never emit mid-invocation, so skip all work here + # (adopting the snapshot, sampling) -- the terminal record is rebuilt + # from the invocation-end snapshot. This keeps the checkpoint path free + # of Insight work outside on-change mode. + if self._emit_mode != EmitMode.ON_CHANGE: + return 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, - ) + # on-change mode schedules an updated RUNNING record on each change so + # mid-invocation progress is observable. The schedule call returns + # immediately -- rendering/export happens off the checkpoint thread -- so + # a slow exporter never blocks workflow progress. + self._schedule_record( + arn, + state, + status="RUNNING", + end_time=None, + output_raw=None, + error=None, + ) def on_invocation_end(self, info: InvocationEndInfo) -> None: arn = info.execution_arn @@ -311,10 +376,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: 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( + # must omit endTime/durationMs. Passing end_time=None makes + # _schedule_record drop both fields. Output and error likewise + # belong only to a terminal record. + self._schedule_record( arn, state, status=status, @@ -323,6 +388,14 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: error=info.error if is_terminal else None, ) + # Drain and flush the touched lanes under one shared timeout, but only if + # this invocation actually scheduled something. A no-op invocation (e.g. + # on-complete + non-terminal end) touches no lane and needs no flush, + # which also avoids spinning up an idle worker just to flush nothing. + # Read the flag under the lock, then call the scheduler outside it. + if self._was_scheduled(state): + self._scheduler.end_invocation(self._export_timeout) + # 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 @@ -373,7 +446,7 @@ def _build_operations( records.append(entry) return records - def _emit( + def _schedule_record( self, execution_arn: str, state: _ExecutionState, @@ -383,6 +456,31 @@ def _emit( output_raw: str | None, error: Any, ) -> None: + record = self._build_record( + execution_arn, + state, + status=status, + end_time=end_time, + output_raw=output_raw, + error=error, + ) + # Hand the canonical record to the scheduler; per-exporter copy, render, + # truncation, export and flush all run on the lane workers, never here. + self._scheduler.schedule(execution_arn, record) + # Set the flag under the plugin lock AFTER the scheduler call so the lock + # never covers scheduler work. + self._mark_scheduled(state) + + def _build_record( + self, + execution_arn: str, + state: _ExecutionState, + *, + status: str, + end_time: Any, + output_raw: str | None, + error: Any, + ) -> dict[str, Any]: arn = state.parsed_arn start_time = state.start_time duration = _duration_ms(start_time, end_time) @@ -434,22 +532,7 @@ def _emit( if error is not None: record["error"] = {"name": error.type, "message": error.message} record["operations"] = self._build_operations(operations) - - for exporter in self._exporters: - try: - shaped = truncate_record( - record, exporter.max_record_size_bytes, exporter.render - ) - exporter.export(shaped) - except Exception as exc: # noqa: BLE001 - one exporter must not break others / the execution - # NOTE (parity gap, same as JS Promise.allSettled): exporter - # failures are swallowed so instrumentation never breaks the - # execution. A silently broken exporter is indistinguishable - # from success; we at least log to stderr. - print( - f"[workflow-insight] exporter {type(exporter).__name__} failed: {exc}", - file=sys.stderr, - ) # noqa: T201 + return 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..c1c59600 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 @@ -11,9 +11,11 @@ from __future__ import annotations +import math +import threading from dataclasses import dataclass, field from enum import StrEnum -from typing import Any, Callable, Literal, Protocol +from typing import Any, Callable, Literal, Protocol, Sequence class EmitMode(StrEnum): @@ -64,6 +66,22 @@ def export(self, record: dict[str, Any]) -> None: ... # pragma: no cover def flush(self) -> None: ... # pragma: no cover +def _validate_exporter_instances(exporters: Sequence[InsightExporter]) -> None: + """Reject one exporter object appearing more than once by identity.""" + seen: list[InsightExporter] = [] + for exporter in exporters: + if any(exporter is other for other in seen): + raise ValueError( + "exporters must not contain the same exporter instance more " + f"than once ({type(exporter).__name__} appears multiple " + "times); each configured exporter runs on its own single " + "background worker, so one instance shared across lanes " + "would schedule duplicate, timing-dependent exports. Use " + "two distinct instances if you need two destinations." + ) + seen.append(exporter) + + @dataclass(frozen=True) class OperationOverride: """Per-operation override matched by ``operation_name``. @@ -106,6 +124,11 @@ class WorkflowInsightConfig: emit_mode: EmitMode | EmitModeInput | None = None operation_detail: OperationDetail | OperationDetailInput | None = None content: ContentConfig | None = None + # Single shared deadline (seconds) for the invocation-end drain + flush of + # every touched exporter lane. Mirrors the JS plugin's best-effort bound: on + # timeout the workflow response is returned and delivery degrades to + # best-effort. Must be a finite number greater than zero. + export_timeout_seconds: float = 5.0 def __post_init__(self) -> None: # Normalize accepted string inputs to enum members so the plugin always @@ -119,3 +142,34 @@ def __post_init__(self) -> None: object.__setattr__( self, "operation_detail", OperationDetail(self.operation_detail) ) + _validate_exporter_instances(self.exporters) + self._validate_export_timeout() + + def _validate_export_timeout(self) -> None: + # A finite, strictly-positive number within Event.wait's platform bound. + # Reject ``bool`` (an ``int`` subtype), values that overflow float + # conversion, NaN, infinity, zero, negatives, and values above + # ``threading.TIMEOUT_MAX``. Invalid input must fail at construction, + # before invocation-end cleanup can be interrupted by ``OverflowError``. + timeout = self.export_timeout_seconds + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): + raise ValueError( + f"export_timeout_seconds must be a number, got {type(timeout).__name__}" + ) + try: + normalized = float(timeout) + except (OverflowError, TypeError, ValueError) as exc: + raise ValueError( + "export_timeout_seconds must be representable as a finite float" + ) from exc + if ( + not math.isfinite(normalized) + or normalized <= 0 + or normalized > threading.TIMEOUT_MAX + ): + raise ValueError( + "export_timeout_seconds must be finite, greater than zero, and " + f"no greater than threading.TIMEOUT_MAX ({threading.TIMEOUT_MAX}), " + f"got {timeout!r}" + ) + object.__setattr__(self, "export_timeout_seconds", normalized) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_config.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_config.py index 623dbc6a..82d68f13 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_config.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_config.py @@ -11,8 +11,15 @@ from __future__ import annotations +import gc +import threading +import time +import weakref +from typing import Any + import pytest +import aws_durable_execution_sdk_python_insight.plugin as insight_plugin_module from aws_durable_execution_sdk_python_insight import ( EmitMode, OperationDetail, @@ -120,3 +127,271 @@ def test_readme_usage_call_shape_constructs_plugin(): ) plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) assert plugin._exporters == [exporter] + + +# -- export_timeout_seconds validation --------------------------------------- + + +def test_export_timeout_defaults_to_five_seconds(): + config = WorkflowInsightConfig() + assert config.export_timeout_seconds == 5.0 + assert workflow_insight(config)._export_timeout == 5.0 + + +@pytest.mark.parametrize("value", [0.1, 1, 2.5, 30, threading.TIMEOUT_MAX]) +def test_export_timeout_accepts_finite_positive_numbers(value): + config = WorkflowInsightConfig(export_timeout_seconds=value) + assert config.export_timeout_seconds == float(value) + assert workflow_insight(config)._export_timeout == float(value) + + +@pytest.mark.parametrize( + "value", + [ + 0, + 0.0, + -1, + -0.5, + float("nan"), + float("inf"), + float("-inf"), + threading.TIMEOUT_MAX * 2, + 10**1000, + True, # bool is an int subtype but must be rejected explicitly + False, + "5", # non-numeric + None, + ], +) +def test_export_timeout_rejects_invalid_values(value): + with pytest.raises(ValueError): + WorkflowInsightConfig(export_timeout_seconds=value) + + +# -- exporter-instance identity validation ----------------------------------- + + +class _StubExporter: + """Minimal exporter that is intentionally NOT hashable/comparable. + + ``__eq__``/``__hash__`` are disabled so the identity check cannot lean on + equality or hashing -- it must compare object identity (``is``) only. + """ + + max_record_size_bytes: int | None = None + + __hash__ = None # type: ignore[assignment] + + def __eq__(self, other): # pragma: no cover - must never be called + raise AssertionError("identity validation must not use __eq__") + + def render(self, record): # pragma: no cover - not exercised here + return record + + def export(self, record): # pragma: no cover - not exercised here + return None + + def flush(self): # pragma: no cover - not exercised here + return None + + +def test_same_exporter_instance_twice_raises_value_error(): + exporter = _StubExporter() + with pytest.raises(ValueError, match="same exporter instance"): + WorkflowInsightConfig(exporters=[exporter, exporter]) + + +def test_same_exporter_instance_among_others_raises(): + dup = _StubExporter() + with pytest.raises(ValueError, match="same exporter instance"): + WorkflowInsightConfig(exporters=[_StubExporter(), dup, _StubExporter(), dup]) + + +def test_two_distinct_same_class_instances_accepted_each_own_lane(): + a = _StubExporter() + b = _StubExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[a, b])) + # Both distinct instances are kept verbatim, in order. + assert plugin._exporters == [a, b] + # Each distinct instance gets its own lane (one background worker each). + lanes = plugin._scheduler._lanes + assert len(lanes) == 2 + assert [lane._exporter for lane in lanes] == [a, b] + + +def test_exporter_instance_cannot_be_shared_across_live_plugins(): + exporter = _StubExporter() + config = WorkflowInsightConfig(exporters=[exporter]) + owner = workflow_insight(config) + + with pytest.raises(ValueError, match="shared across Workflow Insight"): + workflow_insight(config) + + assert owner._exporters == [exporter] + + +def test_exporter_instance_can_be_reused_after_owner_is_collected(): + exporter = _StubExporter() + config = WorkflowInsightConfig(exporters=[exporter]) + owner = workflow_insight(config) + del owner + gc.collect() + + replacement = workflow_insight(config) + assert replacement._exporters == [exporter] + + +def test_exporter_ownership_persists_while_lane_worker_is_alive(): + exporter = _StubExporter() + config = WorkflowInsightConfig(exporters=[exporter]) + owner = workflow_insight(config) + lane = owner._scheduler._lanes[0] + owner._scheduler.schedule( + "arn:aws:lambda:us-east-1:1:function:f:1/durable-execution/e/i", + {"executionArn": "arn:e", "status": "RUNNING", "operations": []}, + ) + deadline = time.monotonic() + 5.0 + while not lane._worker_alive() and time.monotonic() < deadline: + time.sleep(0.005) + assert lane._worker_alive() + + del owner + gc.collect() + with pytest.raises(ValueError, match="shared across Workflow Insight"): + workflow_insight(config) + + lane.request_stop_when_idle() + deadline = time.monotonic() + 5.0 + while lane._worker_alive() and time.monotonic() < deadline: + time.sleep(0.005) + assert not lane._worker_alive() + del lane + gc.collect() + + replacement = workflow_insight(config) + assert replacement._exporters == [exporter] + + +def test_exporter_plugin_cycle_is_not_rooted_by_ownership_registry(): + exporter = _StubExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + lane = plugin._scheduler._lanes[0] + exporter.plugin = plugin + exporter_ref = weakref.ref(exporter) + plugin_ref = weakref.ref(plugin) + lane_ref = weakref.ref(lane) + assert all( + owner_ref.__callback__ is None + for owner_ref in insight_plugin_module._exporter_owners + ) + + del lane + del plugin + del exporter + gc.collect() + + assert exporter_ref() is None + assert plugin_ref() is None + assert lane_ref() is None + + +def test_exporter_finalization_runs_after_ownership_lock_is_released(): + owner_acquired = threading.Event() + ordinary_owner_released = threading.Event() + finalized = threading.Event() + finalizer_lock_results: list[bool] = [] + + class FinalizingExporter(_StubExporter): + def __del__(self) -> None: + acquired = insight_plugin_module._exporter_owner_lock.acquire(timeout=1.0) + finalizer_lock_results.append(acquired) + if acquired: + insight_plugin_module._exporter_owner_lock.release() + finalized.set() + + class CoordinatedLaneRef: + def __init__(self, lane_ref: weakref.ReferenceType[Any]) -> None: + self._lane_ref = lane_ref + + def __call__(self) -> Any: + owner = self._lane_ref() + owner_acquired.set() + assert ordinary_owner_released.wait(5.0) + return owner + + with insight_plugin_module._exporter_owner_lock: + saved_owners = list(insight_plugin_module._exporter_owners) + insight_plugin_module._exporter_owners.clear() + + exporter = FinalizingExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + lane = plugin._scheduler._lanes[0] + lane_ref = weakref.ref(lane) + plugin_holder = [plugin] + coordinated_ref = CoordinatedLaneRef(lane_ref) + with insight_plugin_module._exporter_owner_lock: + insight_plugin_module._exporter_owners[:] = [coordinated_ref] # type: ignore[list-item] + + del lane + del plugin + del exporter + + def release_ordinary_owner() -> None: + assert owner_acquired.wait(5.0) + plugin_holder.clear() + gc.collect() + ordinary_owner_released.set() + + release_thread = threading.Thread(target=release_ordinary_owner) + release_thread.start() + try: + insight_plugin_module._claim_exporter_lanes([]) + release_thread.join(5.0) + assert not release_thread.is_alive() + assert finalized.wait(5.0) + assert finalizer_lock_results == [True] + assert lane_ref() is None + finally: + ordinary_owner_released.set() + release_thread.join(5.0) + with insight_plugin_module._exporter_owner_lock: + insight_plugin_module._exporter_owners[:] = saved_owners + + +def test_default_exporter_unaffected_by_instance_check(): + # No exporters configured -> single default LambdaLogExporter, one lane; the + # identity check never trips on the empty list. + plugin = workflow_insight(WorkflowInsightConfig()) + assert len(plugin._exporters) == 1 + assert len(plugin._scheduler._lanes) == 1 + + +def test_mutating_original_exporter_list_to_duplicate_raises_at_plugin_creation(): + exporter = _StubExporter() + exporters = [exporter] + config = WorkflowInsightConfig(exporters=exporters) + exporters.append(exporter) + + with pytest.raises(ValueError, match="same exporter instance"): + workflow_insight(config) + + +def test_mutating_config_exporters_to_duplicate_raises_at_plugin_creation(): + exporter = _StubExporter() + config = WorkflowInsightConfig(exporters=[exporter]) + config.exporters.append(exporter) + + with pytest.raises(ValueError, match="same exporter instance"): + workflow_insight(config) + + +def test_mutating_exporter_list_with_distinct_instance_is_accepted(): + first = _StubExporter() + second = _StubExporter() + exporters = [first] + config = WorkflowInsightConfig(exporters=exporters) + exporters.append(second) + + plugin = workflow_insight(config) + assert plugin._exporters == [first, second] + assert [lane._exporter for lane in plugin._scheduler._lanes] == [first, second] 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 new file mode 100644 index 00000000..8b438e1f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -0,0 +1,760 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the asynchronous export scheduler (``_export_scheduler``). + +These drive the scheduler directly with plain record dicts and purpose-built +exporter doubles. Synchronization uses events/predicates (not sleeps) so the +coalescing, fairness, drain, flush-ordering, timeout and thread-lifecycle +invariants are asserted deterministically rather than by timing luck. +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _ExportScheduler, +) + + +ARN_A = "arn:aws:lambda:us-west-2:1:function:f:$LATEST/durable-execution/exec-a/inv-1" +ARN_B = "arn:aws:lambda:us-west-2:1:function:f:$LATEST/durable-execution/exec-b/inv-1" +ARN_C = "arn:aws:lambda:us-west-2:1:function:f:$LATEST/durable-execution/exec-c/inv-1" +ARN_D = "arn:aws:lambda:us-west-2:1:function:f:$LATEST/durable-execution/exec-d/inv-1" + + +def _wait_until(predicate, timeout: float = 5.0, interval: float = 0.005) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(interval) + return predicate() + + +def _rec(arn: str, value: str, *, status: str = "RUNNING") -> dict[str, Any]: + return {"executionArn": arn, "status": status, "v": value, "operations": []} + + +def _insight_thread_count() -> int: + return sum( + 1 for t in threading.enumerate() if t.name.startswith("workflow-insight-export") + ) + + +def _lane_worker_count(lane) -> int: + """Count live worker threads that belong to *this* lane by identity. + + Each lane names its worker ``workflow-insight-export-{id(lane)}``, so this is + scoped to the given lane and is unaffected by daemon workers other tests may + still be winding down -- unlike a process-global thread-count delta. + """ + name = f"workflow-insight-export-{id(lane)}" + return sum(1 for t in threading.enumerate() if t.name == name and t.is_alive()) + + +class RecordingExporter: + """Records every export/flush in call order (fast, non-blocking).""" + + def __init__(self, max_record_size_bytes: int | None = None) -> None: + self.max_record_size_bytes = max_record_size_bytes + self.calls: list[tuple[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.calls.append(("export", record.get("v"))) + + def flush(self) -> None: + with self._lock: + self.calls.append(("flush", None)) + + def exported_values(self) -> list[Any]: + with self._lock: + return [v for kind, v in self.calls if kind == "export"] + + +class BlockingExporter: + """Blocks inside ``export`` until released; signals when an export starts.""" + + def __init__(self, max_record_size_bytes: int | None = None) -> None: + self.max_record_size_bytes = max_record_size_bytes + self._release = threading.Event() + self.started = threading.Event() + self.exported: list[Any] = [] + self.flushed = 0 + self._lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + self.started.set() + self._release.wait(5.0) + with self._lock: + self.exported.append(record.get("v")) + + def flush(self) -> None: + with self._lock: + self.flushed += 1 + + def release(self) -> None: + self._release.set() + + def exported_values(self) -> list[Any]: + with self._lock: + return list(self.exported) + + +class BlockingFlushExporter(RecordingExporter): + """Exports normally but blocks inside ``flush`` until released. + + Lets a test drive the worker until it has already popped a flush barrier and + is stuck mid-``flush`` -- the "already in flight" cancellation race. + """ + + def __init__(self, max_record_size_bytes: int | None = None) -> None: + super().__init__(max_record_size_bytes) + self.flush_started = threading.Event() + self._flush_release = threading.Event() + + def flush(self) -> None: + self.flush_started.set() + self._flush_release.wait(5.0) + super().flush() + + def release_flush(self) -> None: + self._flush_release.set() + + +class BlockingBufferedExporter: + """Blocks export and publishes buffered records only when flush runs.""" + + max_record_size_bytes = None + + def __init__(self) -> None: + self.started = threading.Event() + self._release = threading.Event() + self.buffered: list[Any] = [] + self.published: list[Any] = [] + self.flushed = 0 + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + self.started.set() + self._release.wait(5.0) + self.buffered.append(record.get("v")) + + def flush(self) -> None: + self.flushed += 1 + self.published.extend(self.buffered) + self.buffered.clear() + + def release(self) -> None: + self._release.set() + + +class FailingExporter: + """Raises in both export and flush.""" + + def __init__(self, max_record_size_bytes: int | None = None) -> None: + self.max_record_size_bytes = max_record_size_bytes + self.export_calls = 0 + self.flush_calls = 0 + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + self.export_calls += 1 + raise RuntimeError("export boom") + + def flush(self) -> None: + self.flush_calls += 1 + raise RuntimeError("flush boom") + + +class _Uncopyable(dict[str, str]): + """A JSON-serializable payload whose ``deepcopy`` raises.""" + + def __init__(self) -> None: + super().__init__({"value": "safe"}) + + def __deepcopy__(self, memo: dict[int, Any]) -> Any: + raise RuntimeError("uncopyable payload") + + +# -- lazy worker creation / one worker per exporter -------------------------- + + +def test_no_worker_before_first_schedule(): + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + assert lane._worker is None + assert not lane._worker_alive() + + +def test_worker_created_lazily_on_first_schedule(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(lane._worker_alive) + exporter.release() + scheduler.end_invocation(5.0) + assert _wait_until(lambda: not lane._worker_alive()) + + +def test_one_worker_per_exporter(): + e1, e2 = BlockingExporter(), BlockingExporter() + scheduler = _ExportScheduler([e1, e2]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(lambda: e1.started.is_set() and e2.started.is_set()) + assert _lane_worker_count(scheduler._lanes[0]) == 1 + assert _lane_worker_count(scheduler._lanes[1]) == 1 + e1.release() + e2.release() + scheduler.end_invocation(5.0) + assert _wait_until( + lambda: not any(lane._worker_alive() for lane in scheduler._lanes) + ) + + +def test_worker_start_failure_disables_lane_and_fails_barrier(monkeypatch): + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + + def fail_start(self): + raise RuntimeError("cannot start new thread") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + + assert lane._disabled is True + assert lane._pending_count() == 0 + assert lane._queue_len() == 0 + assert scheduler.end_invocation(0.1) is False + assert exporter.exported_values() == [] + + +def test_repeated_scheduling_does_not_grow_threads(): + base = _insight_thread_count() + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + for i in range(100): + scheduler.schedule(ARN_A, _rec(ARN_A, f"v{i}", status="RUNNING")) + # A single lane never runs more than one worker at a time. + assert _insight_thread_count() - base <= 1 + scheduler.end_invocation(5.0) + assert _wait_until(lambda: not lane._worker_alive()) + + +# -- coalescing / fairness / isolation --------------------------------------- + + +def test_same_execution_coalescing_exports_inflight_then_latest(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) # a1 is in flight + # While a1 is in flight, a2 and a3 arrive and coalesce to the latest (a3). + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + scheduler.schedule(ARN_A, _rec(ARN_A, "a3")) + exporter.release() + assert _wait_until(lambda: exporter.exported_values() == ["a1", "a3"]) + scheduler.end_invocation(5.0) + + +def test_latest_pending_replaces_across_executions(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) + scheduler.schedule(ARN_C, _rec(ARN_C, "c1")) + scheduler.schedule(ARN_D, _rec(ARN_D, "d1")) + + assert scheduler._lanes[0]._pending_count() == 1 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["a1", "d1"] + + +def test_terminal_record_supersedes_pending_running(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "r1", status="RUNNING")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_A, _rec(ARN_A, "r2", status="RUNNING")) + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + exporter.release() + assert _wait_until(lambda: exporter.exported_values() == ["r1", "final"]) + scheduler.end_invocation(5.0) + + +# -- copy failure isolation --------------------------------------------------- + + +def test_uncopyable_custom_value_reaches_exporter_render(): + class CustomRenderExporter(RecordingExporter): + def __init__(self) -> None: + super().__init__(max_record_size_bytes=10_000) + self.rendered_values: list[str] = [] + + def render(self, record: dict[str, Any]) -> Any: + value = record["payload"]["value"] + self.rendered_values.append(value) + return {"value": value} + + exporter = CustomRenderExporter() + scheduler = _ExportScheduler([exporter]) + record = _rec(ARN_A, "before-render") + record["payload"] = _Uncopyable() + + scheduler.schedule(ARN_A, record) + scheduler.end_invocation(5.0) + + assert exporter.rendered_values == ["safe"] + assert exporter.exported_values() == ["before-render"] + + +def test_uncopyable_custom_value_does_not_alias_record_containers(): + class MutatingRenderExporter(RecordingExporter): + def render(self, record: dict[str, Any]) -> Any: + record["mutated"] = True + return record + + exporter = MutatingRenderExporter() + scheduler = _ExportScheduler([exporter]) + record = _rec(ARN_A, "custom") + record["payload"] = _Uncopyable() + + scheduler.schedule(ARN_A, record) + scheduler.end_invocation(5.0) + + assert "mutated" not in record + assert exporter.exported_values() == ["custom"] + + +# -- non-blocking hook return / fast-vs-slow isolation ----------------------- + + +def test_schedule_returns_immediately_while_exporter_blocked(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + # The exporter is now blocked mid-export; a further schedule must not block. + start = time.monotonic() + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + assert time.monotonic() - start < 0.5 + exporter.release() + scheduler.end_invocation(5.0) + + +def test_fast_lane_proceeds_while_other_lane_blocked(): + blocked = BlockingExporter() + fast = RecordingExporter() + scheduler = _ExportScheduler([blocked, fast]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + # Fast lane delivers even though the blocked lane is stuck on the same record. + assert _wait_until(lambda: fast.exported_values() == ["a1"]) + assert blocked.exported_values() == [] + blocked.release() + scheduler.end_invocation(5.0) + + +# -- drain / flush ordering --------------------------------------------------- + + +def test_drain_waits_for_final_export(): + class SlowExporter(RecordingExporter): + def export(self, record: dict[str, Any]) -> None: + time.sleep(0.2) + super().export(record) + + exporter = SlowExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + ok = scheduler.end_invocation(5.0) + assert ok is True + assert exporter.exported_values() == ["final"] + + +def test_flush_happens_after_export(): + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + scheduler.end_invocation(5.0) + kinds = [kind for kind, _ in exporter.calls] + assert kinds == ["export", "flush"] + + +def test_export_and_flush_exceptions_are_isolated(): + failing = FailingExporter() + good = RecordingExporter() + scheduler = _ExportScheduler([failing, good]) + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + # Must not raise even though one exporter fails in both export and flush. + ok = scheduler.end_invocation(5.0) + assert ok is True + assert failing.export_calls == 1 + assert failing.flush_calls == 1 + # The healthy exporter still delivered and flushed. + assert good.exported_values() == ["final"] + assert ("flush", None) in good.calls + + +# -- shared timeout ----------------------------------------------------------- + + +def test_shared_timeout_bounds_invocation_end_delay(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + start = time.monotonic() + ok = scheduler.end_invocation(0.2) + elapsed = time.monotonic() - start + assert ok is False # degraded to best-effort + assert elapsed < 2.0 # bounded by the shared deadline, not the blocked export + exporter.release() # let the daemon drain and exit + # Wait for the released worker to actually stop so it cannot leak into a + # later test's baseline thread count. + assert _wait_until(lambda: not lane._worker_alive()) + + +def test_shared_timeout_across_multiple_lanes_is_not_additive(): + e1, e2 = BlockingExporter(), BlockingExporter() + scheduler = _ExportScheduler([e1, e2]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(lambda: e1.started.is_set() and e2.started.is_set()) + start = time.monotonic() + ok = scheduler.end_invocation(0.3) + elapsed = time.monotonic() - start + assert ok is False + # One shared deadline covers both lanes, so total wait is ~0.3s, not 0.6s. + assert elapsed < 0.9 + e1.release() + e2.release() + # Wait for both released workers to actually stop so neither leaks into a + # later test's baseline thread count. + assert _wait_until( + lambda: not any(lane._worker_alive() for lane in scheduler._lanes) + ) + + +# -- worker lifecycle --------------------------------------------------------- + + +def test_blocked_worker_is_not_replaced(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + worker = lane._worker + assert worker is not None and worker.is_alive() + # The blocked lane already has exactly one live worker of its own. + assert _lane_worker_count(lane) == 1 + # More scheduling and an invocation-end (which enqueues a flush + requests + # stop) must not spawn a replacement while the worker is blocked. + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) + scheduler.end_invocation(0.1) + # Identity: the lane still holds the SAME blocked worker -- no replacement + # thread was swapped in -- and it is still the only live worker for this + # lane. Both checks are scoped to this lane, so they cannot flake on daemon + # workers other tests are winding down. + assert lane._worker is worker + assert worker.is_alive() + assert _lane_worker_count(lane) == 1 + exporter.release() + + +def test_idle_worker_exits_after_drain(): + base = _insight_thread_count() + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + scheduler.end_invocation(5.0) + assert _wait_until(lambda: not lane._worker_alive()) + assert _wait_until(lambda: _insight_thread_count() <= base) + + +def test_repeated_invocation_cycles_do_not_leak_threads(): + base = _insight_thread_count() + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + for i in range(20): + scheduler.schedule(ARN_A, _rec(ARN_A, f"final-{i}", status="SUCCEEDED")) + scheduler.end_invocation(5.0) + assert _wait_until(lambda: not lane._worker_alive()) + assert _wait_until(lambda: _insight_thread_count() <= base) + assert len(exporter.exported_values()) == 20 + + +# -- structurally bounded pending slot ---------------------------------------- + + +def test_latest_pending_slot_stays_bounded_during_burst(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + for index in range(100): + arn = f"{ARN_B}-{index}" + scheduler.schedule(arn, _rec(arn, f"pending-{index}")) + + assert lane._pending_count() == 1 + assert lane._queue_len() == 1 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight", "pending-99"] + + +def test_schedule_does_not_inspect_custom_values(): + inspected = threading.Event() + + class OpaqueValue: + def __sizeof__(self) -> int: + inspected.set() + raise AssertionError("checkpoint scheduling must not inspect values") + + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + record = _rec(ARN_B, "opaque") + record["payload"] = OpaqueValue() + + scheduler.schedule(ARN_B, record) + + assert inspected.is_set() is False + assert scheduler._lanes[0]._pending_count() == 1 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight", "opaque"] + + +def test_timed_out_barrier_flushes_eventually_and_worker_exits(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + ok = scheduler.end_invocation(0.1) + assert ok is False + # The caller returns on time, but one detached flush remains queued so a + # buffered exporter can publish before the worker exits idle. + assert lane._queued_flush_count() == 1 + exporter.release() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.flushed == 1 + + +def test_timed_out_buffered_export_is_published_eventually(): + exporter = BlockingBufferedExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "terminal", status="SUCCEEDED")) + assert _wait_until(exporter.started.is_set) + + assert scheduler.end_invocation(0.1) is False + assert exporter.published == [] + exporter.release() + + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.published == ["terminal"] + assert exporter.flushed == 1 + + +def test_repeated_timeouts_behind_blocked_exporter_stay_bounded(): + """Warm timeouts coalesce to one eventual flush on the same worker.""" + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + worker = lane._worker + assert worker is not None and worker.is_alive() + + for i in range(50): + scheduler.schedule(ARN_A, _rec(ARN_A, f"a{i + 2}")) + assert scheduler.end_invocation(0.02) is False + # One latest detached flush plus at most one coalesced record token. + assert lane._queued_flush_count() == 1 + assert lane._queue_len() <= 2 + + assert lane._queue_len() <= 2 + assert lane._pending_count() <= 1 + assert lane._queued_flush_count() == 1 + assert lane._worker is worker + assert worker.is_alive() + assert _lane_worker_count(lane) == 1 + assert exporter.flushed == 0 + + exporter.release() + assert _wait_until(lambda: not lane._worker_alive()) + exported = exporter.exported_values() + assert exported[0] == "a1" + assert len(exported) <= 2 + assert exporter.flushed == 1 + + +def test_cancel_flush_replaces_queued_barrier_with_detached_flush(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + barrier = lane.enqueue_flush() + assert lane._queued_flush_count() == 1 + + lane.cancel_flush(barrier) + + assert lane._queued_flush_count() == 1 + assert barrier.canceled is True + assert barrier.is_done() + exporter.release() + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.flushed == 1 + assert exporter.exported_values() == ["a1"] + + +def test_cancel_flush_after_pop_lets_worker_complete_barrier(): + """Already-popped race: the worker has taken the barrier and is mid-flush, + so cancel_flush only marks it cancelled and leaves completion to the worker. + The in-flight flush is not interrupted.""" + exporter = BlockingFlushExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + barrier = lane.enqueue_flush() + # Worker exports a1, pops the barrier, and enters flush (now in flight). + assert _wait_until(exporter.flush_started.is_set) + assert lane._queued_flush_count() == 0 # already popped from the queue + assert not barrier.is_done() # worker still inside flush + # Cancelling now must NOT complete it here (the worker owns completion) and + # must NOT interrupt the in-flight flush. + lane.cancel_flush(barrier) + assert barrier.canceled is True + assert not barrier.is_done() + # Release the in-flight flush; the worker completes the barrier itself. + exporter.release_flush() + assert _wait_until(barrier.is_done) + # The flush already in flight ran to completion exactly once (not killed). + assert exporter.calls.count(("flush", None)) == 1 + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + + +def test_cancel_popped_barrier_preserves_later_detached_flush(): + exporter = BlockingFlushExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + older = lane.enqueue_flush() + assert _wait_until(exporter.flush_started.is_set) + + later = lane.enqueue_flush() + lane.cancel_flush(later) + assert lane._queued_flush_count() == 1 + lane.cancel_flush(older) + + assert lane._queued_flush_count() == 1 + assert later.is_done() + exporter.release_flush() + assert _wait_until(older.is_done) + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.calls.count(("flush", None)) == 2 + + +def test_cancel_queued_barrier_preserves_later_detached_flush(): + exporter = BlockingBufferedExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + + older = lane.enqueue_flush() + scheduler.schedule(ARN_B, _rec(ARN_B, "a2")) + later = lane.enqueue_flush() + lane.cancel_flush(later) + lane.cancel_flush(older) + + assert older.is_done() + assert later.is_done() + assert lane._queue_len() == 2 + assert lane._queued_flush_count() == 1 + exporter.release() + lane.request_stop_when_idle() + + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.published == ["a1", "a2"] + assert exporter.flushed == 1 + + +def test_sizing_resource_failure_does_not_call_exporter(monkeypatch): + import aws_durable_execution_sdk_python_insight.truncation as truncation_module + + def fail_dumps(*args: Any, **kwargs: Any) -> str: # noqa: ARG001 + raise MemoryError("sizing exhausted") + + monkeypatch.setattr(truncation_module.json, "dumps", fail_dumps) + exporter = RecordingExporter(max_record_size_bytes=1_000) + scheduler = _ExportScheduler([exporter]) + + scheduler.schedule(ARN_A, _rec(ARN_A, "resource-failure")) + scheduler.end_invocation(5.0) + + assert exporter.exported_values() == [] + + +def test_deeply_nested_record_falls_back_to_original_for_custom_render(): + class DeepRenderExporter(RecordingExporter): + def __init__(self) -> None: + super().__init__(max_record_size_bytes=1_000) + self.rendered_depth = 0 + + def render(self, record: dict[str, Any]) -> Any: + value = record["payload"] + while isinstance(value, list): + self.rendered_depth += 1 + value = value[0] + return {"value": value} + + depth = 2_000 + payload: Any = "leaf" + for _ in range(depth): + payload = [payload] + record = _rec(ARN_A, "deep") + record["payload"] = payload + exporter = DeepRenderExporter() + scheduler = _ExportScheduler([exporter]) + + scheduler.schedule(ARN_A, record) + scheduler.end_invocation(5.0) + + assert exporter.rendered_depth == depth + assert exporter.exported_values() == ["deep"] 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 485bcaab..3cefcd32 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 @@ -365,10 +365,15 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): assert rec["durationMs"] is not None and rec["durationMs"] >= 0 -# -- on-change emits an updated record per change (comment 2) ---------------- +# -- on-change schedules RUNNING records, coalescing intermediates (comment 2) -- -def test_on_change_emits_running_on_each_change(): +def test_on_change_schedules_running_and_delivers_terminal(): + # Under the async scheduler, rapid cumulative RUNNING snapshots for one + # execution may coalesce (the design explicitly allows a lane to observe only + # a subset of intermediate records). The invariants that always hold: the + # terminal record is delivered last, every earlier record is RUNNING, and the + # terminal record carries the full, de-duplicated operation set. exporter = CaptureExporter() plugin = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") @@ -376,23 +381,23 @@ def test_on_change_emits_running_on_each_change(): op1 = _step("s1", op_id="1") op2 = _step("s2", op_id="2") - plugin.on_invocation_start(_start(operations={})) # RUNNING #1 (start) + plugin.on_invocation_start(_start(operations={})) # RUNNING (start) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1) ) - ) # RUNNING #2 + ) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op2), operations=_ops(op1, op2) ) - ) # RUNNING #3 - plugin.on_invocation_end(_end(operations=_ops(op1, op2))) # SUCCEEDED #4 + ) + plugin.on_invocation_end(_end(operations=_ops(op1, op2))) # SUCCEEDED (terminal) statuses = [r["status"] for r in exporter.records] - assert statuses == ["RUNNING", "RUNNING", "RUNNING", "SUCCEEDED"] - # The record emitted after the 2nd change already carries both operations. - assert [op["name"] for op in exporter.records[2]["operations"]] == ["s1", "s2"] + assert statuses, "at least the terminal record must be delivered" + assert statuses[-1] == "SUCCEEDED" + assert set(statuses[:-1]) <= {"RUNNING"} final = exporter.records[-1] assert [op["name"] for op in final["operations"]] == ["s1", "s2"] # No duplicate operation entries within a record (no end/change double-count). diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py new file mode 100644 index 00000000..e0a9c432 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py @@ -0,0 +1,393 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Plugin-level tests for the asynchronous export path. + +These drive the plugin through the real SDK hook dataclasses and assert the +scheduler-backed behavior the design requires: non-``on-change`` modes do no +operation-change work, a blocked exporter never blocks a hook, the +invocation-end drain is bounded by ``export_timeout_seconds``, and a buffered +exporter only publishes after the invocation-end flush. +""" + +from __future__ import annotations + +import datetime +import threading +import time +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python.lambda_service import ( + OperationStatus, + OperationSubType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationChangeInfo, + OperationEndInfo, + OperationInfo, + OperationType, +) + +from aws_durable_execution_sdk_python_insight import ( + WorkflowInsightConfig, + workflow_insight, +) + + +ARN = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-1/inv-1" +ARN_A = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-a/inv-1" +ARN_B = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-b/inv-1" +T0 = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) +T1 = datetime.datetime(2026, 1, 1, 0, 0, 1, tzinfo=datetime.UTC) + + +def _wait_until(predicate, timeout: float = 5.0, interval: float = 0.005) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(interval) + return predicate() + + +def _step(name: str, op_id: str) -> OperationInfo: + return OperationEndInfo( + operation_id=op_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=name, + parent_id=None, + start_time=T0, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=T1, + result=None, + error=None, + attempt=1, + ) + + +def _ops(*ops: OperationInfo) -> dict[str, OperationInfo]: + return {op.operation_id: op for op in ops} + + +def _start(operations: dict[str, OperationInfo]) -> InvocationStartInfo: + return InvocationStartInfo( + request_id=None, + execution_arn=ARN, + is_first_invocation=True, + execution_start_time=T0, + execution_input="World", + operations=operations, + ) + + +def _end(operations: dict[str, OperationInfo]) -> InvocationEndInfo: + return InvocationEndInfo( + request_id=None, + execution_arn=ARN, + is_first_invocation=True, + execution_start_time=T0, + status=InvocationStatus.SUCCEEDED, + error=None, + execution_result='"Hello, World!"', + operations=operations, + ) + + +def _start_arn(arn: str, operations: dict[str, OperationInfo]) -> InvocationStartInfo: + return InvocationStartInfo( + request_id=None, + execution_arn=arn, + is_first_invocation=True, + execution_start_time=T0, + execution_input="World", + operations=operations, + ) + + +def _end_arn(arn: str, operations: dict[str, OperationInfo]) -> InvocationEndInfo: + return InvocationEndInfo( + request_id=None, + execution_arn=arn, + is_first_invocation=True, + execution_start_time=T0, + status=InvocationStatus.SUCCEEDED, + error=None, + execution_result='"Hello, World!"', + operations=operations, + ) + + +class _BlockingExporter: + def __init__(self) -> None: + self.max_record_size_bytes: int | None = None + self._release = threading.Event() + self.started = threading.Event() + self.exported: list[dict[str, Any]] = [] + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self.started.set() + self._release.wait(5.0) + self.exported.append(record) + + def flush(self) -> None: + return None + + def release(self) -> None: + self._release.set() + + +class _BufferedExporter: + """Buffers exports and only publishes them when ``flush`` is called.""" + + def __init__(self) -> None: + self.max_record_size_bytes: int | None = None + self._buffer: list[dict[str, Any]] = [] + self.published: list[dict[str, Any]] = [] + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self._buffer.append(record) + + def flush(self) -> None: + self.published.extend(self._buffer) + self._buffer.clear() + + +# -- non-on-change modes do no operation-change work ------------------------- + + +def test_non_on_change_mode_skips_operation_change_work(): + exporter = _BufferedExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter]) + ) # on-complete + op = _step("s", "1") + plugin.on_invocation_start(_start({})) + # An operation-change in a non-on-change mode must not create/adopt state. + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + state = plugin._state.get(ARN) + assert state is not None and state.operations == {} # snapshot not adopted + assert state.scheduled is False # nothing scheduled on the change + + +def test_oversized_timeout_rejected_before_exporter_worker_starts(): + exporter = _BlockingExporter() + with pytest.raises(ValueError, match="threading.TIMEOUT_MAX"): + workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + export_timeout_seconds=threading.TIMEOUT_MAX * 2, + ) + ) + assert not exporter.started.is_set() + + +def test_worker_start_failure_does_not_skip_invocation_cleanup(monkeypatch): + exporter = _BufferedExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + op = _step("s", "1") + plugin.on_invocation_start(_start({})) + + def fail_start(self): + raise RuntimeError("cannot start new thread") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + plugin.on_invocation_end(_end(_ops(op))) + + assert plugin._state == {} + assert exporter.published == [] + assert plugin._scheduler._lanes[0]._disabled is True + + +# -- a blocked exporter never blocks a hook ---------------------------------- + + +def test_operation_change_returns_immediately_with_blocked_exporter(): + exporter = _BlockingExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + plugin.on_invocation_start(_start({})) # schedules RUNNING; worker blocks on it + assert _wait_until(exporter.started.is_set) + op = _step("s", "1") + start = time.monotonic() + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + assert time.monotonic() - start < 0.5 # returned without waiting on the export + exporter.release() + plugin.on_invocation_end(_end(_ops(op))) + + +# -- invocation-end drain is bounded by export_timeout_seconds ---------------- + + +def test_invocation_end_bounded_by_export_timeout(): + exporter = _BlockingExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], export_timeout_seconds=0.2) + ) + op = _step("s", "1") + plugin.on_invocation_start(_start({})) + assert not exporter.started.is_set() # on-complete: nothing scheduled at start + start = time.monotonic() + plugin.on_invocation_end(_end(_ops(op))) # schedules terminal; worker blocks + elapsed = time.monotonic() - start + assert elapsed < 2.0 # bounded by the 0.2s shared deadline + assert plugin._state == {} # state cleared even on degraded delivery + exporter.release() + + +# -- buffered exporter publishes only after the invocation-end flush ---------- + + +def test_buffered_exporter_publishes_after_flush(): + exporter = _BufferedExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + op = _step("s", "1") + plugin.on_invocation_start(_start({})) + plugin.on_invocation_end(_end(_ops(op))) # drains + flushes before returning + assert len(exporter.published) == 1 + assert exporter.published[0]["status"] == "SUCCEEDED" + + +# -- warm-container cross-invocation isolation -------------------------------- + + +class _OrderedBlockingExporter: + """Blocks the first export until released; records export order and flushes. + + Once released, subsequent exports return immediately (the release event stays + set), so a lane can drain a backlog without re-blocking. + """ + + def __init__(self) -> None: + self.max_record_size_bytes: int | None = None + self._release = threading.Event() + self.started = threading.Event() + self._lock = threading.Lock() + self.exported: list[str] = [] + self.flush_calls = 0 + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self.started.set() + self._release.wait(5.0) + with self._lock: + self.exported.append(record.get("executionArn", "")) + + def flush(self) -> None: + with self._lock: + self.flush_calls += 1 + + def release(self) -> None: + self._release.set() + + def exported_arns(self) -> list[str]: + with self._lock: + return list(self.exported) + + +def test_warm_container_cross_invocation_isolation_and_ordering(): + # One warm plugin instance handles two executions on the same exporter lane. + # Execution A blocks the lane and times out at invocation end; execution B + # schedules behind it. Both invocation-end waits stay bounded, B is neither + # lost nor merged into A, and after unblock the lane drains FIFO and flushes. + exporter = _OrderedBlockingExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], export_timeout_seconds=0.2) + ) + op = _step("s", "1") + + # -- Execution A: terminal end schedules a record; the lane blocks on it. -- + plugin.on_invocation_start(_start_arn(ARN_A, {})) + a_start = time.monotonic() + plugin.on_invocation_end(_end_arn(ARN_A, _ops(op))) # blocks; times out + a_elapsed = time.monotonic() - a_start + assert a_elapsed < 2.0 # bounded by the shared 0.2s deadline, not the export + assert _wait_until(exporter.started.is_set) # A is in flight + assert exporter.exported_arns() == [] # still blocked -> nothing delivered + + # -- Execution B arrives on the warm container while A is blocked. -------- + plugin.on_invocation_start(_start_arn(ARN_B, {})) + exporter.release() # let the lane drain A first + assert _wait_until(lambda: exporter.exported_arns() == [ARN_A]) + + b_start = time.monotonic() + plugin.on_invocation_end(_end_arn(ARN_B, _ops(op))) # bounded; drains + flush + b_elapsed = time.monotonic() - b_start + assert b_elapsed < 2.0 + + # B was delivered as its own record after A (FIFO), never merged into A. + assert _wait_until(lambda: exporter.exported_arns() == [ARN_A, ARN_B]) + # B's invocation-end flush completed (its barrier was not cancelled). + assert _wait_until(lambda: exporter.flush_calls >= 1) + assert plugin._state == {} # both executions cleared their state + + +# -- scheduled flag is read/written under the plugin lock -------------------- + + +def test_scheduled_flag_lock_helpers_track_scheduling(): + # The ``scheduled`` gate is now mutated/read through the plugin lock like + # every other _ExecutionState field. Pin the observable behavior: it starts + # False, flips True once a record is scheduled, and the lock-guarded read + # helper agrees with the raw attribute. The SDK serializes hooks, so this is + # a defensive/consistency check rather than a concurrency race test. + exporter = _BufferedExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + op = _step("s", "1") + + plugin.on_invocation_start(_start({})) # on-complete: nothing scheduled yet + state = plugin._state[ARN] + assert state.scheduled is False + assert plugin._was_scheduled(state) is False + + plugin.on_invocation_end(_end(_ops(op))) # schedules the terminal record + # State is cleared at invocation end, but the local reference still reflects + # the flip performed via the lock helper before the drain. + assert state.scheduled is True + assert plugin._was_scheduled(state) is True + assert len(exporter.published) == 1 + + +def test_no_op_invocation_leaves_scheduled_false(): + # on-complete + non-terminal (PENDING/RETRY) end schedules nothing, so the + # gate stays False and no flush/lane work is triggered. + exporter = _BufferedExporter() + plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin.on_invocation_start(_start({})) + state = plugin._state[ARN] + pending_end = InvocationEndInfo( + request_id=None, + execution_arn=ARN, + is_first_invocation=True, + execution_start_time=T0, + status=InvocationStatus.PENDING, + error=None, + execution_result=None, + operations={}, + ) + plugin.on_invocation_end(pending_end) + assert state.scheduled is False + assert exporter.published == []