From 852c880cc3a65b6623b799ef4cc6589f466b041e Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Tue, 1 Sep 2026 23:08:13 +0000 Subject: [PATCH 01/23] feat(insight): async coalescing export scheduler Move all exporter work off the SDK checkpoint thread. A new private _ExportScheduler owns one lazy daemon worker per exporter lane; per-exporter copy, render, truncation, export() and flush() now run there, so a slow exporter never blocks workflow progress. Per lane: at most one in-flight record and one latest pending record per execution ARN. Cumulative snapshots for the same ARN coalesce (the in-flight record is never cancelled); updating a pending ARN moves it to the back for FIFO fairness across ARNs; pending ARNs are capped with oldest-eviction. A blocked worker is retained and never replaced, and idle workers exit after the drain, so threads cannot grow unbounded. on_operation_change returns immediately unless emit mode is on-change. Invocation end schedules the final record, then drains and flushes the touched lanes under one shared deadline; on timeout the workflow response is returned and delivery degrades to best-effort. Exceptions in render/export/flush are isolated and logged. Add WorkflowInsightConfig.export_timeout_seconds (default 5.0), validated as a finite number greater than zero (rejects bool, NaN, infinity, and non-positive values). Add scheduler, plugin-async, and config unit tests plus updated on-change coalescing coverage; refresh the README note. No core SDK changes. --- .../README.md | 13 +- .../_export_scheduler.py | 308 ++++++++++++++ .../plugin.py | 97 +++-- .../types.py | 23 ++ .../tests/test_config.py | 37 ++ .../tests/test_export_scheduler.py | 388 ++++++++++++++++++ .../tests/test_plugin.py | 23 +- .../tests/test_plugin_async.py | 212 ++++++++++ 8 files changed, 1055 insertions(+), 46 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py create mode 100644 packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py create mode 100644 packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index ed47c053..6ea44943 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -55,10 +55,15 @@ 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. 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..d6f27c61 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -0,0 +1,308 @@ +# 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, stale + barriers are cancelled, and any blocked worker stays daemonized (a synchronous + Python ``export()`` cannot be safely killed). +* 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) + self._cond = threading.Condition() + # 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() + # The arn whose record was popped and is being exported right now. + self._inflight_arn: str | None = None + 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() + + # -- 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 + self._inflight_arn = payload + + if kind == _RECORD and record is not None: + self._export_one(record) + with self._cond: + self._inflight_arn = None + 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: two lanes share the same canonical record, + # and truncation/export must never mutate what another lane sees. + try: + local = copy.deepcopy(record) + except Exception: # noqa: BLE001 - a non-copyable payload must not break export + 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 len(self._pending) + + +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.enqueue_flush() for lane in self._lanes] + deadline = time.monotonic() + timeout_seconds + degraded = False + for barrier in barriers: + remaining = deadline - time.monotonic() + if not barrier.wait(remaining): + barrier.canceled = True + 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..e469d43c 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() @@ -257,7 +266,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 +276,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 @@ -314,7 +330,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # 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( + self._schedule_record( arn, state, status=status, @@ -323,6 +339,13 @@ 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. + if state.scheduled: + 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 +396,7 @@ def _build_operations( records.append(entry) return records - def _emit( + def _schedule_record( self, execution_arn: str, state: _ExecutionState, @@ -383,6 +406,29 @@ 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) + state.scheduled = True + + 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 +480,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..f59e5159 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,20 @@ def __post_init__(self) -> None: object.__setattr__( self, "operation_detail", OperationDetail(self.operation_detail) ) + self._validate_export_timeout() + + 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..dc5d85a2 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,40 @@ 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) 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..8717c58e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -0,0 +1,388 @@ +# 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") + ) + + +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 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") + + +# -- 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) + + +# -- 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]) + 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 + + +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() + + +# -- worker lifecycle --------------------------------------------------------- + + +def test_blocked_worker_is_not_replaced(): + base = _insight_thread_count() + 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 + # 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) + assert lane._worker is worker + assert _insight_thread_count() - base == 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 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..aad2def3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py @@ -0,0 +1,212 @@ +# 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" +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, + ) + + +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" From 8fd62942d7a17166a1a98b747082267b7abc1cb3 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 1 Sep 2026 23:23:24 +0000 Subject: [PATCH 02/23] fix(insight): preserve exporter isolation Skip a lane record when copy.deepcopy fails instead of aliasing the shared canonical record. The alias let this lane's truncation mutate the object other lanes still read, breaking workflow isolation. A copy failure is now logged through the module logger and the lane keeps draining, matching render/truncation failure handling. Also remove the dead _inflight_arn lane field (written, never read). Tests: deepcopy-failure skips the record, does not call the exporter, logs the failure, and the lane continues to export a later valid record; a non-aliasing regression guards in-place mutation; a warm-container cross-invocation test proves bounded invocation-end waits, no A/B merge, and FIFO drain + flush after unblock. --- .../_export_scheduler.py | 24 +++-- .../tests/test_export_scheduler.py | 65 +++++++++++ .../tests/test_plugin_async.py | 102 ++++++++++++++++++ 3 files changed, 182 insertions(+), 9 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index d6f27c61..df673b82 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -102,8 +102,6 @@ def __init__( # 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() - # The arn whose record was popped and is being exported right now. - self._inflight_arn: str | None = None self._stop_when_idle = False self._worker: threading.Thread | None = None @@ -201,12 +199,9 @@ def _run_worker(self) -> None: record = self._pending.pop(payload, None) if record is None: continue - self._inflight_arn = payload if kind == _RECORD and record is not None: self._export_one(record) - with self._cond: - self._inflight_arn = None else: # _FLUSH barrier: _FlushBarrier = payload if not barrier.canceled: @@ -215,12 +210,23 @@ def _run_worker(self) -> None: def _export_one(self, record: dict[str, Any]) -> None: exporter = self._exporter - # Copy for exporter isolation: two lanes share the same canonical record, - # and truncation/export must never mutate what another lane sees. + # 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: # noqa: BLE001 - a non-copyable payload must not break export - local = 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 diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 8717c58e..d031e94e 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,6 +11,7 @@ from __future__ import annotations +import logging import threading import time from typing import Any @@ -121,6 +122,13 @@ def flush(self) -> None: 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 -------------------------- @@ -210,6 +218,63 @@ def test_terminal_record_supersedes_pending_running(): 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 ----------------------- 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 index aad2def3..ad2675e8 100644 --- 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 @@ -38,6 +38,8 @@ 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) @@ -96,6 +98,30 @@ def _end(operations: dict[str, OperationInfo]) -> InvocationEndInfo: ) +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 @@ -210,3 +236,79 @@ def test_buffered_exporter_publishes_after_flush(): 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 From cde166dd03046d444b09ee65a6ba8050c6b0ffa4 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 1 Sep 2026 23:35:34 +0000 Subject: [PATCH 03/23] test(insight): harden worker lifecycle checks Make the two shared-timeout tests wait deterministically for their released lane workers to stop before returning, so their daemon workers cannot exit between a later test's baseline capture and its assertion. Replace the fragile process-global thread-count delta in test_blocked_worker_is_not_replaced with lane-local worker identity, aliveness, and a lane-scoped worker count. This proves the blocked lane never spawns a replacement without depending on global thread state. Product code is unchanged. --- .../tests/test_export_scheduler.py | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index d031e94e..01365072 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -46,6 +46,17 @@ def _insight_thread_count() -> int: ) +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).""" @@ -350,6 +361,7 @@ def test_export_and_flush_exceptions_are_isolated(): 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() @@ -358,6 +370,9 @@ def test_shared_timeout_bounds_invocation_end_delay(): 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(): @@ -373,26 +388,38 @@ def test_shared_timeout_across_multiple_lanes_is_not_additive(): 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(): - base = _insight_thread_count() 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 _insight_thread_count() - base == 1 + assert worker.is_alive() + assert _lane_worker_count(lane) == 1 exporter.release() From 9cc5446f9bce72a901c0b48c3fd75086fea2af9c Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 2 Sep 2026 20:53:36 +0000 Subject: [PATCH 04/23] fix(insight): clean timed-out barriers Cancelled flush barriers no longer pile up behind a blocked exporter. end_invocation now pairs each barrier with its lane and, on timeout, calls _ExporterLane.cancel_flush(barrier): under the lane lock it marks the barrier cancelled and pulls its still-queued _FLUSH marker out, completing it there. If the worker already popped the marker the flush is left to the worker; an in-flight synchronous flush is not killed. This keeps queue and barrier state bounded across many warm invocations while preserving record ordering, normal flush, the shared deadline, blocked-worker retention, and bounded pending state. Also switch the lane Condition from the default RLock to an explicit non-reentrant Lock; the lane never re-acquires _cond while holding it. Tests: deterministic repeated-timeout test (blocked exporter across many warm invocations) plus queued-vs-already-popped cancellation race tests. --- .../_export_scheduler.py | 57 +++++++- .../tests/test_export_scheduler.py | 123 ++++++++++++++++++ 2 files changed, 173 insertions(+), 7 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index df673b82..4e2e97b9 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -24,9 +24,11 @@ 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, stale - barriers are cancelled, and any blocked worker stays daemonized (a synchronous - Python ``export()`` cannot be safely killed). + 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. """ @@ -96,7 +98,13 @@ def __init__( ) -> None: self._exporter = exporter self._max_pending = max(1, max_pending_executions) - self._cond = threading.Condition() + # 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 @@ -136,6 +144,30 @@ def request_stop_when_idle(self) -> None: 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: @@ -267,6 +299,14 @@ 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.""" @@ -295,13 +335,16 @@ def end_invocation(self, timeout_seconds: float) -> bool: to stop once idle. Returns ``True`` if every barrier completed within the deadline, ``False`` if delivery degraded to best-effort on timeout. """ - barriers = [lane.enqueue_flush() for lane in self._lanes] + barriers = [(lane, lane.enqueue_flush()) for lane in self._lanes] deadline = time.monotonic() + timeout_seconds degraded = False - for barrier in barriers: + for lane, barrier in barriers: remaining = deadline - time.monotonic() if not barrier.wait(remaining): - barrier.canceled = True + # 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() diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 01365072..374a9b83 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -113,6 +113,27 @@ def exported_values(self) -> list[Any]: 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.""" @@ -478,3 +499,105 @@ def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): 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()) From ca425b6312c9051fbe434fc4e0076d75cc9c1230 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 2 Sep 2026 23:34:29 +0000 Subject: [PATCH 05/23] fix(insight): validate exporters, lock scheduled Reject the same exporter instance appearing more than once in WorkflowInsightConfig.exporters with a clear ValueError, compared by object identity (not equality/hash) during config normalization. Two distinct instances of the same class stay valid and each keeps its own lane; the default exporter is unaffected. Preserves the one-thread-per-distinct-instance safety and avoids duplicate, timing-dependent scheduling. Route _ExecutionState.scheduled mutation and read through the plugin _lock via _mark_scheduled/_was_scheduled, consistent with the other state fields. The lock is released before any scheduler/end_invocation or exporter work, so no new lock ordering or deadlock is introduced. Add tests: same instance twice raises; two distinct same-class instances each get a lane; default exporter unaffected; scheduled flag tracks scheduling. --- .../README.md | 13 ++-- .../plugin.py | 22 ++++++- .../types.py | 24 ++++++++ .../tests/test_config.py | 59 +++++++++++++++++++ .../tests/test_plugin_async.py | 48 +++++++++++++++ 5 files changed, 160 insertions(+), 6 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 6ea44943..1dca3682 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -58,10 +58,15 @@ Behavior is validated cross-SDK by the `insight` conformance suite > **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. 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 +> 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. 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 e469d43c..be1a93c8 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 @@ -248,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: @@ -343,7 +358,8 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # 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. - if state.scheduled: + # 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 @@ -417,7 +433,9 @@ def _schedule_record( # 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) - state.scheduled = True + # 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, 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 f59e5159..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 @@ -125,8 +125,32 @@ 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 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 dc5d85a2..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 @@ -157,3 +157,62 @@ def test_export_timeout_accepts_finite_positive_numbers(value): 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_plugin_async.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin_async.py index ad2675e8..103b4982 100644 --- 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 @@ -312,3 +312,51 @@ def test_warm_container_cross_invocation_isolation_and_ordering(): # 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 == [] From 653a861a154941ae86676855d70d20f72c4d8290 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 3 Sep 2026 21:33:55 +0000 Subject: [PATCH 06/23] docs(insight): fix scheduler comment --- .../src/aws_durable_execution_sdk_python_insight/plugin.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index be1a93c8..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 @@ -342,9 +342,9 @@ 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. + # 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, From 8698f575e8105e314f81eafee6bc227a58249a06 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 8 Sep 2026 21:10:28 +0000 Subject: [PATCH 07/23] fix(insight): bound export timeout --- .../types.py | 28 ++++++++++++++----- .../tests/test_config.py | 8 ++++-- .../tests/test_plugin_async.py | 14 ++++++++++ 3 files changed, 41 insertions(+), 9 deletions(-) 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 ad1c421a..edef7ac1 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 @@ -12,6 +12,7 @@ from __future__ import annotations import math +import threading from dataclasses import dataclass, field from enum import StrEnum from typing import Any, Callable, Literal, Protocol @@ -152,17 +153,30 @@ def _validate_exporters(self) -> None: 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. + # 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__}" ) - if not math.isfinite(timeout) or timeout <= 0: + try: + normalized = float(timeout) + except (OverflowError, TypeError, ValueError) as exc: raise ValueError( - "export_timeout_seconds must be a finite number greater than " - f"zero, got {timeout!r}" + "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 2b13e3b3..561bea15 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,6 +11,8 @@ from __future__ import annotations +import threading + import pytest from aws_durable_execution_sdk_python_insight import ( @@ -131,10 +133,10 @@ def test_export_timeout_defaults_to_five_seconds(): assert workflow_insight(config)._export_timeout == 5.0 -@pytest.mark.parametrize("value", [0.1, 1, 2.5, 30]) +@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 == value + assert config.export_timeout_seconds == float(value) assert workflow_insight(config)._export_timeout == float(value) @@ -148,6 +150,8 @@ def test_export_timeout_accepts_finite_positive_numbers(value): 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 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 index 103b4982..250f7239 100644 --- 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 @@ -17,6 +17,8 @@ import time from typing import Any +import pytest + from aws_durable_execution_sdk_python.lambda_service import ( OperationStatus, OperationSubType, @@ -184,6 +186,18 @@ def test_non_on_change_mode_skips_operation_change_work(): 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() + + # -- a blocked exporter never blocks a hook ---------------------------------- From 7ceb4226866b3e54e607622354a12c58ac43778f Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 8 Sep 2026 21:38:12 +0000 Subject: [PATCH 08/23] fix(insight): revalidate exporter snapshot --- .../plugin.py | 12 ++++-- .../types.py | 43 ++++++++----------- .../tests/test_config.py | 31 +++++++++++++ 3 files changed, 57 insertions(+), 29 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index ac6e7dae..608fb24c 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 @@ -57,6 +57,7 @@ OperationDetail, OperationOverride, WorkflowInsightConfig, + _validate_exporter_instances, ) @@ -202,11 +203,14 @@ 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) 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 edef7ac1..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 @@ -15,7 +15,7 @@ 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): @@ -66,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``. @@ -126,32 +142,9 @@ def __post_init__(self) -> None: object.__setattr__( self, "operation_detail", OperationDetail(self.operation_detail) ) - self._validate_exporters() + _validate_exporter_instances(self.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 within Event.wait's platform bound. # Reject ``bool`` (an ``int`` subtype), values that overflow float 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 561bea15..2423d710 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 @@ -220,3 +220,34 @@ def test_default_exporter_unaffected_by_instance_check(): 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] From e6f101650c673b6f7e22f734d2003911bdc207de Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 20:50:21 +0000 Subject: [PATCH 09/23] docs(insight): shorten scheduler overview --- .../_export_scheduler.py | 45 +++++++------------ 1 file changed, 15 insertions(+), 30 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 4e2e97b9..006c19a4 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -1,36 +1,21 @@ # 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. +"""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 the latest pending snapshot per execution ARN and processes ARNs + round-robin. +* Drops the oldest pending execution when the lane-wide limit is reached, + keeping memory bounded. +* 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 From 603743098a751bab770351ace63cbb734f926782 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 22:05:52 +0000 Subject: [PATCH 10/23] test(insight): scope worker counts to lanes --- .../tests/test_export_scheduler.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 374a9b83..5244e8ff 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -180,18 +180,22 @@ def test_worker_created_lazily_on_first_schedule(): 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(): - 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) + 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_repeated_scheduling_does_not_grow_threads(): From bf9071a2e1c9e96b6ee6808d6fe49d3450c66983 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 23:30:18 +0000 Subject: [PATCH 11/23] fix(insight): enforce exporter lane bounds --- .../README.md | 10 +- .../_export_scheduler.py | 92 +++++++++++++++---- .../plugin.py | 28 ++++++ .../tests/test_config.py | 23 +++++ .../tests/test_export_scheduler.py | 26 +++++- 5 files changed, 154 insertions(+), 25 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 1dca3682..070293e1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -59,10 +59,12 @@ Behavior is validated cross-SDK by the `insight` conformance suite > 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 +> 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 +> 1,024 pending executions and 16 MB of estimated canonical JSON; it drops the +> oldest pending snapshot when either bound is reached. 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 diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 006c19a4..fcefc14e 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -11,8 +11,8 @@ * Keeps the latest pending snapshot per execution ARN and processes ARNs round-robin. -* Drops the oldest pending execution when the lane-wide limit is reached, - keeping memory bounded. +* Drops the oldest pending snapshot when the execution-count or byte budget is + reached, keeping memory bounded. * 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. @@ -27,7 +27,10 @@ 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.truncation import ( + json_byte_size, + truncate_record, +) from aws_durable_execution_sdk_python_insight.types import InsightExporter @@ -39,6 +42,12 @@ # bounded regardless of how long a worker stays blocked. _DEFAULT_MAX_PENDING_EXECUTIONS = 1024 +# Canonical JSON-byte estimate retained by one blocked lane. Lambda functions +# can be configured with 128 MiB, so keep the instrumentation backlog well below +# that floor. Python object overhead is higher than JSON bytes; this is a +# conservative budget signal, not an exact heap measurement. +_DEFAULT_MAX_PENDING_BYTES = 16_000_000 + # Queue entry kinds. _RECORD = "record" _FLUSH = "flush" @@ -80,9 +89,12 @@ def __init__( exporter: InsightExporter, *, max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, + max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, ) -> None: self._exporter = exporter self._max_pending = max(1, max_pending_executions) + self._max_pending_bytes = max(1, max_pending_bytes) + self._pending_bytes = 0 # 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 @@ -92,27 +104,42 @@ def __init__( 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() + # arn -> (latest pending record, canonical JSON-byte estimate). Insertion + # order is both record age and fairness order because replacing an ARN + # moves it to the back. + self._pending: OrderedDict[str, tuple[dict[str, Any], int]] = 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: + def schedule( + self, + execution_arn: str, + record: dict[str, Any], + record_size: int | None, + ) -> None: with self._cond: self._stop_when_idle = False + size = ( + self._max_pending_bytes + 1 + if record_size is None + else max(0, record_size) + ) 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 + _, old_size = self._pending[execution_arn] + self._pending_bytes -= old_size + self._pending[execution_arn] = (record, size) + self._pending_bytes += size self._pending.move_to_end(execution_arn) self._move_record_token_to_back(execution_arn) else: - self._pending[execution_arn] = record + self._pending[execution_arn] = (record, size) + self._pending_bytes += size self._queue.append((_RECORD, execution_arn)) - self._enforce_pending_cap() + self._enforce_pending_caps() self._ensure_worker_locked() self._cond.notify() @@ -165,17 +192,32 @@ def _move_record_token_to_back(self, execution_arn: str) -> None: # 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: + def _enforce_pending_caps(self) -> None: while len(self._pending) > self._max_pending: - old_arn, _ = self._pending.popitem(last=False) - self._remove_record_token(old_arn) + old_arn, _ = self._drop_oldest_pending() _logger.warning( - "workflow-insight: export lane for %s is full " - "(cap=%d); dropping pending record for %s", + "workflow-insight: export lane for %s reached its execution cap " + "(%d); dropping pending record for %s", type(self._exporter).__name__, self._max_pending, old_arn, ) + while self._pending_bytes > self._max_pending_bytes and self._pending: + old_arn, dropped_size = self._drop_oldest_pending() + _logger.warning( + "workflow-insight: export lane for %s reached its pending byte " + "budget (%d); dropping %d-byte pending record for %s", + type(self._exporter).__name__, + self._max_pending_bytes, + dropped_size, + old_arn, + ) + + def _drop_oldest_pending(self) -> tuple[str, int]: + old_arn, (_, old_size) = self._pending.popitem(last=False) + self._pending_bytes -= old_size + self._remove_record_token(old_arn) + return old_arn, old_size def _remove_record_token(self, execution_arn: str) -> None: for index, (kind, payload) in enumerate(self._queue): @@ -213,9 +255,11 @@ def _run_worker(self) -> None: kind, payload = self._queue.popleft() record: dict[str, Any] | None = None if kind == _RECORD: - record = self._pending.pop(payload, None) - if record is None: + pending = self._pending.pop(payload, None) + if pending is None: continue + record, record_size = pending + self._pending_bytes -= record_size if kind == _RECORD and record is not None: self._export_one(record) @@ -284,6 +328,10 @@ def _pending_count(self) -> int: with self._cond: return len(self._pending) + def _pending_bytes_count(self) -> int: + with self._cond: + return self._pending_bytes + def _queue_len(self) -> int: with self._cond: return len(self._queue) @@ -301,16 +349,22 @@ def __init__( exporters: list[InsightExporter], *, max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, + max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, ) -> None: self._lanes = [ - _ExporterLane(exporter, max_pending_executions=max_pending_executions) + _ExporterLane( + exporter, + max_pending_executions=max_pending_executions, + max_pending_bytes=max_pending_bytes, + ) 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.""" + record_size = json_byte_size(record) for lane in self._lanes: - lane.schedule(execution_arn, record) + lane.schedule(execution_arn, record, record_size) def end_invocation(self, timeout_seconds: float) -> bool: """Drain and flush every touched lane under one shared timeout. 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 608fb24c..a07fc26a 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 @@ -34,6 +34,7 @@ import json import math import threading +import weakref from typing import Any, Callable from aws_durable_execution_sdk_python.plugin import ( @@ -71,6 +72,32 @@ InvocationStatus.RETRY: "RUNNING", } +_exporter_owner_lock = threading.Lock() +_exporter_owners: list[tuple[InsightExporter, weakref.ReferenceType[object]]] = [] + + +def _claim_exporters(owner: object, exporters: list[InsightExporter]) -> None: + """Give each exporter object to at most one live plugin instance.""" + + def release(owner_ref: weakref.ReferenceType[object]) -> None: + with _exporter_owner_lock: + _exporter_owners[:] = [ + entry for entry in _exporter_owners if entry[1] is not owner_ref + ] + + owner_ref = weakref.ref(owner, release) + with _exporter_owner_lock: + _exporter_owners[:] = [ + entry for entry in _exporter_owners if entry[1]() is not None + ] + for exporter in exporters: + if any(existing is exporter for existing, _ in _exporter_owners): + raise ValueError( + "the same exporter instance cannot be shared across " + "Workflow Insight plugin instances" + ) + _exporter_owners.extend((exporter, owner_ref) for exporter in exporters) + def _parse_execution_arn(execution_arn: str) -> dict[str, str]: # arn::lambda:::function::/durable-execution// @@ -220,6 +247,7 @@ def __init__(self, config: WorkflowInsightConfig) -> None: self._scheduler = _ExportScheduler(self._exporters) self._state: dict[str, _ExecutionState] = {} self._lock = threading.Lock() + _claim_exporters(self, self._exporters) # -- sampling / state ----------------------------------------------------- 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 2423d710..a98385d0 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,6 +11,7 @@ from __future__ import annotations +import gc import threading import pytest @@ -214,6 +215,28 @@ def test_two_distinct_same_class_instances_accepted_each_own_lane(): 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_default_exporter_unaffected_by_instance_check(): # No exporters configured -> single default LambdaLogExporter, one lane; the # identity check never trips on the empty list. diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 5244e8ff..ca2d1bf0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -154,8 +154,11 @@ def flush(self) -> None: raise RuntimeError("flush boom") -class _Uncopyable: - """A payload whose ``deepcopy`` raises, to force a per-record copy failure.""" +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") @@ -490,6 +493,25 @@ def test_pending_execution_cap_evicts_oldest(): scheduler.end_invocation(5.0) +def test_pending_byte_budget_evicts_oldest_large_record(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 1_500)) + scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 1_500)) + + assert lane._pending_count() == 1 + assert lane._pending_bytes_count() <= 2_000 + exporter.release() + scheduler.end_invocation(5.0) + exported = exporter.exported_values() + assert exported[0] == "inflight" + assert exported[1] == "c" * 1_500 + + def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From 399a606df1f7d1f6a9222a5459616e370bf31f56 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 23:59:49 +0000 Subject: [PATCH 12/23] fix(insight): harden lane admission and ownership --- .../_export_scheduler.py | 24 ++++++-- .../plugin.py | 17 +++--- .../truncation.py | 2 +- .../tests/test_config.py | 32 +++++++++++ .../tests/test_export_scheduler.py | 56 +++++++++++++++++++ 5 files changed, 117 insertions(+), 14 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index fcefc14e..9a20af52 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -121,11 +121,25 @@ def schedule( ) -> None: with self._cond: self._stop_when_idle = False - size = ( - self._max_pending_bytes + 1 - if record_size is None - else max(0, record_size) - ) + if record_size is None: + _logger.warning( + "workflow-insight: cannot measure pending record for %s on " + "%s; dropping this record", + execution_arn, + type(self._exporter).__name__, + ) + return + size = max(0, record_size) + if size > self._max_pending_bytes: + _logger.warning( + "workflow-insight: pending record for %s on %s exceeds the " + "byte budget (%d > %d); dropping this record", + execution_arn, + type(self._exporter).__name__, + size, + self._max_pending_bytes, + ) + return 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. 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 a07fc26a..78082cc6 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 @@ -76,27 +76,28 @@ _exporter_owners: list[tuple[InsightExporter, weakref.ReferenceType[object]]] = [] -def _claim_exporters(owner: object, exporters: list[InsightExporter]) -> None: - """Give each exporter object to at most one live plugin instance.""" +def _claim_exporter_lanes(lanes: list[Any]) -> None: + """Give each exporter object to at most one live scheduler lane.""" - def release(owner_ref: weakref.ReferenceType[object]) -> None: + def release(lane_ref: weakref.ReferenceType[object]) -> None: with _exporter_owner_lock: _exporter_owners[:] = [ - entry for entry in _exporter_owners if entry[1] is not owner_ref + entry for entry in _exporter_owners if entry[1] is not lane_ref ] - owner_ref = weakref.ref(owner, release) + claims = [(lane._exporter, lane) for lane in lanes] with _exporter_owner_lock: _exporter_owners[:] = [ entry for entry in _exporter_owners if entry[1]() is not None ] - for exporter in exporters: + for exporter, _ in claims: if any(existing is exporter for existing, _ in _exporter_owners): raise ValueError( "the same exporter instance cannot be shared across " "Workflow Insight plugin instances" ) - _exporter_owners.extend((exporter, owner_ref) for exporter in exporters) + for exporter, lane in claims: + _exporter_owners.append((exporter, weakref.ref(lane, release))) def _parse_execution_arn(execution_arn: str) -> dict[str, str]: @@ -247,7 +248,7 @@ def __init__(self, config: WorkflowInsightConfig) -> None: self._scheduler = _ExportScheduler(self._exporters) self._state: dict[str, _ExecutionState] = {} self._lock = threading.Lock() - _claim_exporters(self, self._exporters) + _claim_exporter_lanes(self._scheduler._lanes) # -- sampling / state ----------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py index d826d2e0..aee615b6 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py @@ -28,7 +28,7 @@ def json_byte_size(value: Any) -> int | None: return len( json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode("utf-8") ) - except (TypeError, ValueError): + except Exception: # noqa: BLE001 - sizing failure must never break instrumentation return None 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 a98385d0..3bba291e 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 @@ -13,6 +13,7 @@ import gc import threading +import time import pytest @@ -237,6 +238,37 @@ def test_exporter_instance_can_be_reused_after_owner_is_collected(): 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_default_exporter_unaffected_by_instance_check(): # No exporters configured -> single default LambdaLogExporter, one lane; the # identity check never trips on the empty list. diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index ca2d1bf0..f770cc99 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,6 +11,7 @@ from __future__ import annotations +import json import logging import threading import time @@ -512,6 +513,61 @@ def test_pending_byte_budget_evicts_oldest_large_record(): assert exported[1] == "c" * 1_500 +def test_unmeasurable_record_does_not_evict_existing_backlog(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) + scheduler.schedule(ARN_C, _rec(ARN_C, "c1")) + + unmeasurable = _rec(ARN_D, "bad") + unmeasurable["payload"] = {"not-json"} + scheduler.schedule(ARN_D, unmeasurable) + + assert lane._pending_count() == 2 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight", "b1", "c1"] + + +def test_individually_over_budget_record_does_not_evict_existing_backlog(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 700)) + scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 700)) + scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000)) + + assert lane._pending_count() == 2 + assert lane._pending_bytes_count() <= 2_000 + exporter.release() + scheduler.end_invocation(5.0) + exported = exporter.exported_values() + assert exported[0] == "inflight" + assert exported[1:] == ["b" * 700, "c" * 700] + + +def test_record_sizing_exception_does_not_escape_schedule(monkeypatch): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + def fail_sizing(*args, **kwargs): + raise RecursionError("record nesting is too deep") + + monkeypatch.setattr(json, "dumps", fail_sizing) + scheduler.schedule(ARN_B, _rec(ARN_B, "too-deep")) + + assert scheduler._lanes[0]._pending_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From cab5f526407f2b3ae911a66b66ba708a0c919dfb Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 00:45:52 +0000 Subject: [PATCH 13/23] fix(insight): contain scheduler admission failures --- .../README.md | 2 +- .../_export_scheduler.py | 94 +++++++++++++++---- .../tests/test_export_scheduler.py | 60 ++++++++---- .../tests/test_plugin_async.py | 17 ++++ 4 files changed, 134 insertions(+), 39 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 070293e1..17b5617b 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -63,7 +63,7 @@ Behavior is validated cross-SDK by the `insight` conformance suite > `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 -> 1,024 pending executions and 16 MB of estimated canonical JSON; it drops the +> 1,024 pending executions and 16 MB of estimated retained memory; it drops the > oldest pending snapshot when either bound is reached. 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 diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 9a20af52..96122e3a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -22,15 +22,13 @@ import copy import logging +import sys import threading import time from collections import OrderedDict, deque from typing import Any -from aws_durable_execution_sdk_python_insight.truncation import ( - json_byte_size, - truncate_record, -) +from aws_durable_execution_sdk_python_insight.truncation import truncate_record from aws_durable_execution_sdk_python_insight.types import InsightExporter @@ -42,12 +40,41 @@ # bounded regardless of how long a worker stays blocked. _DEFAULT_MAX_PENDING_EXECUTIONS = 1024 -# Canonical JSON-byte estimate retained by one blocked lane. Lambda functions +# Estimated Python object memory retained by one blocked lane. Lambda functions # can be configured with 128 MiB, so keep the instrumentation backlog well below -# that floor. Python object overhead is higher than JSON bytes; this is a -# conservative budget signal, not an exact heap measurement. +# that floor. This is a conservative budget signal, not an exact heap measurement. _DEFAULT_MAX_PENDING_BYTES = 16_000_000 + +def _estimate_retained_size(value: Any) -> int: + """Estimate retained Python memory without serializing or calling render().""" + total = 0 + seen: set[int] = set() + stack: list[Any] = [value] + while stack: + item = stack.pop() + identity = id(item) + if identity in seen: + continue + seen.add(identity) + try: + total += sys.getsizeof(item) + except Exception: # noqa: BLE001 - estimation must never break a hook + total += 1_024 + continue + if isinstance(item, dict): + stack.extend(item.keys()) + stack.extend(item.values()) + elif isinstance(item, (list, tuple, set, frozenset, deque)): + stack.extend(item) + else: + try: + stack.append(vars(item)) + except Exception: # noqa: BLE001 - custom objects are best-effort + pass + return total + + # Queue entry kinds. _RECORD = "record" _FLUSH = "flush" @@ -61,11 +88,12 @@ class _FlushBarrier: a later, still-blocked worker skips the now-pointless flush. """ - __slots__ = ("_event", "canceled") + __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() @@ -104,12 +132,13 @@ def __init__( 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, canonical JSON-byte estimate). Insertion + # arn -> (latest pending record, retained-memory estimate). Insertion # order is both record age and fairness order because replacing an ARN # moves it to the back. self._pending: OrderedDict[str, tuple[dict[str, Any], int]] = OrderedDict() self._stop_when_idle = False self._worker: threading.Thread | None = None + self._disabled = False # -- producer API (checkpoint / invocation-end threads) ------------------- @@ -117,18 +146,12 @@ def schedule( self, execution_arn: str, record: dict[str, Any], - record_size: int | None, + record_size: int, ) -> None: with self._cond: - self._stop_when_idle = False - if record_size is None: - _logger.warning( - "workflow-insight: cannot measure pending record for %s on " - "%s; dropping this record", - execution_arn, - type(self._exporter).__name__, - ) + if self._disabled: return + self._stop_when_idle = False size = max(0, record_size) if size > self._max_pending_bytes: _logger.warning( @@ -160,6 +183,11 @@ def schedule( 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() @@ -239,11 +267,32 @@ def _remove_record_token(self, execution_arn: str) -> None: del self._queue[index] return + def _disable_locked(self, exc: Exception) -> None: + self._disabled = True + self._worker = None + self._pending.clear() + self._pending_bytes = 0 + for kind, payload in self._queue: + if kind == _FLUSH: + 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, @@ -251,7 +300,10 @@ def _ensure_worker_locked(self) -> None: daemon=True, ) self._worker = worker - worker.start() + try: + worker.start() + except Exception as exc: # noqa: BLE001 - instrumentation must not break hooks + self._disable_locked(exc) # -- worker (single daemon thread) --------------------------------------- @@ -376,7 +428,7 @@ def __init__( def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: """Fan a canonical record out to every lane. Returns immediately.""" - record_size = json_byte_size(record) + record_size = _estimate_retained_size(record) for lane in self._lanes: lane.schedule(execution_arn, record, record_size) @@ -399,6 +451,8 @@ def end_invocation(self, timeout_seconds: float) -> bool: # 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: diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index f770cc99..d51ed0e0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,7 +11,6 @@ from __future__ import annotations -import json import logging import threading import time @@ -165,6 +164,13 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Any: raise RuntimeError("uncopyable payload") +class _Unsized: + """A payload whose custom ``__sizeof__`` raises.""" + + def __sizeof__(self) -> int: + raise RuntimeError("size unavailable") + + # -- lazy worker creation / one worker per exporter -------------------------- @@ -202,6 +208,25 @@ def test_one_worker_per_exporter(): ) +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._pending_bytes_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() @@ -496,7 +521,7 @@ def test_pending_execution_cap_evicts_oldest(): def test_pending_byte_budget_evicts_oldest_large_record(): exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000) + scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) lane = scheduler._lanes[0] scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) assert _wait_until(exporter.started.is_set) @@ -505,7 +530,7 @@ def test_pending_byte_budget_evicts_oldest_large_record(): scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 1_500)) assert lane._pending_count() == 1 - assert lane._pending_bytes_count() <= 2_000 + assert lane._pending_bytes_count() <= 3_000 exporter.release() scheduler.end_invocation(5.0) exported = exporter.exported_values() @@ -513,7 +538,7 @@ def test_pending_byte_budget_evicts_oldest_large_record(): assert exported[1] == "c" * 1_500 -def test_unmeasurable_record_does_not_evict_existing_backlog(): +def test_non_json_record_reaches_exporter_without_evicting_backlog(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) lane = scheduler._lanes[0] @@ -522,19 +547,19 @@ def test_unmeasurable_record_does_not_evict_existing_backlog(): scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) scheduler.schedule(ARN_C, _rec(ARN_C, "c1")) - unmeasurable = _rec(ARN_D, "bad") - unmeasurable["payload"] = {"not-json"} - scheduler.schedule(ARN_D, unmeasurable) + non_json = _rec(ARN_D, "custom") + non_json["payload"] = {"not-json"} + scheduler.schedule(ARN_D, non_json) - assert lane._pending_count() == 2 + assert lane._pending_count() == 3 exporter.release() scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight", "b1", "c1"] + assert exporter.exported_values() == ["inflight", "b1", "c1", "custom"] def test_individually_over_budget_record_does_not_evict_existing_backlog(): exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000) + scheduler = _ExportScheduler([exporter], max_pending_bytes=3_500) lane = scheduler._lanes[0] scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) assert _wait_until(exporter.started.is_set) @@ -543,7 +568,7 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog(): scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000)) assert lane._pending_count() == 2 - assert lane._pending_bytes_count() <= 2_000 + assert lane._pending_bytes_count() <= 3_500 exporter.release() scheduler.end_invocation(5.0) exported = exporter.exported_values() @@ -551,21 +576,20 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog(): assert exported[1:] == ["b" * 700, "c" * 700] -def test_record_sizing_exception_does_not_escape_schedule(monkeypatch): +def test_record_sizing_exception_does_not_escape_schedule(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) assert _wait_until(exporter.started.is_set) - def fail_sizing(*args, **kwargs): - raise RecursionError("record nesting is too deep") - - monkeypatch.setattr(json, "dumps", fail_sizing) - scheduler.schedule(ARN_B, _rec(ARN_B, "too-deep")) + record = _rec(ARN_B, "custom-sized") + record["payload"] = _Unsized() + scheduler.schedule(ARN_B, record) - assert scheduler._lanes[0]._pending_count() == 0 + assert scheduler._lanes[0]._pending_count() == 1 exporter.release() scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight", "custom-sized"] def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): 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 index 250f7239..e0a9c432 100644 --- 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 @@ -198,6 +198,23 @@ def test_oversized_timeout_rejected_before_exporter_worker_starts(): 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 ---------------------------------- From d4979550e449f960a83c282053aff37bf1be8667 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 01:22:00 +0000 Subject: [PATCH 14/23] fix(insight): preserve bounded custom records --- .../_export_scheduler.py | 77 ++++++++--- .../tests/test_export_scheduler.py | 125 ++++++++++++------ 2 files changed, 141 insertions(+), 61 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 96122e3a..51faa28d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -61,18 +61,54 @@ def _estimate_retained_size(value: Any) -> int: total += sys.getsizeof(item) except Exception: # noqa: BLE001 - estimation must never break a hook total += 1_024 + try: + if isinstance(item, dict): + stack.extend(item.keys()) + stack.extend(item.values()) + elif isinstance(item, (list, tuple, set, frozenset, deque)): + stack.extend(item) + else: + try: + stack.append(vars(item)) + except Exception: # noqa: BLE001 - custom objects may use slots + pass + for cls in type(item).__mro__: + slots = vars(cls).get("__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + for slot in slots: + if slot in {"__dict__", "__weakref__"}: + continue + if slot.startswith("__") and not slot.endswith("__"): + slot = f"_{cls.__name__.lstrip('_')}{slot}" + try: + stack.append(getattr(item, slot)) + except Exception: # noqa: BLE001 - unset/custom slots are best-effort + pass + except Exception: # noqa: BLE001 - traversal must never break a hook + pass + return total + + +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 - if isinstance(item, dict): + seen.add(identity) + if type(item) is dict: stack.extend(item.keys()) stack.extend(item.values()) - elif isinstance(item, (list, tuple, set, frozenset, deque)): + elif type(item) in {list, tuple, set, frozenset, deque}: stack.extend(item) else: - try: - stack.append(vars(item)) - except Exception: # noqa: BLE001 - custom objects are best-effort - pass - return total + memo[identity] = item + return copy.deepcopy(record, memo) # Queue entry kinds. @@ -154,13 +190,21 @@ def schedule( self._stop_when_idle = False size = max(0, record_size) if size > self._max_pending_bytes: + superseded = self._pending.pop(execution_arn, None) + if superseded is not None: + _, superseded_size = superseded + self._pending_bytes -= superseded_size + self._remove_record_token(execution_arn) _logger.warning( "workflow-insight: pending record for %s on %s exceeds the " - "byte budget (%d > %d); dropping this record", + "byte budget (%d > %d); dropping this record%s", execution_arn, type(self._exporter).__name__, size, self._max_pending_bytes, + " and its superseded pending snapshot" + if superseded is not None + else "", ) return if execution_arn in self._pending: @@ -337,18 +381,15 @@ def _run_worker(self) -> None: 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. + # 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.deepcopy(record) - except Exception as exc: # noqa: BLE001 - a non-copyable payload must not alias the shared record or break the lane + local = _copy_record_containers(record) + except Exception as exc: # noqa: BLE001 - malformed containers must not break the lane _logger.warning( - "workflow-insight: record copy failed for exporter %s; " + "workflow-insight: record container copy failed for exporter %s; " "skipping export for this record: %s", type(exporter).__name__, exc, diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index d51ed0e0..1686e0fd 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,7 +11,6 @@ from __future__ import annotations -import logging import threading import time from typing import Any @@ -164,6 +163,18 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Any: raise RuntimeError("uncopyable payload") +class _SlottedPayload: + __slots__ = ("payload",) + + def __init__(self, payload: Any) -> None: + self.payload = payload + + +class _UnsizedSlottedPayload(_SlottedPayload): + def __sizeof__(self) -> int: + raise RuntimeError("size unavailable") + + class _Unsized: """A payload whose custom ``__sizeof__`` raises.""" @@ -286,58 +297,45 @@ def test_terminal_record_supersedes_pending_running(): # -- copy failure isolation --------------------------------------------------- -def test_deepcopy_failure_skips_record_and_lane_continues(caplog): - exporter = RecordingExporter() +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]) - # 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" - ) + 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_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. + +def test_uncopyable_custom_value_does_not_alias_record_containers(): class MutatingRenderExporter(RecordingExporter): def render(self, record: dict[str, Any]) -> Any: - record["mutated"] = True # would corrupt an aliased shared record + record["mutated"] = True 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"]) + record = _rec(ARN_A, "custom") + record["payload"] = _Uncopyable() + + scheduler.schedule(ARN_A, record) 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"] + + assert "mutated" not in record + assert exporter.exported_values() == ["custom"] # -- non-blocking hook return / fast-vs-slow isolation ----------------------- @@ -576,6 +574,47 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog(): assert exported[1:] == ["b" * 700, "c" * 700] +def test_over_budget_replacement_removes_superseded_same_arn_only(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=3_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_A, _rec(ARN_A, "stale-running")) + scheduler.schedule(ARN_B, _rec(ARN_B, "unrelated")) + scheduler.schedule( + ARN_A, + _rec(ARN_A, "terminal" * 500, status="SUCCEEDED"), + ) + + assert lane._pending_count() == 1 + assert lane._pending_bytes_count() <= 3_500 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight", "unrelated"] + + +def test_retained_size_traverses_slots_after_shallow_size_failure(): + for payload in ( + _SlottedPayload("x" * 4_000), + _UnsizedSlottedPayload("x" * 4_000), + ): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + record = _rec(ARN_B, "opaque") + record["payload"] = payload + scheduler.schedule(ARN_B, record) + + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["inflight"] + + def test_record_sizing_exception_does_not_escape_schedule(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From b3c9f1a70602ae6c5ff26b61b374353c62183342 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 01:54:06 +0000 Subject: [PATCH 15/23] fix(insight): preserve timed-out buffered exports --- .../_export_scheduler.py | 48 +++--- .../tests/test_export_scheduler.py | 156 +++++++++++++----- 2 files changed, 137 insertions(+), 67 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 51faa28d..57b1cad0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -46,7 +46,7 @@ _DEFAULT_MAX_PENDING_BYTES = 16_000_000 -def _estimate_retained_size(value: Any) -> int: +def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: """Estimate retained Python memory without serializing or calling render().""" total = 0 seen: set[int] = set() @@ -61,8 +61,12 @@ def _estimate_retained_size(value: Any) -> int: total += sys.getsizeof(item) except Exception: # noqa: BLE001 - estimation must never break a hook total += 1_024 + if max_size is not None and total > max_size: + return max_size + 1 try: - if isinstance(item, dict): + if isinstance(item, memoryview): + stack.append(item.obj) + elif isinstance(item, dict): stack.extend(item.keys()) stack.extend(item.values()) elif isinstance(item, (list, tuple, set, frozenset, deque)): @@ -243,26 +247,19 @@ def request_stop_when_idle(self) -> None: 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. - """ + """Stop waiting for a timed-out barrier while retaining one later flush.""" with self._cond: barrier.canceled = True + # Keep at most one detached flush. Moving it to this barrier's + # position makes it cover all work scheduled before the latest + # timeout without accumulating one marker per warm invocation. + for index in range(len(self._queue) - 1, -1, -1): + kind, payload = self._queue[index] + if kind == _FLUSH and payload is None: + del self._queue[index] for index, (kind, payload) in enumerate(self._queue): if kind == _FLUSH and payload is barrier: - del self._queue[index] + self._queue[index] = (_FLUSH, None) barrier.complete() return @@ -317,7 +314,7 @@ def _disable_locked(self, exc: Exception) -> None: self._pending.clear() self._pending_bytes = 0 for kind, payload in self._queue: - if kind == _FLUSH: + if kind == _FLUSH and payload is not None: barrier: _FlushBarrier = payload barrier.canceled = True barrier.failed = True @@ -374,10 +371,10 @@ def _run_worker(self) -> None: 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() + 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 @@ -458,18 +455,19 @@ def __init__( max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, ) -> None: + self._max_pending_bytes = max(1, max_pending_bytes) self._lanes = [ _ExporterLane( exporter, max_pending_executions=max_pending_executions, - max_pending_bytes=max_pending_bytes, + max_pending_bytes=self._max_pending_bytes, ) 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.""" - record_size = _estimate_retained_size(record) + record_size = _estimate_retained_size(record, self._max_pending_bytes) for lane in self._lanes: lane.schedule(execution_arn, record, record_size) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 1686e0fd..e6486802 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -133,6 +133,35 @@ 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.""" @@ -182,6 +211,16 @@ def __sizeof__(self) -> int: raise RuntimeError("size unavailable") +class _TrackedLargeList(list[Any]): + def __init__(self) -> None: + super().__init__([None] * 10_000) + self.iterated = False + + def __iter__(self): + self.iterated = True + return super().__iter__() + + # -- lazy worker creation / one worker per exporter -------------------------- @@ -615,6 +654,42 @@ def test_retained_size_traverses_slots_after_shallow_size_failure(): assert exporter.exported_values() == ["inflight"] +def test_retained_size_saturates_before_traversing_large_shallow_container(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + payload = _TrackedLargeList() + record = _rec(ARN_B, "large-shallow") + record["payload"] = payload + scheduler.schedule(ARN_B, record) + + assert payload.iterated is False + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + +def test_retained_size_counts_memoryview_backing_buffer(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + record = _rec(ARN_B, "memoryview") + record["payload"] = memoryview(bytearray(4_000)) + scheduler.schedule(ARN_B, record) + + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + def test_record_sizing_exception_does_not_escape_schedule(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) @@ -631,93 +706,90 @@ def test_record_sizing_exception_does_not_escape_schedule(): assert exporter.exported_values() == ["inflight", "custom-sized"] -def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): +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) # times out -> barrier cancelled + ok = scheduler.end_invocation(0.1) assert ok is False - # Once the exporter unblocks, the worker drains the cancelled barrier - # (skipping the pointless flush) and exits idle -- no permanent leak. + # 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 == 0 # cancelled barrier did not flush + 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(): - """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. - """ + """Warm timeouts coalesce to one eventual flush on the same worker.""" 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 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() == 0 - # The blocked worker was never replaced. + assert lane._queued_flush_count() == 1 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 + assert exported[0] == "a1" + assert len(exported) <= 2 + assert exporter.flushed == 1 -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.""" +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) # worker blocked in export + assert _wait_until(exporter.started.is_set) 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 lane._queued_flush_count() == 1 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.flushed == 1 assert exporter.exported_values() == ["a1"] From 9a9cb0ef6bfa9e60f1f51f1a7adbb1afc61a38b6 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:14:32 +0000 Subject: [PATCH 16/23] fix(insight): bound opaque retained graphs --- .../_export_scheduler.py | 51 +++++++++---------- .../plugin.py | 27 +++++----- .../tests/test_config.py | 20 ++++++++ .../tests/test_export_scheduler.py | 41 ++++++++++++++- 4 files changed, 96 insertions(+), 43 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 57b1cad0..7fae566f 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -21,10 +21,12 @@ from __future__ import annotations import copy +import gc import logging import sys import threading import time +import types from collections import OrderedDict, deque from typing import Any @@ -46,8 +48,17 @@ _DEFAULT_MAX_PENDING_BYTES = 16_000_000 +_RETAINED_GRAPH_BOUNDARIES = ( + type, + types.ModuleType, + types.FunctionType, + types.BuiltinFunctionType, + types.CodeType, +) + + def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: - """Estimate retained Python memory without serializing or calling render().""" + """Estimate a bounded retained graph without serializing or calling render().""" total = 0 seen: set[int] = set() stack: list[Any] = [value] @@ -64,33 +75,17 @@ def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: if max_size is not None and total > max_size: return max_size + 1 try: - if isinstance(item, memoryview): - stack.append(item.obj) - elif isinstance(item, dict): - stack.extend(item.keys()) - stack.extend(item.values()) - elif isinstance(item, (list, tuple, set, frozenset, deque)): - stack.extend(item) - else: - try: - stack.append(vars(item)) - except Exception: # noqa: BLE001 - custom objects may use slots - pass - for cls in type(item).__mro__: - slots = vars(cls).get("__slots__", ()) - if isinstance(slots, str): - slots = (slots,) - for slot in slots: - if slot in {"__dict__", "__weakref__"}: - continue - if slot.startswith("__") and not slot.endswith("__"): - slot = f"_{cls.__name__.lstrip('_')}{slot}" - try: - stack.append(getattr(item, slot)) - except Exception: # noqa: BLE001 - unset/custom slots are best-effort - pass - except Exception: # noqa: BLE001 - traversal must never break a hook - pass + referents = gc.get_referents(item) + except Exception: # noqa: BLE001 - estimation must never break a hook + continue + for referent in referents: + # Type/module/function/code objects lead into process-global graphs, + # not memory retained specifically by this record. Bound methods, + # partial args, generator iterators, slots, buffers, and container + # subclasses remain traversable through their other referents. + if isinstance(referent, _RETAINED_GRAPH_BOUNDARIES): + continue + stack.append(referent) return total 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 78082cc6..5a2a2c55 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 @@ -73,31 +73,32 @@ } _exporter_owner_lock = threading.Lock() -_exporter_owners: list[tuple[InsightExporter, weakref.ReferenceType[object]]] = [] +_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.""" - def release(lane_ref: weakref.ReferenceType[object]) -> None: + def release(lane_ref: weakref.ReferenceType[Any]) -> None: with _exporter_owner_lock: _exporter_owners[:] = [ - entry for entry in _exporter_owners if entry[1] is not lane_ref + existing for existing in _exporter_owners if existing is not lane_ref ] - claims = [(lane._exporter, lane) for lane in lanes] with _exporter_owner_lock: _exporter_owners[:] = [ - entry for entry in _exporter_owners if entry[1]() is not None + lane_ref for lane_ref in _exporter_owners if lane_ref() is not None ] - for exporter, _ in claims: - if any(existing is exporter for existing, _ in _exporter_owners): - raise ValueError( - "the same exporter instance cannot be shared across " - "Workflow Insight plugin instances" - ) - for exporter, lane in claims: - _exporter_owners.append((exporter, weakref.ref(lane, release))) + for lane in lanes: + exporter = lane._exporter + for lane_ref in _exporter_owners: + owner = lane_ref() + if owner is not None and owner._exporter is exporter: + raise ValueError( + "the same exporter instance cannot be shared across " + "Workflow Insight plugin instances" + ) + _exporter_owners.extend(weakref.ref(lane, release) for lane in lanes) def _parse_execution_arn(execution_arn: str) -> dict[str, str]: 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 3bba291e..01bf2773 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 @@ -14,6 +14,7 @@ import gc import threading import time +import weakref import pytest @@ -269,6 +270,25 @@ def test_exporter_ownership_persists_while_lane_worker_is_alive(): 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) + + 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_default_exporter_unaffected_by_instance_check(): # No exporters configured -> single default LambdaLogExporter, one lane; the # identity check never trips on the empty list. diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index e6486802..fb94b8cb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,6 +11,7 @@ from __future__ import annotations +import functools import threading import time from typing import Any @@ -596,7 +597,7 @@ def test_non_json_record_reaches_exporter_without_evicting_backlog(): def test_individually_over_budget_record_does_not_evict_existing_backlog(): exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_500) + scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) lane = scheduler._lanes[0] scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) assert _wait_until(exporter.started.is_set) @@ -605,7 +606,7 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog(): scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000)) assert lane._pending_count() == 2 - assert lane._pending_bytes_count() <= 3_500 + assert lane._pending_bytes_count() <= 3_000 exporter.release() scheduler.end_invocation(5.0) exported = exporter.exported_values() @@ -706,6 +707,42 @@ def test_record_sizing_exception_does_not_escape_schedule(): assert exporter.exported_values() == ["inflight", "custom-sized"] +def test_retained_size_traverses_filtered_opaque_referents(): + class HiddenList(list[Any]): + def __init__(self, value: Any) -> None: + super().__init__([value]) + self.iterated = False + + def __iter__(self): + self.iterated = True + return iter(()) + + backing_buffers = [bytearray(4_000) for _ in range(3)] + hidden = HiddenList(backing_buffers[2]) + payloads = [ + functools.partial(lambda value: value, backing_buffers[0]), + (value for value in (backing_buffers[1],)), + hidden, + ] + + for payload in payloads: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + record = _rec(ARN_B, "opaque-referent") + record["payload"] = payload + scheduler.schedule(ARN_B, record) + + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + assert hidden.iterated is False + + def test_timed_out_barrier_flushes_eventually_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From 4bd5e14f415ee5f249d3e4ff6501b023bcf916f7 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:38:40 +0000 Subject: [PATCH 17/23] fix(insight): bound retained traversal work --- .../_export_scheduler.py | 148 +++++++++++++++--- .../plugin.py | 10 +- .../tests/test_config.py | 5 + .../tests/test_export_scheduler.py | 32 ++++ 4 files changed, 167 insertions(+), 28 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 7fae566f..7d1dd1be 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -21,7 +21,8 @@ from __future__ import annotations import copy -import gc +import functools +import itertools import logging import sys import threading @@ -48,44 +49,149 @@ _DEFAULT_MAX_PENDING_BYTES = 16_000_000 -_RETAINED_GRAPH_BOUNDARIES = ( +_UNSUPPORTED_RETAINED_GRAPH = object() +_ATOMIC_RETAINED_TYPES = ( + str, + bytes, + bytearray, + int, + float, + complex, + bool, + type(None), + range, + slice, type, types.ModuleType, - types.FunctionType, - types.BuiltinFunctionType, types.CodeType, + types.WrapperDescriptorType, + types.MethodDescriptorType, ) +class _RetainedChildren: + __slots__ = ("iterator",) + + def __init__(self, items: Any) -> None: + self.iterator = iter(items) + + +def _custom_retained_children(value: Any) -> list[Any]: + children: list[Any] = [] + try: + children.append(object.__getattribute__(value, "__dict__")) + except Exception: # noqa: BLE001 - custom objects may use slots only + pass + for cls in type(value).__mro__: + slots = vars(cls).get("__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + for slot in slots: + if slot in {"__dict__", "__weakref__"}: + continue + if slot.startswith("__") and not slot.endswith("__"): + slot = f"_{cls.__name__.lstrip('_')}{slot}" + try: + children.append(object.__getattribute__(value, slot)) + except Exception: # noqa: BLE001 - unset/custom slots are best-effort + pass + return children + + +def _retained_children(value: Any) -> Any: + custom = _custom_retained_children(value) + if isinstance(value, dict): + return itertools.chain(dict.__iter__(value), dict.values(value), custom) + if isinstance(value, list): + return itertools.chain(list.__iter__(value), custom) + if isinstance(value, tuple): + return itertools.chain(tuple.__iter__(value), custom) + if isinstance(value, set): + return itertools.chain(set.__iter__(value), custom) + if isinstance(value, frozenset): + return itertools.chain(frozenset.__iter__(value), custom) + if isinstance(value, deque): + return itertools.chain(deque.__iter__(value), custom) + if isinstance(value, memoryview): + return (value.obj,) + if isinstance(value, functools.partial): + return (value.func, value.args, value.keywords) + if isinstance(value, types.FunctionType): + closure = [] + for cell in value.__closure__ or (): + try: + closure.append(cell.cell_contents) + except ValueError: + pass + return itertools.chain(closure, (value.__defaults__, value.__kwdefaults__)) + if isinstance(value, types.MethodType): + return (value.__self__, value.__func__) + if isinstance(value, types.BuiltinFunctionType): + owner = value.__self__ + return () if owner is None or isinstance(owner, types.ModuleType) else (owner,) + if isinstance(value, types.MethodWrapperType): + return (value.__self__,) + if isinstance(value, types.GeneratorType): + frame = value.gi_frame + return () if frame is None else (frame.f_locals, value.gi_yieldfrom) + if isinstance(value, _ATOMIC_RETAINED_TYPES): + return () + if custom: + return custom + return _UNSUPPORTED_RETAINED_GRAPH + + +def _retained_shallow_size(value: Any) -> int: + try: + size = sys.getsizeof(value) + except Exception: # noqa: BLE001 - estimation must never break a hook + size = 1_024 + try: + if isinstance(value, dict): + size = max(size, dict.__sizeof__(value)) + elif isinstance(value, list): + size = max(size, list.__sizeof__(value)) + elif isinstance(value, tuple): + size = max(size, tuple.__sizeof__(value)) + elif isinstance(value, set): + size = max(size, set.__sizeof__(value)) + elif isinstance(value, frozenset): + size = max(size, frozenset.__sizeof__(value)) + elif isinstance(value, deque): + size = max(size, deque.__sizeof__(value)) + except Exception: # noqa: BLE001 - base sizing remains best-effort + pass + return size + + def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: - """Estimate a bounded retained graph without serializing or calling render().""" + """Estimate retained memory with bounded, non-overridable traversal.""" total = 0 seen: set[int] = set() stack: list[Any] = [value] while stack: item = stack.pop() + if isinstance(item, _RetainedChildren): + try: + child = next(item.iterator) + except StopIteration: + continue + except Exception: # noqa: BLE001 - fail closed on malformed iterators + return max_size + 1 if max_size is not None else total + 1_024 + stack.append(item) + stack.append(child) + continue identity = id(item) if identity in seen: continue seen.add(identity) - try: - total += sys.getsizeof(item) - except Exception: # noqa: BLE001 - estimation must never break a hook - total += 1_024 + total += _retained_shallow_size(item) if max_size is not None and total > max_size: return max_size + 1 - try: - referents = gc.get_referents(item) - except Exception: # noqa: BLE001 - estimation must never break a hook - continue - for referent in referents: - # Type/module/function/code objects lead into process-global graphs, - # not memory retained specifically by this record. Bound methods, - # partial args, generator iterators, slots, buffers, and container - # subclasses remain traversable through their other referents. - if isinstance(referent, _RETAINED_GRAPH_BOUNDARIES): - continue - stack.append(referent) + children = _retained_children(item) + if children is _UNSUPPORTED_RETAINED_GRAPH: + return max_size + 1 if max_size is not None else total + 1_024 + stack.append(_RetainedChildren(children)) return total 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 5a2a2c55..1079973b 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 @@ -79,12 +79,6 @@ def _claim_exporter_lanes(lanes: list[Any]) -> None: """Give each exporter object to at most one live scheduler lane.""" - def release(lane_ref: weakref.ReferenceType[Any]) -> None: - with _exporter_owner_lock: - _exporter_owners[:] = [ - existing for existing in _exporter_owners if existing is not lane_ref - ] - with _exporter_owner_lock: _exporter_owners[:] = [ lane_ref for lane_ref in _exporter_owners if lane_ref() is not None @@ -98,7 +92,9 @@ def release(lane_ref: weakref.ReferenceType[Any]) -> None: "the same exporter instance cannot be shared across " "Workflow Insight plugin instances" ) - _exporter_owners.extend(weakref.ref(lane, release) for lane in lanes) + # 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) def _parse_execution_arn(execution_arn: str) -> dict[str, str]: 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 01bf2773..5944b7eb 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 @@ -18,6 +18,7 @@ import pytest +import aws_durable_execution_sdk_python_insight.plugin as insight_plugin_module from aws_durable_execution_sdk_python_insight import ( EmitMode, OperationDetail, @@ -278,6 +279,10 @@ def test_exporter_plugin_cycle_is_not_rooted_by_ownership_registry(): 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 diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index fb94b8cb..43ee2efd 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -14,6 +14,7 @@ import functools import threading import time +import types from typing import Any from aws_durable_execution_sdk_python_insight._export_scheduler import ( @@ -217,6 +218,9 @@ def __init__(self) -> None: super().__init__([None] * 10_000) self.iterated = False + def __sizeof__(self) -> int: + return 1 + def __iter__(self): self.iterated = True return super().__iter__() @@ -743,6 +747,34 @@ def __iter__(self): assert hidden.iterated is False +def test_retained_size_counts_closures_and_bound_builtin_owners(): + closure_buffer = bytearray(4_000) + + def closure() -> bytearray: + return closure_buffer + + bound_owner = [bytearray(4_000)] + wrapper_owner = [bytearray(4_000)] + method_wrapper = wrapper_owner.__str__ + assert isinstance(method_wrapper, types.MethodWrapperType) + payloads = [closure, bound_owner.append, method_wrapper] + + for payload in payloads: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + record = _rec(ARN_B, "retained-callable") + record["payload"] = payload + scheduler.schedule(ARN_B, record) + + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + def test_timed_out_barrier_flushes_eventually_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From 3a78affbb4beb7c25b916c32f6a9d81bd33f714e Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:58:24 +0000 Subject: [PATCH 18/23] fix(insight): preserve bounded renderer values --- .../_export_scheduler.py | 52 +++++++++++++++---- .../tests/test_export_scheduler.py | 47 ++++++++++++++++- 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 7d1dd1be..cec1822e 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -21,6 +21,8 @@ from __future__ import annotations import copy +import datetime +import decimal import functools import itertools import logging @@ -28,6 +30,7 @@ import threading import time import types +import uuid from collections import OrderedDict, deque from typing import Any @@ -69,6 +72,16 @@ ) +_SAFE_OPAQUE_RETAINED_TYPES = ( + datetime.date, + datetime.datetime, + datetime.time, + datetime.timedelta, + decimal.Decimal, + uuid.UUID, +) + + class _RetainedChildren: __slots__ = ("iterator",) @@ -113,9 +126,9 @@ def _retained_children(value: Any) -> Any: if isinstance(value, deque): return itertools.chain(deque.__iter__(value), custom) if isinstance(value, memoryview): - return (value.obj,) + return itertools.chain((value.obj,), custom) if isinstance(value, functools.partial): - return (value.func, value.args, value.keywords) + return itertools.chain((value.func, value.args, value.keywords), custom) if isinstance(value, types.FunctionType): closure = [] for cell in value.__closure__ or (): @@ -123,19 +136,29 @@ def _retained_children(value: Any) -> Any: closure.append(cell.cell_contents) except ValueError: pass - return itertools.chain(closure, (value.__defaults__, value.__kwdefaults__)) + return itertools.chain( + closure, (value.__defaults__, value.__kwdefaults__), custom + ) if isinstance(value, types.MethodType): - return (value.__self__, value.__func__) + return itertools.chain((value.__self__, value.__func__), custom) if isinstance(value, types.BuiltinFunctionType): owner = value.__self__ - return () if owner is None or isinstance(owner, types.ModuleType) else (owner,) + retained = ( + () if owner is None or isinstance(owner, types.ModuleType) else (owner,) + ) + return itertools.chain(retained, custom) if isinstance(value, types.MethodWrapperType): - return (value.__self__,) + return itertools.chain((value.__self__,), custom) if isinstance(value, types.GeneratorType): frame = value.gi_frame - return () if frame is None else (frame.f_locals, value.gi_yieldfrom) - if isinstance(value, _ATOMIC_RETAINED_TYPES): - return () + generator_children = ( + () if frame is None else (frame.f_locals, value.gi_yieldfrom) + ) + return itertools.chain(generator_children, custom) + if type(value) in _ATOMIC_RETAINED_TYPES: + return custom + if type(value) in _SAFE_OPAQUE_RETAINED_TYPES: + return custom if custom: return custom return _UNSUPPORTED_RETAINED_GRAPH @@ -568,7 +591,16 @@ def __init__( def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: """Fan a canonical record out to every lane. Returns immediately.""" - record_size = _estimate_retained_size(record, self._max_pending_bytes) + try: + record_size = _estimate_retained_size(record, self._max_pending_bytes) + except Exception as exc: # noqa: BLE001 - inspection must never break a hook + _logger.warning( + "workflow-insight: retained-size inspection failed for %s; " + "rejecting this record safely: %s", + execution_arn, + exc, + ) + record_size = self._max_pending_bytes + 1 for lane in self._lanes: lane.schedule(execution_arn, record, record_size) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 43ee2efd..827b25ab 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,12 +11,15 @@ from __future__ import annotations +import datetime import functools import threading import time import types from typing import Any +import aws_durable_execution_sdk_python_insight._export_scheduler as scheduler_module + from aws_durable_execution_sdk_python_insight._export_scheduler import ( _ExportScheduler, ) @@ -721,12 +724,18 @@ def __iter__(self): self.iterated = True return iter(()) - backing_buffers = [bytearray(4_000) for _ in range(3)] + class PayloadPartial(functools.partial): + pass + + backing_buffers = [bytearray(4_000) for _ in range(4)] hidden = HiddenList(backing_buffers[2]) + partial_with_payload = PayloadPartial(lambda value: value, "small") + partial_with_payload.payload = backing_buffers[3] payloads = [ functools.partial(lambda value: value, backing_buffers[0]), (value for value in (backing_buffers[1],)), hidden, + partial_with_payload, ] for payload in payloads: @@ -775,6 +784,42 @@ def closure() -> bytearray: scheduler.end_invocation(5.0) +def test_safe_opaque_datetime_reaches_custom_renderer(): + class DateRenderExporter(RecordingExporter): + def __init__(self) -> None: + super().__init__(max_record_size_bytes=10_000) + self.rendered: list[str] = [] + + def render(self, record: dict[str, Any]) -> Any: + value = record["payload"].isoformat() + self.rendered.append(value) + return {"value": value} + + exporter = DateRenderExporter() + scheduler = _ExportScheduler([exporter]) + record = _rec(ARN_A, "date") + record["payload"] = datetime.date(2026, 9, 10) + scheduler.schedule(ARN_A, record) + scheduler.end_invocation(5.0) + + assert exporter.rendered == ["2026-09-10"] + assert exporter.exported_values() == ["date"] + + +def test_retained_size_inspection_failure_does_not_escape_schedule(monkeypatch): + def fail_estimate(value: Any, max_size: int | None = None) -> int: + raise RuntimeError("inspection failed") + + monkeypatch.setattr(scheduler_module, "_estimate_retained_size", fail_estimate) + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + + scheduler.schedule(ARN_A, _rec(ARN_A, "rejected")) + + assert scheduler._lanes[0]._pending_count() == 0 + assert scheduler._lanes[0]._worker is None + + def test_timed_out_barrier_flushes_eventually_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From 7b5c6a3383009fd90a86136b6e2622d45636fa24 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 03:17:11 +0000 Subject: [PATCH 19/23] fix(insight): bound in-flight export retention --- .../_export_scheduler.py | 64 ++++++++++----- .../tests/test_export_scheduler.py | 82 ++++++++++++++++++- 2 files changed, 121 insertions(+), 25 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index cec1822e..5e0d4298 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -26,7 +26,6 @@ import functools import itertools import logging -import sys import threading import time import types @@ -63,7 +62,6 @@ bool, type(None), range, - slice, type, types.ModuleType, types.CodeType, @@ -125,6 +123,8 @@ def _retained_children(value: Any) -> Any: return itertools.chain(frozenset.__iter__(value), custom) if isinstance(value, deque): return itertools.chain(deque.__iter__(value), custom) + if isinstance(value, slice): + return itertools.chain((value.start, value.stop, value.step), custom) if isinstance(value, memoryview): return itertools.chain((value.obj,), custom) if isinstance(value, functools.partial): @@ -165,26 +165,25 @@ def _retained_children(value: Any) -> Any: def _retained_shallow_size(value: Any) -> int: - try: - size = sys.getsizeof(value) - except Exception: # noqa: BLE001 - estimation must never break a hook - size = 1_024 + """Return shallow size without dispatching to user-defined ``__sizeof__``.""" try: if isinstance(value, dict): - size = max(size, dict.__sizeof__(value)) - elif isinstance(value, list): - size = max(size, list.__sizeof__(value)) - elif isinstance(value, tuple): - size = max(size, tuple.__sizeof__(value)) - elif isinstance(value, set): - size = max(size, set.__sizeof__(value)) - elif isinstance(value, frozenset): - size = max(size, frozenset.__sizeof__(value)) - elif isinstance(value, deque): - size = max(size, deque.__sizeof__(value)) - except Exception: # noqa: BLE001 - base sizing remains best-effort - pass - return size + return dict.__sizeof__(value) + if isinstance(value, list): + return list.__sizeof__(value) + if isinstance(value, tuple): + return tuple.__sizeof__(value) + if isinstance(value, set): + return set.__sizeof__(value) + if isinstance(value, frozenset): + return frozenset.__sizeof__(value) + if isinstance(value, deque): + return deque.__sizeof__(value) + if type(value) in _ATOMIC_RETAINED_TYPES + _SAFE_OPAQUE_RETAINED_TYPES: + return value.__sizeof__() + return object.__sizeof__(value) + except Exception: # noqa: BLE001 - estimation must never break a hook + return 1_024 def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: @@ -287,6 +286,7 @@ def __init__( self._max_pending = max(1, max_pending_executions) self._max_pending_bytes = max(1, max_pending_bytes) self._pending_bytes = 0 + self._inflight_bytes = 0 # 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 @@ -374,6 +374,12 @@ def cancel_flush(self, barrier: _FlushBarrier) -> None: """Stop waiting for a timed-out barrier while retaining one later flush.""" with self._cond: barrier.canceled = True + 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 # Keep at most one detached flush. Moving it to this barrier's # position makes it cover all work scheduled before the latest # timeout without accumulating one marker per warm invocation. @@ -409,7 +415,10 @@ def _enforce_pending_caps(self) -> None: self._max_pending, old_arn, ) - while self._pending_bytes > self._max_pending_bytes and self._pending: + while ( + self._pending_bytes + self._inflight_bytes > self._max_pending_bytes + and self._pending + ): old_arn, dropped_size = self._drop_oldest_pending() _logger.warning( "workflow-insight: export lane for %s reached its pending byte " @@ -437,6 +446,7 @@ def _disable_locked(self, exc: Exception) -> None: self._worker = None self._pending.clear() self._pending_bytes = 0 + self._inflight_bytes = 0 for kind, payload in self._queue: if kind == _FLUSH and payload is not None: barrier: _FlushBarrier = payload @@ -485,15 +495,21 @@ def _run_worker(self) -> None: return kind, payload = self._queue.popleft() record: dict[str, Any] | None = None + record_size = 0 if kind == _RECORD: pending = self._pending.pop(payload, None) if pending is None: continue record, record_size = pending self._pending_bytes -= record_size + self._inflight_bytes += record_size if kind == _RECORD and record is not None: - self._export_one(record) + try: + self._export_one(record) + finally: + with self._cond: + self._inflight_bytes -= record_size else: # _FLUSH barrier: _FlushBarrier | None = payload self._flush() @@ -560,6 +576,10 @@ def _pending_bytes_count(self) -> int: with self._cond: return self._pending_bytes + def _retained_bytes_count(self) -> int: + with self._cond: + return self._pending_bytes + self._inflight_bytes + def _queue_len(self) -> int: with self._cond: return len(self._queue) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 827b25ab..aea400f5 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -604,16 +604,16 @@ def test_non_json_record_reaches_exporter_without_evicting_backlog(): def test_individually_over_budget_record_does_not_evict_existing_backlog(): exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) + scheduler = _ExportScheduler([exporter], max_pending_bytes=5_000) lane = scheduler._lanes[0] scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) assert _wait_until(exporter.started.is_set) scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 700)) scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 700)) - scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000)) + scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 5_000)) assert lane._pending_count() == 2 - assert lane._pending_bytes_count() <= 3_000 + assert lane._pending_bytes_count() <= 5_000 exporter.release() scheduler.end_invocation(5.0) exported = exporter.exported_values() @@ -932,3 +932,79 @@ def test_cancel_flush_after_pop_lets_worker_complete_barrier(): assert exporter.calls.count(("flush", None)) == 1 lane.request_stop_when_idle() assert _wait_until(lambda: not lane._worker_alive()) + + +def test_retained_size_counts_slice_referents(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + + record = _rec(ARN_B, "slice") + record["payload"] = slice(bytearray(4_000), None) + scheduler.schedule(ARN_B, record) + + assert lane._pending_count() == 0 + assert lane._pending_bytes_count() == 0 + exporter.release() + scheduler.end_invocation(5.0) + + +def test_retained_size_does_not_dispatch_custom_sizeof(): + called = threading.Event() + + class CustomSized: + def __sizeof__(self) -> int: + called.set() + raise AssertionError("custom __sizeof__ must not run") + + exporter = RecordingExporter() + scheduler = _ExportScheduler([exporter]) + record = _rec(ARN_A, "custom-sized") + record["payload"] = CustomSized() + + scheduler.schedule(ARN_A, record) + scheduler.end_invocation(5.0) + + assert called.is_set() is False + assert exporter.exported_values() == ["custom-sized"] + + +def test_inflight_record_remains_charged_until_export_returns(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a" * 1_500)) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 1_500)) + + assert lane._pending_count() == 0 + assert lane._retained_bytes_count() <= 3_000 + exporter.release() + scheduler.end_invocation(5.0) + assert exporter.exported_values() == ["a" * 1_500] + assert lane._retained_bytes_count() == 0 + + +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 From f596264592d0ca9ff0ed56ac344ca3c67f153a9b Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 17:28:12 +0000 Subject: [PATCH 20/23] refactor(insight): simplify pending export scheduling --- .../README.md | 11 +- .../_export_scheduler.py | 363 ++--------------- .../tests/test_export_scheduler.py | 372 ++---------------- 3 files changed, 58 insertions(+), 688 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 17b5617b..1c3b5f24 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -63,12 +63,11 @@ Behavior is validated cross-SDK by the `insight` conformance suite > `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 -> 1,024 pending executions and 16 MB of estimated retained memory; it drops the -> oldest pending snapshot when either bound is reached. 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 +> 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. diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 5e0d4298..911209d7 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -9,10 +9,8 @@ Each exporter lane: -* Keeps the latest pending snapshot per execution ARN and processes ARNs - round-robin. -* Drops the oldest pending snapshot when the execution-count or byte budget is - reached, keeping memory bounded. +* 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. @@ -21,16 +19,10 @@ from __future__ import annotations import copy -import datetime -import decimal -import functools -import itertools import logging import threading import time -import types -import uuid -from collections import OrderedDict, deque +from collections import deque from typing import Any from aws_durable_execution_sdk_python_insight.truncation import truncate_record @@ -39,183 +31,6 @@ _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 - -# Estimated Python object memory retained by one blocked lane. Lambda functions -# can be configured with 128 MiB, so keep the instrumentation backlog well below -# that floor. This is a conservative budget signal, not an exact heap measurement. -_DEFAULT_MAX_PENDING_BYTES = 16_000_000 - - -_UNSUPPORTED_RETAINED_GRAPH = object() -_ATOMIC_RETAINED_TYPES = ( - str, - bytes, - bytearray, - int, - float, - complex, - bool, - type(None), - range, - type, - types.ModuleType, - types.CodeType, - types.WrapperDescriptorType, - types.MethodDescriptorType, -) - - -_SAFE_OPAQUE_RETAINED_TYPES = ( - datetime.date, - datetime.datetime, - datetime.time, - datetime.timedelta, - decimal.Decimal, - uuid.UUID, -) - - -class _RetainedChildren: - __slots__ = ("iterator",) - - def __init__(self, items: Any) -> None: - self.iterator = iter(items) - - -def _custom_retained_children(value: Any) -> list[Any]: - children: list[Any] = [] - try: - children.append(object.__getattribute__(value, "__dict__")) - except Exception: # noqa: BLE001 - custom objects may use slots only - pass - for cls in type(value).__mro__: - slots = vars(cls).get("__slots__", ()) - if isinstance(slots, str): - slots = (slots,) - for slot in slots: - if slot in {"__dict__", "__weakref__"}: - continue - if slot.startswith("__") and not slot.endswith("__"): - slot = f"_{cls.__name__.lstrip('_')}{slot}" - try: - children.append(object.__getattribute__(value, slot)) - except Exception: # noqa: BLE001 - unset/custom slots are best-effort - pass - return children - - -def _retained_children(value: Any) -> Any: - custom = _custom_retained_children(value) - if isinstance(value, dict): - return itertools.chain(dict.__iter__(value), dict.values(value), custom) - if isinstance(value, list): - return itertools.chain(list.__iter__(value), custom) - if isinstance(value, tuple): - return itertools.chain(tuple.__iter__(value), custom) - if isinstance(value, set): - return itertools.chain(set.__iter__(value), custom) - if isinstance(value, frozenset): - return itertools.chain(frozenset.__iter__(value), custom) - if isinstance(value, deque): - return itertools.chain(deque.__iter__(value), custom) - if isinstance(value, slice): - return itertools.chain((value.start, value.stop, value.step), custom) - if isinstance(value, memoryview): - return itertools.chain((value.obj,), custom) - if isinstance(value, functools.partial): - return itertools.chain((value.func, value.args, value.keywords), custom) - if isinstance(value, types.FunctionType): - closure = [] - for cell in value.__closure__ or (): - try: - closure.append(cell.cell_contents) - except ValueError: - pass - return itertools.chain( - closure, (value.__defaults__, value.__kwdefaults__), custom - ) - if isinstance(value, types.MethodType): - return itertools.chain((value.__self__, value.__func__), custom) - if isinstance(value, types.BuiltinFunctionType): - owner = value.__self__ - retained = ( - () if owner is None or isinstance(owner, types.ModuleType) else (owner,) - ) - return itertools.chain(retained, custom) - if isinstance(value, types.MethodWrapperType): - return itertools.chain((value.__self__,), custom) - if isinstance(value, types.GeneratorType): - frame = value.gi_frame - generator_children = ( - () if frame is None else (frame.f_locals, value.gi_yieldfrom) - ) - return itertools.chain(generator_children, custom) - if type(value) in _ATOMIC_RETAINED_TYPES: - return custom - if type(value) in _SAFE_OPAQUE_RETAINED_TYPES: - return custom - if custom: - return custom - return _UNSUPPORTED_RETAINED_GRAPH - - -def _retained_shallow_size(value: Any) -> int: - """Return shallow size without dispatching to user-defined ``__sizeof__``.""" - try: - if isinstance(value, dict): - return dict.__sizeof__(value) - if isinstance(value, list): - return list.__sizeof__(value) - if isinstance(value, tuple): - return tuple.__sizeof__(value) - if isinstance(value, set): - return set.__sizeof__(value) - if isinstance(value, frozenset): - return frozenset.__sizeof__(value) - if isinstance(value, deque): - return deque.__sizeof__(value) - if type(value) in _ATOMIC_RETAINED_TYPES + _SAFE_OPAQUE_RETAINED_TYPES: - return value.__sizeof__() - return object.__sizeof__(value) - except Exception: # noqa: BLE001 - estimation must never break a hook - return 1_024 - - -def _estimate_retained_size(value: Any, max_size: int | None = None) -> int: - """Estimate retained memory with bounded, non-overridable traversal.""" - total = 0 - seen: set[int] = set() - stack: list[Any] = [value] - while stack: - item = stack.pop() - if isinstance(item, _RetainedChildren): - try: - child = next(item.iterator) - except StopIteration: - continue - except Exception: # noqa: BLE001 - fail closed on malformed iterators - return max_size + 1 if max_size is not None else total + 1_024 - stack.append(item) - stack.append(child) - continue - identity = id(item) - if identity in seen: - continue - seen.add(identity) - total += _retained_shallow_size(item) - if max_size is not None and total > max_size: - return max_size + 1 - children = _retained_children(item) - if children is _UNSUPPORTED_RETAINED_GRAPH: - return max_size + 1 if max_size is not None else total + 1_024 - stack.append(_RetainedChildren(children)) - return total - def _copy_record_containers(record: dict[str, Any]) -> dict[str, Any]: """Copy built-in containers while treating custom values as opaque leaves.""" @@ -275,18 +90,8 @@ class _ExporterLane: the queue; scheduling threads are producers that wake it via ``notify``. """ - def __init__( - self, - exporter: InsightExporter, - *, - max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, - max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, - ) -> None: + def __init__(self, exporter: InsightExporter) -> None: self._exporter = exporter - self._max_pending = max(1, max_pending_executions) - self._max_pending_bytes = max(1, max_pending_bytes) - self._pending_bytes = 0 - self._inflight_bytes = 0 # 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 @@ -294,61 +99,27 @@ def __init__( # 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). + # Ordered work list: entries are (_RECORD, None) or (_FLUSH, barrier). self._queue: deque[tuple[str, Any]] = deque() - # arn -> (latest pending record, retained-memory estimate). Insertion - # order is both record age and fairness order because replacing an ARN - # moves it to the back. - self._pending: OrderedDict[str, tuple[dict[str, Any], int]] = OrderedDict() + # 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, - execution_arn: str, - record: dict[str, Any], - record_size: int, - ) -> None: + def schedule(self, record: dict[str, Any]) -> None: with self._cond: if self._disabled: return self._stop_when_idle = False - size = max(0, record_size) - if size > self._max_pending_bytes: - superseded = self._pending.pop(execution_arn, None) - if superseded is not None: - _, superseded_size = superseded - self._pending_bytes -= superseded_size - self._remove_record_token(execution_arn) - _logger.warning( - "workflow-insight: pending record for %s on %s exceeds the " - "byte budget (%d > %d); dropping this record%s", - execution_arn, - type(self._exporter).__name__, - size, - self._max_pending_bytes, - " and its superseded pending snapshot" - if superseded is not None - else "", - ) - return - 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. - _, old_size = self._pending[execution_arn] - self._pending_bytes -= old_size - self._pending[execution_arn] = (record, size) - self._pending_bytes += size - self._pending.move_to_end(execution_arn) - self._move_record_token_to_back(execution_arn) + if self._pending is None: + self._queue.append((_RECORD, None)) else: - self._pending[execution_arn] = (record, size) - self._pending_bytes += size - self._queue.append((_RECORD, execution_arn)) - self._enforce_pending_caps() + self._move_record_token_to_back() + self._pending = record self._ensure_worker_locked() self._cond.notify() @@ -395,58 +166,17 @@ def cancel_flush(self, barrier: _FlushBarrier) -> None: # -- 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_caps(self) -> None: - while len(self._pending) > self._max_pending: - old_arn, _ = self._drop_oldest_pending() - _logger.warning( - "workflow-insight: export lane for %s reached its execution cap " - "(%d); dropping pending record for %s", - type(self._exporter).__name__, - self._max_pending, - old_arn, - ) - while ( - self._pending_bytes + self._inflight_bytes > self._max_pending_bytes - and self._pending - ): - old_arn, dropped_size = self._drop_oldest_pending() - _logger.warning( - "workflow-insight: export lane for %s reached its pending byte " - "budget (%d); dropping %d-byte pending record for %s", - type(self._exporter).__name__, - self._max_pending_bytes, - dropped_size, - old_arn, - ) - - def _drop_oldest_pending(self) -> tuple[str, int]: - old_arn, (_, old_size) = self._pending.popitem(last=False) - self._pending_bytes -= old_size - self._remove_record_token(old_arn) - return old_arn, old_size - - def _remove_record_token(self, execution_arn: str) -> None: - for index, (kind, payload) in enumerate(self._queue): - if kind == _RECORD and payload == execution_arn: + 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.clear() - self._pending_bytes = 0 - self._inflight_bytes = 0 + self._pending = None for kind, payload in self._queue: if kind == _FLUSH and payload is not None: barrier: _FlushBarrier = payload @@ -495,21 +225,14 @@ def _run_worker(self) -> None: return kind, payload = self._queue.popleft() record: dict[str, Any] | None = None - record_size = 0 if kind == _RECORD: - pending = self._pending.pop(payload, None) - if pending is None: + record = self._pending + self._pending = None + if record is None: continue - record, record_size = pending - self._pending_bytes -= record_size - self._inflight_bytes += record_size if kind == _RECORD and record is not None: - try: - self._export_one(record) - finally: - with self._cond: - self._inflight_bytes -= record_size + self._export_one(record) else: # _FLUSH barrier: _FlushBarrier | None = payload self._flush() @@ -570,15 +293,7 @@ def _worker_alive(self) -> bool: def _pending_count(self) -> int: with self._cond: - return len(self._pending) - - def _pending_bytes_count(self) -> int: - with self._cond: - return self._pending_bytes - - def _retained_bytes_count(self) -> int: - with self._cond: - return self._pending_bytes + self._inflight_bytes + return int(self._pending is not None) def _queue_len(self) -> int: with self._cond: @@ -592,37 +307,13 @@ def _queued_flush_count(self) -> int: 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, - max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, - ) -> None: - self._max_pending_bytes = max(1, max_pending_bytes) - self._lanes = [ - _ExporterLane( - exporter, - max_pending_executions=max_pending_executions, - max_pending_bytes=self._max_pending_bytes, - ) - for exporter in exporters - ] + 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: + def schedule(self, _execution_arn: str, record: dict[str, Any]) -> None: """Fan a canonical record out to every lane. Returns immediately.""" - try: - record_size = _estimate_retained_size(record, self._max_pending_bytes) - except Exception as exc: # noqa: BLE001 - inspection must never break a hook - _logger.warning( - "workflow-insight: retained-size inspection failed for %s; " - "rejecting this record safely: %s", - execution_arn, - exc, - ) - record_size = self._max_pending_bytes + 1 for lane in self._lanes: - lane.schedule(execution_arn, record, record_size) + lane.schedule(record) def end_invocation(self, timeout_seconds: float) -> bool: """Drain and flush every touched lane under one shared timeout. diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index aea400f5..86fd95fd 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -11,15 +11,10 @@ from __future__ import annotations -import datetime -import functools import threading import time -import types from typing import Any -import aws_durable_execution_sdk_python_insight._export_scheduler as scheduler_module - from aws_durable_execution_sdk_python_insight._export_scheduler import ( _ExportScheduler, ) @@ -197,38 +192,6 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Any: raise RuntimeError("uncopyable payload") -class _SlottedPayload: - __slots__ = ("payload",) - - def __init__(self, payload: Any) -> None: - self.payload = payload - - -class _UnsizedSlottedPayload(_SlottedPayload): - def __sizeof__(self) -> int: - raise RuntimeError("size unavailable") - - -class _Unsized: - """A payload whose custom ``__sizeof__`` raises.""" - - def __sizeof__(self) -> int: - raise RuntimeError("size unavailable") - - -class _TrackedLargeList(list[Any]): - def __init__(self) -> None: - super().__init__([None] * 10_000) - self.iterated = False - - def __sizeof__(self) -> int: - return 1 - - def __iter__(self): - self.iterated = True - return super().__iter__() - - # -- lazy worker creation / one worker per exporter -------------------------- @@ -279,7 +242,6 @@ def fail_start(self): assert lane._disabled is True assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 assert lane._queue_len() == 0 assert scheduler.end_invocation(0.1) is False assert exporter.exported_values() == [] @@ -314,19 +276,20 @@ def test_same_execution_coalescing_exports_inflight_then_latest(): scheduler.end_invocation(5.0) -def test_different_executions_isolated_and_fair(): +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) # 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] + 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() - # 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) + assert exporter.exported_values() == ["a1", "d1"] def test_terminal_record_supersedes_pending_running(): @@ -546,278 +509,49 @@ def test_repeated_invocation_cycles_do_not_leak_threads(): assert len(exporter.exported_values()) == 20 -# -- pending cap / cancelled barrier cleanup --------------------------------- +# -- structurally bounded pending slot ---------------------------------------- -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_pending_byte_budget_evicts_oldest_large_record(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - - scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 1_500)) - scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 1_500)) - - assert lane._pending_count() == 1 - assert lane._pending_bytes_count() <= 3_000 - exporter.release() - scheduler.end_invocation(5.0) - exported = exporter.exported_values() - assert exported[0] == "inflight" - assert exported[1] == "c" * 1_500 - - -def test_non_json_record_reaches_exporter_without_evicting_backlog(): +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) - scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) - scheduler.schedule(ARN_C, _rec(ARN_C, "c1")) - - non_json = _rec(ARN_D, "custom") - non_json["payload"] = {"not-json"} - scheduler.schedule(ARN_D, non_json) - assert lane._pending_count() == 3 - exporter.release() - scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight", "b1", "c1", "custom"] - - -def test_individually_over_budget_record_does_not_evict_existing_backlog(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=5_000) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 700)) - scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 700)) - scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 5_000)) - - assert lane._pending_count() == 2 - assert lane._pending_bytes_count() <= 5_000 - exporter.release() - scheduler.end_invocation(5.0) - exported = exporter.exported_values() - assert exported[0] == "inflight" - assert exported[1:] == ["b" * 700, "c" * 700] - - -def test_over_budget_replacement_removes_superseded_same_arn_only(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - scheduler.schedule(ARN_A, _rec(ARN_A, "stale-running")) - scheduler.schedule(ARN_B, _rec(ARN_B, "unrelated")) - scheduler.schedule( - ARN_A, - _rec(ARN_A, "terminal" * 500, status="SUCCEEDED"), - ) + 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._pending_bytes_count() <= 3_500 + assert lane._queue_len() == 1 exporter.release() scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight", "unrelated"] - - -def test_retained_size_traverses_slots_after_shallow_size_failure(): - for payload in ( - _SlottedPayload("x" * 4_000), - _UnsizedSlottedPayload("x" * 4_000), - ): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - record = _rec(ARN_B, "opaque") - record["payload"] = payload - scheduler.schedule(ARN_B, record) - - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight"] + assert exporter.exported_values() == ["inflight", "pending-99"] -def test_retained_size_saturates_before_traversing_large_shallow_container(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - - payload = _TrackedLargeList() - record = _rec(ARN_B, "large-shallow") - record["payload"] = payload - scheduler.schedule(ARN_B, record) - - assert payload.iterated is False - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) - - -def test_retained_size_counts_memoryview_backing_buffer(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - - record = _rec(ARN_B, "memoryview") - record["payload"] = memoryview(bytearray(4_000)) - scheduler.schedule(ARN_B, record) - - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) +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") -def test_record_sizing_exception_does_not_escape_schedule(): 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() - record = _rec(ARN_B, "custom-sized") - record["payload"] = _Unsized() 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", "custom-sized"] - - -def test_retained_size_traverses_filtered_opaque_referents(): - class HiddenList(list[Any]): - def __init__(self, value: Any) -> None: - super().__init__([value]) - self.iterated = False - - def __iter__(self): - self.iterated = True - return iter(()) - - class PayloadPartial(functools.partial): - pass - - backing_buffers = [bytearray(4_000) for _ in range(4)] - hidden = HiddenList(backing_buffers[2]) - partial_with_payload = PayloadPartial(lambda value: value, "small") - partial_with_payload.payload = backing_buffers[3] - payloads = [ - functools.partial(lambda value: value, backing_buffers[0]), - (value for value in (backing_buffers[1],)), - hidden, - partial_with_payload, - ] - - for payload in payloads: - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - record = _rec(ARN_B, "opaque-referent") - record["payload"] = payload - scheduler.schedule(ARN_B, record) - - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) - - assert hidden.iterated is False - - -def test_retained_size_counts_closures_and_bound_builtin_owners(): - closure_buffer = bytearray(4_000) - - def closure() -> bytearray: - return closure_buffer - - bound_owner = [bytearray(4_000)] - wrapper_owner = [bytearray(4_000)] - method_wrapper = wrapper_owner.__str__ - assert isinstance(method_wrapper, types.MethodWrapperType) - payloads = [closure, bound_owner.append, method_wrapper] - - for payload in payloads: - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - record = _rec(ARN_B, "retained-callable") - record["payload"] = payload - scheduler.schedule(ARN_B, record) - - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) - - -def test_safe_opaque_datetime_reaches_custom_renderer(): - class DateRenderExporter(RecordingExporter): - def __init__(self) -> None: - super().__init__(max_record_size_bytes=10_000) - self.rendered: list[str] = [] - - def render(self, record: dict[str, Any]) -> Any: - value = record["payload"].isoformat() - self.rendered.append(value) - return {"value": value} - - exporter = DateRenderExporter() - scheduler = _ExportScheduler([exporter]) - record = _rec(ARN_A, "date") - record["payload"] = datetime.date(2026, 9, 10) - scheduler.schedule(ARN_A, record) - scheduler.end_invocation(5.0) - - assert exporter.rendered == ["2026-09-10"] - assert exporter.exported_values() == ["date"] - - -def test_retained_size_inspection_failure_does_not_escape_schedule(monkeypatch): - def fail_estimate(value: Any, max_size: int | None = None) -> int: - raise RuntimeError("inspection failed") - - monkeypatch.setattr(scheduler_module, "_estimate_retained_size", fail_estimate) - exporter = RecordingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - - scheduler.schedule(ARN_A, _rec(ARN_A, "rejected")) - - assert scheduler._lanes[0]._pending_count() == 0 - assert scheduler._lanes[0]._worker is None + assert exporter.exported_values() == ["inflight", "opaque"] def test_timed_out_barrier_flushes_eventually_and_worker_exits(): @@ -934,60 +668,6 @@ def test_cancel_flush_after_pop_lets_worker_complete_barrier(): assert _wait_until(lambda: not lane._worker_alive()) -def test_retained_size_counts_slice_referents(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) - assert _wait_until(exporter.started.is_set) - - record = _rec(ARN_B, "slice") - record["payload"] = slice(bytearray(4_000), None) - scheduler.schedule(ARN_B, record) - - assert lane._pending_count() == 0 - assert lane._pending_bytes_count() == 0 - exporter.release() - scheduler.end_invocation(5.0) - - -def test_retained_size_does_not_dispatch_custom_sizeof(): - called = threading.Event() - - class CustomSized: - def __sizeof__(self) -> int: - called.set() - raise AssertionError("custom __sizeof__ must not run") - - exporter = RecordingExporter() - scheduler = _ExportScheduler([exporter]) - record = _rec(ARN_A, "custom-sized") - record["payload"] = CustomSized() - - scheduler.schedule(ARN_A, record) - scheduler.end_invocation(5.0) - - assert called.is_set() is False - assert exporter.exported_values() == ["custom-sized"] - - -def test_inflight_record_remains_charged_until_export_returns(): - exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "a" * 1_500)) - assert _wait_until(exporter.started.is_set) - - scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 1_500)) - - assert lane._pending_count() == 0 - assert lane._retained_bytes_count() <= 3_000 - exporter.release() - scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["a" * 1_500] - assert lane._retained_bytes_count() == 0 - - def test_cancel_popped_barrier_preserves_later_detached_flush(): exporter = BlockingFlushExporter() scheduler = _ExportScheduler([exporter]) From 41e83739297f63cef45ca9f53d20d65ec936ef9b Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 18:50:07 +0000 Subject: [PATCH 21/23] fix(insight): preserve latest detached flush --- .../_export_scheduler.py | 31 +++++++++++-------- .../tests/test_export_scheduler.py | 25 +++++++++++++++ 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 911209d7..64b38571 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -142,27 +142,32 @@ def request_stop_when_idle(self) -> None: self._cond.notify() def cancel_flush(self, barrier: _FlushBarrier) -> None: - """Stop waiting for a timed-out barrier while retaining one later flush.""" + """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 - # Keep at most one detached flush. Moving it to this barrier's - # position makes it cover all work scheduled before the latest - # timeout without accumulating one marker per warm invocation. - for index in range(len(self._queue) - 1, -1, -1): - kind, payload = self._queue[index] - if kind == _FLUSH and payload is None: - del self._queue[index] - for index, (kind, payload) in enumerate(self._queue): - if kind == _FLUSH and payload is barrier: - self._queue[index] = (_FLUSH, None) - barrier.complete() - 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``) ------------------------------ diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 86fd95fd..1d800379 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -688,3 +688,28 @@ def test_cancel_popped_barrier_preserves_later_detached_flush(): 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 From 867d65298b2d9cb2db7e991d23b9070fab49a7c6 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 21:10:36 +0000 Subject: [PATCH 22/23] fix(insight): preserve records after copy failure --- .../_export_scheduler.py | 6 +-- .../truncation.py | 2 +- .../tests/test_export_scheduler.py | 45 +++++++++++++++++++ 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index 64b38571..f5a94d5d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -252,14 +252,14 @@ def _export_one(self, record: dict[str, Any]) -> None: # requiring custom-renderable values to implement ``deepcopy``. try: local = _copy_record_containers(record) - except Exception as exc: # noqa: BLE001 - malformed containers must not break the lane + except Exception as exc: # noqa: BLE001 - export remains best-effort _logger.warning( "workflow-insight: record container copy failed for exporter %s; " - "skipping export for this record: %s", + "using the original record without lane isolation: %s", type(exporter).__name__, exc, ) - return + local = record try: shaped = truncate_record( local, exporter.max_record_size_bytes, exporter.render diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py index aee615b6..d826d2e0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/truncation.py @@ -28,7 +28,7 @@ def json_byte_size(value: Any) -> int | None: return len( json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode("utf-8") ) - except Exception: # noqa: BLE001 - sizing failure must never break instrumentation + except (TypeError, ValueError): return None diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 1d800379..8b438e1f 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -713,3 +713,48 @@ def test_cancel_queued_barrier_preserves_later_detached_flush(): 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"] From afc10fe129032aa3f2496bcbea3af41b17fe11bd Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 22:46:02 +0000 Subject: [PATCH 23/23] fix(insight): defer exporter finalization past lock --- .../plugin.py | 16 +++-- .../tests/test_config.py | 64 +++++++++++++++++++ 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 1079973b..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 @@ -79,15 +79,17 @@ 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: - _exporter_owners[:] = [ - lane_ref for lane_ref in _exporter_owners if lane_ref() is not None - ] + 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 lane_ref in _exporter_owners: - owner = lane_ref() - if owner is not None and owner._exporter is 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" @@ -95,6 +97,8 @@ def _claim_exporter_lanes(lanes: list[Any]) -> None: # 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]: 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 5944b7eb..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 @@ -15,6 +15,7 @@ import threading import time import weakref +from typing import Any import pytest @@ -294,6 +295,69 @@ def test_exporter_plugin_cycle_is_not_rooted_by_ownership_registry(): 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.