diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index ed47c053..1dca3682 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -55,10 +55,20 @@ 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, every entry in `exporters` must be a **distinct +> instance**: passing the same object twice raises `ValueError` at construction. +> Two separate instances of the same exporter class (e.g. two `S3Exporter`s for +> different buckets) are fine — each gets its own worker. Rapid cumulative +> snapshots for one execution 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..4e2e97b9 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -0,0 +1,357 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Asynchronous, coalescing export scheduler for the Workflow Insight plugin. + +The plugin builds one canonical ``WorkflowInsight`` record on the SDK checkpoint +thread and hands it to :class:`_ExportScheduler`. The scheduler keeps all +exporter-specific work -- per-exporter copy, ``render``, truncation, ``export`` +and ``flush`` -- off the checkpoint thread by running it in a lazily-created +daemon worker, one per exporter ("lane"). Scheduling a record only enqueues it +and returns immediately, so ``on_operation_change`` never blocks on a slow +exporter. + +Design (``workflow-insight-async-export-design.md``): + +* One lazy daemon worker per exporter lane; never more than one live worker per + lane, and a blocked worker is retained -- never replaced -- so threads cannot + grow without bound. +* Per lane, at most one in-flight record and one latest *pending* record per + execution ARN. Records are cumulative snapshots, so a newer pending record for + an ARN replaces the older one (coalescing); an in-flight record is never + cancelled. Updating a pending ARN moves it to the back of the queue for + fairness across ARNs. Pending ARNs are capped; the oldest is evicted when the + cap is exceeded (only reachable behind a blocked/slow exporter). +* Invocation end enqueues one flush barrier per touched lane after the latest + record and waits for all barriers under a single shared timeout deadline. On + timeout the workflow response is returned, degradation is logged, and each + stale barrier is cancelled and its still-queued ``_FLUSH`` marker pulled from + the lane so barriers cannot accumulate behind a blocked worker; any blocked + worker stays daemonized (a synchronous Python ``export()`` cannot be safely + killed) and completes an already-popped barrier itself. +* Idle workers exit after the drain/flush request, so a normal invocation leaves + no lingering thread. +""" + +from __future__ import annotations + +import copy +import logging +import threading +import time +from collections import OrderedDict, 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") + +# Upper bound on distinct executions with a record waiting in a single lane. +# Only reached when a lane's exporter is blocked or slow; the oldest pending +# execution is then evicted (best-effort delivery) so plugin memory stays +# bounded regardless of how long a worker stays blocked. +_DEFAULT_MAX_PENDING_EXECUTIONS = 1024 + +# 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") + + def __init__(self) -> None: + self._event = threading.Event() + self.canceled = 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, + *, + max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, + ) -> None: + self._exporter = exporter + self._max_pending = max(1, max_pending_executions) + # 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, arn) or (_FLUSH, barrier). + self._queue: deque[tuple[str, Any]] = deque() + # arn -> latest pending record (coalesced). Insertion order is the + # fairness order; updating an arn moves it to the back. + self._pending: OrderedDict[str, dict[str, Any]] = OrderedDict() + self._stop_when_idle = False + self._worker: threading.Thread | None = None + + # -- producer API (checkpoint / invocation-end threads) ------------------- + + def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: + with self._cond: + self._stop_when_idle = False + if execution_arn in self._pending: + # Coalesce: replace the pending record and move it to the back so + # a busy execution cannot starve the others. + self._pending[execution_arn] = record + self._pending.move_to_end(execution_arn) + self._move_record_token_to_back(execution_arn) + else: + self._pending[execution_arn] = record + self._queue.append((_RECORD, execution_arn)) + self._enforce_pending_cap() + self._ensure_worker_locked() + self._cond.notify() + + def enqueue_flush(self) -> _FlushBarrier: + barrier = _FlushBarrier() + with self._cond: + 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: + """Cancel a timed-out flush barrier so it cannot pile up behind a + blocked worker. + + Under the lane lock: mark the barrier cancelled and, if its ``_FLUSH`` + marker is still queued, remove that exact marker and complete the + barrier here. Removing it is what keeps queue/barrier state bounded + across many warm invocations behind a blocked exporter -- otherwise one + stale barrier per invocation would accumulate behind the stuck worker. + + If the worker has already popped the marker (the flush is in flight or + about to run) the marker is no longer in the queue: we only set + ``canceled`` and leave completion to the worker, which skips the + now-pointless flush and completes the barrier itself. A synchronous + in-flight ``flush()`` is never interrupted. + """ + with self._cond: + barrier.canceled = True + for index, (kind, payload) in enumerate(self._queue): + if kind == _FLUSH and payload is barrier: + del self._queue[index] + barrier.complete() + return + + # -- queue bookkeeping (must hold ``_cond``) ------------------------------ + + def _move_record_token_to_back(self, execution_arn: str) -> None: + for index, (kind, payload) in enumerate(self._queue): + if kind == _RECORD and payload == execution_arn: + del self._queue[index] + self._queue.append((_RECORD, execution_arn)) + return + # No token means the arn is currently in flight; a fresh token will be + # appended when it leaves flight (the next schedule sees it absent from + # ``_pending``), which yields the "export A then latest" behavior. + + def _enforce_pending_cap(self) -> None: + while len(self._pending) > self._max_pending: + old_arn, _ = self._pending.popitem(last=False) + self._remove_record_token(old_arn) + _logger.warning( + "workflow-insight: export lane for %s is full " + "(cap=%d); dropping pending record for %s", + type(self._exporter).__name__, + self._max_pending, + old_arn, + ) + + def _remove_record_token(self, execution_arn: str) -> None: + for index, (kind, payload) in enumerate(self._queue): + if kind == _RECORD and payload == execution_arn: + del self._queue[index] + return + + 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._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 + worker.start() + + # -- 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.pop(payload, None) + if record is None: + continue + + if kind == _RECORD and record is not None: + self._export_one(record) + else: # _FLUSH + barrier: _FlushBarrier = payload + if not barrier.canceled: + self._flush() + barrier.complete() + + def _export_one(self, record: dict[str, Any]) -> None: + exporter = self._exporter + # Copy for exporter isolation: every lane shares the same canonical + # record, and truncation/export must never mutate what another lane + # sees. If the copy fails we must NOT fall back to the shared record -- + # exporting the alias would let this lane's truncation mutate the object + # other lanes still read, breaking workflow isolation. Treat a copy + # failure like a render/truncation failure: log and skip this record for + # this lane, then continue processing the lane's queue. + try: + local = copy.deepcopy(record) + except Exception as exc: # noqa: BLE001 - a non-copyable payload must not alias the shared record or break the lane + _logger.warning( + "workflow-insight: record copy failed for exporter %s; " + "skipping export for this record: %s", + type(exporter).__name__, + exc, + ) + return + 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 len(self._pending) + + 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], + *, + max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, + ) -> None: + self._lanes = [ + _ExporterLane(exporter, max_pending_executions=max_pending_executions) + 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(execution_arn, 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 + 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..ac6e7dae 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,7 +33,6 @@ import datetime import json import math -import sys import threading from typing import Any, Callable @@ -47,10 +46,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, @@ -161,7 +160,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 +169,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): @@ -205,6 +208,12 @@ def __init__(self, config: WorkflowInsightConfig) -> None: self._exporters: list[InsightExporter] = ( list(config.exporters) if config.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() @@ -239,6 +248,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 +281,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 +291,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 +342,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 +354,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 +412,7 @@ def _build_operations( records.append(entry) return records - def _emit( + def _schedule_record( self, execution_arn: str, state: _ExecutionState, @@ -383,6 +422,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 +498,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..ad1c421a 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,6 +11,7 @@ from __future__ import annotations +import math from dataclasses import dataclass, field from enum import StrEnum from typing import Any, Callable, Literal, Protocol @@ -106,6 +107,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 +125,44 @@ def __post_init__(self) -> None: object.__setattr__( self, "operation_detail", OperationDetail(self.operation_detail) ) + self._validate_exporters() + self._validate_export_timeout() + + def _validate_exporters(self) -> None: + # One background worker (lane) is created per configured exporter, and + # each exporter is assumed to be driven from exactly one lane. Passing + # the SAME instance twice would give one object two lanes racing to + # render/export/flush it, producing duplicate, timing-dependent exports + # and breaking the one-thread-per-distinct-instance safety contract. + # Reject it loudly at construction. Compare by object IDENTITY (``is``), + # never equality/hash: exporters need not be hashable or comparable, and + # two DISTINCT instances of the same class (e.g. two S3Exporters for + # different buckets) are valid and each gets its own lane. + seen: list[InsightExporter] = [] + for exporter in self.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) + + def _validate_export_timeout(self) -> None: + # A finite, strictly-positive number. Reject ``bool`` (a subtype of + # ``int`` that would silently mean 1s / disallowed 0s), NaN, +/-inf, zero + # and negatives -- an invalid timeout must fail loudly at construction, + # not silently disable or unbound the invocation-end drain. + 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__}" + ) + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError( + "export_timeout_seconds must be a finite number greater than " + f"zero, got {timeout!r}" + ) 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..2b13e3b3 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 @@ -120,3 +120,99 @@ 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]) +def test_export_timeout_accepts_finite_positive_numbers(value): + config = WorkflowInsightConfig(export_timeout_seconds=value) + assert config.export_timeout_seconds == 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"), + 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_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 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..374a9b83 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -0,0 +1,603 @@ +# 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 logging +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 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: + """A payload whose ``deepcopy`` raises, to force a per-record copy failure.""" + + 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) + + +def test_one_worker_per_exporter(): + base = _insight_thread_count() + 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 _wait_until(lambda: _insight_thread_count() - base == 2) + e1.release() + e2.release() + scheduler.end_invocation(5.0) + + +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_different_executions_isolated_and_fair(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) # a1 in flight + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) # queued: [B] + scheduler.schedule(ARN_B, _rec(ARN_B, "b2")) # coalesce B -> b2 + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) # queued: [B, A] + exporter.release() + # a1 (in flight) first, then FIFO fairness B before the re-added A, each + # carrying its latest coalesced value. + assert _wait_until(lambda: exporter.exported_values() == ["a1", "b2", "a2"]) + scheduler.end_invocation(5.0) + + +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_deepcopy_failure_skips_record_and_lane_continues(caplog): + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + # A record whose deepcopy raises must be skipped for this lane -- never + # exported by aliasing the shared object -- and the lane must keep draining. + bad = _rec(ARN_A, "bad") + bad["payload"] = _Uncopyable() + good = _rec(ARN_B, "good") + with caplog.at_level( + logging.WARNING, logger="aws_durable_execution_sdk_python_insight" + ): + scheduler.schedule(ARN_A, bad) # queued first: copy fails -> skipped + scheduler.schedule(ARN_B, good) # queued behind it: must still export + # The good record delivering proves the lane continued past the failure; + # a single-lane worker drains FIFO, so "bad" was processed (and skipped) + # before "good" ran. + assert _wait_until(lambda: exporter.exported_values() == ["good"]) + scheduler.end_invocation(5.0) + # The exporter was never called for the un-copyable record. + assert exporter.exported_values() == ["good"] + # The failure was logged through the module logger. + assert any( + "record copy failed" in record.getMessage() + for record in caplog.records + if record.name == "aws_durable_execution_sdk_python_insight" + ) + + +def test_deepcopy_failure_does_not_alias_shared_record(): + # Before the fix a copy failure aliased the shared record and passed it to + # truncate_record -> render, which could mutate the canonical object other + # lanes still read. With the fix the record is skipped before render, so it + # is never aliased or mutated in place. + class MutatingRenderExporter(RecordingExporter): + def render(self, record: dict[str, Any]) -> Any: + record["mutated"] = True # would corrupt an aliased shared record + return record + + exporter = MutatingRenderExporter() + scheduler = _ExportScheduler([exporter]) + bad = _rec(ARN_A, "bad") + bad["payload"] = _Uncopyable() + scheduler.schedule(ARN_A, bad) + # A good record behind it lets us deterministically wait for the lane to + # drain past the bad one (single lane drains FIFO). + scheduler.schedule(ARN_B, _rec(ARN_B, "good")) + assert _wait_until(lambda: exporter.exported_values() == ["good"]) + scheduler.end_invocation(5.0) + # render never ran on the un-copyable record, so the canonical object was + # neither aliased into export nor mutated in place. + assert "mutated" not in bad + assert exporter.exported_values() == ["good"] + + +# -- 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 + + +# -- pending cap / cancelled barrier cleanup --------------------------------- + + +def test_pending_execution_cap_evicts_oldest(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_executions=2) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) # a1 in flight (not pending) + # Three distinct pending executions with cap 2 -> oldest (B) is evicted. + 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 _wait_until(lambda: lane._pending_count() == 2) + exporter.release() + scheduler.end_invocation(5.0) + + +def test_cancelled_barrier_is_cleaned_up_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) # times out -> barrier cancelled + assert ok is False + # Once the exporter unblocks, the worker drains the cancelled barrier + # (skipping the pointless flush) and exits idle -- no permanent leak. + exporter.release() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.flushed == 0 # cancelled barrier did not flush + + +def test_repeated_timeouts_behind_blocked_exporter_stay_bounded(): + """A blocked exporter across many warm invocations must not accumulate + barriers or grow queue state, must keep the SAME worker (no replacement), + must not execute any cancelled flush, and must drain + exit after release. + """ + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + + # First record puts the single worker into a blocked export. + 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() + + # Many warm invocations. Each schedules a coalescing record for the same + # ARN then ends with a short timeout; the barrier always times out because + # the worker is still stuck in the first export. + for i in range(50): + scheduler.schedule(ARN_A, _rec(ARN_A, f"a{i + 2}")) + ok = scheduler.end_invocation(0.02) + assert ok is False # degraded every time -- worker is blocked + # The cancelled barrier is pulled from the queue immediately, so no + # _FLUSH marker lingers behind the blocked worker. + assert lane._queued_flush_count() == 0 + # Queue holds at most the single coalesced record token; it never grows. + assert lane._queue_len() <= 1 + + # Bounded state: one in-flight ARN coalesced to a single pending record, and + # no growing pile of barriers. + assert lane._queue_len() <= 1 + assert lane._pending_count() <= 1 + assert lane._queued_flush_count() == 0 + # The blocked worker was never replaced. + assert lane._worker is worker + assert worker.is_alive() + assert _lane_worker_count(lane) == 1 + # No cancelled flush ran while the worker was blocked. + assert exporter.flushed == 0 + + # Release: the worker drains the latest coalesced record, then exits idle. + exporter.release() + assert _wait_until(lambda: not lane._worker_alive()) + exported = exporter.exported_values() + assert exported[0] == "a1" # the in-flight record delivered first + assert len(exported) <= 2 # a1 plus at most one final coalesced record + # Cancelled barriers never triggered a flush, and the idle-stop path does + # not flush either. + assert exporter.flushed == 0 + + +def test_cancel_flush_removes_queued_barrier_immediately(): + """Queued-barrier race: while the worker is blocked the barrier is still in + the queue, so cancel_flush pulls it out and completes it synchronously -- + without waiting for the worker and without ever flushing.""" + 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 blocked in export + barrier = lane.enqueue_flush() + assert lane._queued_flush_count() == 1 + lane.cancel_flush(barrier) + # Removed from the queue and completed here, without the worker. + assert lane._queued_flush_count() == 0 + assert barrier.canceled is True + assert barrier.is_done() + # Finish the in-flight export and go idle; the pulled barrier never flushed. + exporter.release() + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.flushed == 0 + 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()) 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..103b4982 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py @@ -0,0 +1,362 @@ +# 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 + +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 + + +# -- 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 == []