From 7663297b5b11c39daaeda7276ad297a451a18796 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 19:04:43 +0000 Subject: [PATCH 01/13] fix(insight): buffer on-change export bursts --- .../README.md | 28 +- .../_export_scheduler.py | 483 ++++--------- .../tests/test_export_scheduler.py | 656 ++++++------------ .../tests/test_plugin.py | 16 +- .../tests/test_plugin_async.py | 32 + 5 files changed, 392 insertions(+), 823 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..38f9f02a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -59,18 +59,26 @@ 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, 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 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 +> 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. Each lane keeps a +> bounded FIFO of up to 16 pending snapshots per execution and at most 1,024 +> pending snapshots across the lane, preserving ordinary bursts without relying +> on daemon-thread scheduling. If an exporter remains slower than the producer +> and either bound fills, the oldest pending snapshot is dropped so the newest progress and terminal snapshots are retained. 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. +> +> `flush()` is lane-wide, not execution-scoped: it applies to the configured +> exporter instance's entire buffer. A barrier may therefore publish records +> from another execution that were already buffered, while a record scheduled +> after that barrier is exported after the flush and waits for a later barrier. +> A custom batching exporter that requires execution-level isolation should key +> its buffer by `executionArn` or use a distinct exporter instance per isolated +> stream. ## 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 index 5e0d4298..54afcb54 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,9 @@ 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 a bounded FIFO per execution ARN and processes ARNs round-robin. +* Drops the oldest pending snapshots when the per-execution or lane-wide limit + is reached, favoring the newest progress and terminal snapshots. * 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,15 +20,9 @@ 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 typing import Any @@ -45,198 +38,16 @@ # 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.""" - memo: dict[int, Any] = {} - seen: set[int] = set() - stack: list[Any] = [record] - while stack: - item = stack.pop() - identity = id(item) - if identity in seen: - continue - seen.add(identity) - if type(item) is dict: - stack.extend(item.keys()) - stack.extend(item.values()) - elif type(item) in {list, tuple, set, frozenset, deque}: - stack.extend(item) - else: - memo[identity] = item - return copy.deepcopy(record, memo) +# Upper bound on all pending records in a lane. Keeping this equal to the +# original distinct-execution cap preserves the scheduler's previous worst-case +# record count even though one execution can now retain a short burst. +_DEFAULT_MAX_PENDING_RECORDS = 1024 +# Upper bound on records waiting for one execution in one lane. This is large +# enough to preserve the known 11-progress-plus-terminal burst without relying +# on daemon-thread scheduling, while still bounding memory behind a blocked +# exporter. The in-flight record is not included in this count. +_DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION = 16 # Queue entry kinds. _RECORD = "record" @@ -251,12 +62,11 @@ class _FlushBarrier: a later, still-blocked worker skips the now-pointless flush. """ - __slots__ = ("_event", "canceled", "failed") + __slots__ = ("_event", "canceled") def __init__(self) -> None: self._event = threading.Event() self.canceled = False - self.failed = False def complete(self) -> None: self._event.set() @@ -280,13 +90,15 @@ def __init__( exporter: InsightExporter, *, max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, - max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, + max_pending_records: int = _DEFAULT_MAX_PENDING_RECORDS, + max_pending_records_per_execution: int = ( + _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION + ), ) -> 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 + self._max_pending_records = max(1, max_pending_records) + self._max_pending_per_execution = max(1, max_pending_records_per_execution) # 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 @@ -296,70 +108,43 @@ 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, 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() + # arn -> bounded FIFO of pending records. Insertion order is the ARN + # fairness order; scheduling an existing ARN moves its queue token to + # the back, and the worker requeues an ARN that has more records. + self._pending: OrderedDict[str, deque[dict[str, Any]]] = OrderedDict() 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, execution_arn: str, 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 + pending = self._pending[execution_arn] + if len(pending) >= self._max_pending_per_execution: + pending.popleft() + _logger.warning( + "workflow-insight: pending export FIFO for %s on %s is " + "full (cap=%d); dropping oldest pending record", + execution_arn, + type(self._exporter).__name__, + self._max_pending_per_execution, + ) + pending.append(record) self._pending.move_to_end(execution_arn) self._move_record_token_to_back(execution_arn) else: - self._pending[execution_arn] = (record, size) - self._pending_bytes += size + self._pending[execution_arn] = deque([record]) self._queue.append((_RECORD, execution_arn)) - self._enforce_pending_caps() + self._enforce_pending_execution_cap() + self._enforce_pending_record_cap() self._ensure_worker_locked() self._cond.notify() def enqueue_flush(self) -> _FlushBarrier: barrier = _FlushBarrier() with self._cond: - if self._disabled: - barrier.canceled = True - barrier.failed = True - barrier.complete() - return barrier self._queue.append((_FLUSH, barrier)) self._ensure_worker_locked() self._cond.notify() @@ -371,25 +156,26 @@ 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.""" + """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 - 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) + del self._queue[index] barrier.complete() return @@ -401,39 +187,53 @@ def _move_record_token_to_back(self, execution_arn: str) -> None: 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. + # No token means the ARN's last queued record is currently in flight. + # The next schedule sees it absent from ``_pending`` and appends a fresh + # FIFO plus token, preserving the in-flight record before new work. - def _enforce_pending_caps(self) -> None: + def _requeue_record_before_flush(self, execution_arn: str) -> None: + """Requeue an ARN behind peer records but before its drain barrier. + + A record token represents the ARN's whole pending FIFO at the moment the + barrier is enqueued. Consuming one record must not move the remaining + pre-barrier records behind that barrier, or invocation end could flush + and return while part of its FIFO is still waiting. + """ + for index, (kind, _) in enumerate(self._queue): + if kind == _FLUSH: + self._queue.insert(index, (_RECORD, execution_arn)) + return + self._queue.append((_RECORD, execution_arn)) + + def _enforce_pending_execution_cap(self) -> None: while len(self._pending) > self._max_pending: - old_arn, _ = self._drop_oldest_pending() + old_arn, _ = self._pending.popitem(last=False) + self._remove_record_token(old_arn) _logger.warning( - "workflow-insight: export lane for %s reached its execution cap " - "(%d); dropping pending record for %s", + "workflow-insight: export lane for %s is full " + "(cap=%d); dropping pending records 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 + + def _enforce_pending_record_cap(self) -> None: + while sum(len(records) for records in self._pending.values()) > ( + self._max_pending_records ): - old_arn, dropped_size = self._drop_oldest_pending() + old_arn = next(iter(self._pending)) + records = self._pending[old_arn] + records.popleft() _logger.warning( - "workflow-insight: export lane for %s reached its pending byte " - "budget (%d); dropping %d-byte pending record for %s", + "workflow-insight: export lane for %s reached its pending " + "record cap (%d); dropping oldest pending record for %s", type(self._exporter).__name__, - self._max_pending_bytes, - dropped_size, + self._max_pending_records, 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 + if not records: + del self._pending[old_arn] + self._remove_record_token(old_arn) def _remove_record_token(self, execution_arn: str) -> None: for index, (kind, payload) in enumerate(self._queue): @@ -441,33 +241,11 @@ 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 - self._inflight_bytes = 0 - for kind, payload in self._queue: - if kind == _FLUSH and payload is not None: - barrier: _FlushBarrier = payload - barrier.canceled = True - barrier.failed = True - barrier.complete() - self._queue.clear() - _logger.warning( - "workflow-insight: could not start worker for exporter %s; " - "disabling this lane: %s", - type(self._exporter).__name__, - exc, - ) - def _ensure_worker_locked(self) -> None: # Never create a replacement while a prior worker is alive (a blocked # worker keeps ``_worker`` non-None). A worker that exits cleanly nulls # ``_worker`` under the lock before returning, so this check is a # race-free "start iff there is no live worker". - if self._disabled: - return if self._worker is None or not self._worker.is_alive(): worker = threading.Thread( target=self._run_worker, @@ -475,10 +253,7 @@ def _ensure_worker_locked(self) -> None: daemon=True, ) self._worker = worker - try: - worker.start() - except Exception as exc: # noqa: BLE001 - instrumentation must not break hooks - self._disable_locked(exc) + worker.start() # -- worker (single daemon thread) --------------------------------------- @@ -495,38 +270,41 @@ 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: + pending = self._pending.get(payload) + if not pending: continue - record, record_size = pending - self._pending_bytes -= record_size - self._inflight_bytes += record_size + record = pending.popleft() + if pending: + # Round-robin across ARNs: one record per turn, then the + # ARN goes behind all queue entries already waiting. + self._pending.move_to_end(payload) + self._requeue_record_before_flush(payload) + else: + del self._pending[payload] 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() - if barrier is not None: - barrier.complete() + barrier: _FlushBarrier = payload + if not barrier.canceled: + self._flush() + barrier.complete() def _export_one(self, record: dict[str, Any]) -> None: exporter = self._exporter - # Copy the record's built-in containers for lane isolation, but preserve - # custom values as opaque leaves for exporter-specific rendering. This - # keeps one lane's render/truncation mutations out of other lanes without - # requiring custom-renderable values to implement ``deepcopy``. + # 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_record_containers(record) - except Exception as exc: # noqa: BLE001 - malformed containers must not break the lane + local = copy.deepcopy(record) + except Exception as exc: # noqa: BLE001 - a non-copyable payload must not alias the shared record or break the lane _logger.warning( - "workflow-insight: record container copy failed for exporter %s; " + "workflow-insight: record copy failed for exporter %s; " "skipping export for this record: %s", type(exporter).__name__, exc, @@ -572,13 +350,9 @@ def _pending_count(self) -> int: with self._cond: return len(self._pending) - def _pending_bytes_count(self) -> int: + def _pending_record_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 sum(len(records) for records in self._pending.values()) def _queue_len(self) -> int: with self._cond: @@ -597,32 +371,25 @@ def __init__( exporters: list[InsightExporter], *, max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS, - max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, + max_pending_records: int = _DEFAULT_MAX_PENDING_RECORDS, + max_pending_records_per_execution: int = ( + _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION + ), ) -> 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, + max_pending_records=max_pending_records, + max_pending_records_per_execution=max_pending_records_per_execution, ) 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.""" - 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(execution_arn, record) def end_invocation(self, timeout_seconds: float) -> bool: """Drain and flush every touched lane under one shared timeout. @@ -643,8 +410,6 @@ 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 aea400f5..26dbcbc8 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 @@ -5,21 +5,18 @@ 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 +bounded FIFO, fairness, drain, flush-ordering, timeout and thread-lifecycle invariants are asserted deterministically rather than by timing luck. """ from __future__ import annotations -import datetime -import functools +import json +import logging 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, ) @@ -138,30 +135,19 @@ 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 +class FirstExportBlockingRecorder(RecordingExporter): + """Records call order but blocks the first export until released.""" def __init__(self) -> None: + super().__init__() 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() + if not self.started.is_set(): + self.started.set() + self._release.wait(5.0) + super().export(record) def release(self) -> None: self._release.set() @@ -187,48 +173,13 @@ def flush(self) -> None: raise RuntimeError("flush boom") -class _Uncopyable(dict[str, str]): - """A JSON-serializable payload whose ``deepcopy`` raises.""" - - def __init__(self) -> None: - super().__init__({"value": "safe"}) +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") -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 -------------------------- @@ -266,25 +217,6 @@ 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() @@ -298,19 +230,19 @@ def test_repeated_scheduling_does_not_grow_threads(): assert _wait_until(lambda: not lane._worker_alive()) -# -- coalescing / fairness / isolation --------------------------------------- +# -- bounded FIFO / fairness / isolation ------------------------------------- -def test_same_execution_coalescing_exports_inflight_then_latest(): +def test_same_execution_fifo_exports_all_pending_records(): 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). + # While a1 is in flight, a2 and a3 remain ordered in the pending FIFO. 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"]) + assert _wait_until(lambda: exporter.exported_values() == ["a1", "a2", "a3"]) scheduler.end_invocation(5.0) @@ -320,16 +252,15 @@ def test_different_executions_isolated_and_fair(): 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_B, _rec(ARN_B, "b2")) # B FIFO: [b1, 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"]) + # a1 (in flight) first, then one record per ARN turn: B, A, then B again. + assert _wait_until(lambda: exporter.exported_values() == ["a1", "b1", "a2", "b2"]) scheduler.end_invocation(5.0) -def test_terminal_record_supersedes_pending_running(): +def test_terminal_record_follows_pending_running_records(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) scheduler.schedule(ARN_A, _rec(ARN_A, "r1", status="RUNNING")) @@ -337,52 +268,80 @@ def test_terminal_record_supersedes_pending_running(): 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"]) + assert _wait_until(lambda: exporter.exported_values() == ["r1", "r2", "final"]) scheduler.end_invocation(5.0) -# -- copy failure isolation --------------------------------------------------- +def test_pending_fifo_cap_drops_oldest_and_retains_terminal(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_records_per_execution=2) + scheduler.schedule(ARN_A, _rec(ARN_A, "inflight")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_A, _rec(ARN_A, "oldest")) + scheduler.schedule(ARN_A, _rec(ARN_A, "newer")) + scheduler.schedule(ARN_A, _rec(ARN_A, "final", status="SUCCEEDED")) + exporter.release() + assert _wait_until( + lambda: exporter.exported_values() == ["inflight", "newer", "final"] + ) + scheduler.end_invocation(5.0) -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] = [] +# -- copy failure isolation --------------------------------------------------- - def render(self, record: dict[str, Any]) -> Any: - value = record["payload"]["value"] - self.rendered_values.append(value) - return {"value": value} - exporter = CustomRenderExporter() +def test_deepcopy_failure_skips_record_and_lane_continues(caplog): + exporter = RecordingExporter() scheduler = _ExportScheduler([exporter]) - record = _rec(ARN_A, "before-render") - record["payload"] = _Uncopyable() - - scheduler.schedule(ARN_A, record) - scheduler.end_invocation(5.0) - - assert exporter.rendered_values == ["safe"] - assert exporter.exported_values() == ["before-render"] + # 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_uncopyable_custom_value_does_not_alias_record_containers(): +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 + record["mutated"] = True # would corrupt an aliased shared record return record exporter = MutatingRenderExporter() scheduler = _ExportScheduler([exporter]) - record = _rec(ARN_A, "custom") - record["payload"] = _Uncopyable() - - scheduler.schedule(ARN_A, record) + 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) - - assert "mutated" not in record - assert exporter.exported_values() == ["custom"] + # 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 ----------------------- @@ -439,6 +398,70 @@ def test_flush_happens_after_export(): assert kinds == ["export", "flush"] +def test_flush_barrier_waits_for_entire_pending_fifo(): + exporter = FirstExportBlockingRecorder() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + scheduler.schedule(ARN_A, _rec(ARN_A, "a3")) + barrier = lane.enqueue_flush() + exporter.release() + + assert barrier.wait(5.0) + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.calls == [ + ("export", "a1"), + ("export", "a2"), + ("export", "a3"), + ("flush", None), + ] + + +def test_record_scheduled_after_barrier_waits_for_later_flush(): + exporter = FirstExportBlockingRecorder() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + + barrier = lane.enqueue_flush() + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + exporter.release() + + assert barrier.wait(5.0) + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.calls == [ + ("export", "a1"), + ("flush", None), + ("export", "a2"), + ] + + +def test_lane_flush_includes_other_execution_records_already_queued(): + exporter = FirstExportBlockingRecorder() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) + barrier = lane.enqueue_flush() + exporter.release() + + assert barrier.wait(5.0) + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.calls == [ + ("export", "a1"), + ("export", "b1"), + ("flush", None), + ] + + def test_export_and_flush_exceptions_are_isolated(): failing = FailingExporter() good = RecordingExporter() @@ -564,26 +587,23 @@ def test_pending_execution_cap_evicts_oldest(): scheduler.end_invocation(5.0) -def test_pending_byte_budget_evicts_oldest_large_record(): +def test_pending_record_cap_preserves_original_lane_memory_bound(): exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000) + scheduler = _ExportScheduler([exporter], max_pending_records=3) 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 + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) # a1 in flight (not pending) + scheduler.schedule(ARN_B, _rec(ARN_B, "b1")) + scheduler.schedule(ARN_C, _rec(ARN_C, "c1")) + scheduler.schedule(ARN_D, _rec(ARN_D, "d1")) + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + assert lane._pending_record_count() == 3 exporter.release() scheduler.end_invocation(5.0) - exported = exporter.exported_values() - assert exported[0] == "inflight" - assert exported[1] == "c" * 1_500 + assert exporter.exported_values() == ["a1", "c1", "d1", "a2"] -def test_non_json_record_reaches_exporter_without_evicting_backlog(): +def test_unmeasurable_record_does_not_evict_existing_backlog(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) lane = scheduler._lanes[0] @@ -592,28 +612,28 @@ def test_non_json_record_reaches_exporter_without_evicting_backlog(): 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) + unmeasurable = _rec(ARN_D, "bad") + unmeasurable["payload"] = {"not-json"} + scheduler.schedule(ARN_D, unmeasurable) - assert lane._pending_count() == 3 + assert lane._pending_count() == 2 exporter.release() scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight", "b1", "c1", "custom"] + 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=5_000) + 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" * 5_000)) + scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000)) assert lane._pending_count() == 2 - assert lane._pending_bytes_count() <= 5_000 + assert lane._pending_bytes_count() <= 2_000 exporter.release() scheduler.end_invocation(5.0) exported = exporter.exported_values() @@ -621,289 +641,113 @@ 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_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(): +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) - record = _rec(ARN_B, "custom-sized") - record["payload"] = _Unsized() - scheduler.schedule(ARN_B, record) - - 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 + def fail_sizing(*args, **kwargs): + raise RecursionError("record nesting is too deep") - 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")) + monkeypatch.setattr(json, "dumps", fail_sizing) + scheduler.schedule(ARN_B, _rec(ARN_B, "too-deep")) assert scheduler._lanes[0]._pending_count() == 0 - assert scheduler._lanes[0]._worker is None + exporter.release() + scheduler.end_invocation(5.0) -def test_timed_out_barrier_flushes_eventually_and_worker_exits(): +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) + ok = scheduler.end_invocation(0.1) # times out -> barrier cancelled assert ok is False - # The caller returns on time, but one detached flush remains queued so a - # buffered exporter can publish before the worker exits idle. - assert lane._queued_flush_count() == 1 - exporter.release() - assert _wait_until(lambda: not lane._worker_alive()) - assert exporter.flushed == 1 - - -def test_timed_out_buffered_export_is_published_eventually(): - exporter = BlockingBufferedExporter() - scheduler = _ExportScheduler([exporter]) - lane = scheduler._lanes[0] - scheduler.schedule(ARN_A, _rec(ARN_A, "terminal", status="SUCCEEDED")) - assert _wait_until(exporter.started.is_set) - - assert scheduler.end_invocation(0.1) is False - assert exporter.published == [] + # 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.published == ["terminal"] - assert exporter.flushed == 1 + assert exporter.flushed == 0 # cancelled barrier did not flush def test_repeated_timeouts_behind_blocked_exporter_stay_bounded(): - """Warm timeouts coalesce to one eventual flush on the same worker.""" + """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 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}")) - 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 + 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 one token for the ARN; records are bounded inside its FIFO. + assert lane._queue_len() <= 1 + assert lane._pending_record_count() <= 16 + + # Bounded state: one in-flight ARN with a bounded pending FIFO, and no + # growing pile of barriers. + assert lane._queue_len() <= 1 assert lane._pending_count() <= 1 - assert lane._queued_flush_count() == 1 + assert lane._pending_record_count() <= 16 + 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 newest bounded window, then exits idle. exporter.release() assert _wait_until(lambda: not lane._worker_alive()) exported = exporter.exported_values() - assert exported[0] == "a1" - assert len(exported) <= 2 - assert exporter.flushed == 1 + assert exported[0] == "a1" # the in-flight record delivered first + assert len(exported) <= 17 # a1 plus at most 16 pending FIFO records + assert exported[-1] == "a51" # newest snapshot was retained + # Cancelled barriers never triggered a flush, and the idle-stop path does + # not flush either. + assert exporter.flushed == 0 -def test_cancel_flush_replaces_queued_barrier_with_detached_flush(): +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) + 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) - - assert lane._queued_flush_count() == 1 + # 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 == 1 + assert exporter.flushed == 0 assert exporter.exported_values() == ["a1"] @@ -932,79 +776,3 @@ 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 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 3cefcd32..6a6f2b0e 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,15 +365,13 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): assert rec["durationMs"] is not None and rec["durationMs"] >= 0 -# -- on-change schedules RUNNING records, coalescing intermediates (comment 2) -- +# -- on-change preserves ordinary RUNNING-record bursts (comment 2) ----------- -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. +def test_on_change_schedules_running_on_each_change_and_delivers_terminal(): + # The bounded FIFO preserves this ordinary burst without relying on the + # daemon worker receiving a turn between hooks: invocation start, both + # operation changes, then the terminal snapshot must all be observable. exporter = CaptureExporter() plugin = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") @@ -395,9 +393,7 @@ def test_on_change_schedules_running_and_delivers_terminal(): plugin.on_invocation_end(_end(operations=_ops(op1, op2))) # SUCCEEDED (terminal) statuses = [r["status"] for r in exporter.records] - assert statuses, "at least the terminal record must be delivered" - assert statuses[-1] == "SUCCEEDED" - assert set(statuses[:-1]) <= {"RUNNING"} + assert statuses == ["RUNNING", "RUNNING", "RUNNING", "SUCCEEDED"] 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 index e0a9c432..5cc278c2 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 @@ -237,6 +237,38 @@ def test_operation_change_returns_immediately_with_blocked_exporter(): plugin.on_invocation_end(_end(_ops(op))) +def test_tight_on_change_burst_preserves_all_records(): + """Fast consecutive hooks do not depend on daemon-thread scheduling. + + This mirrors the reported parallel/map fan-in shape: invocation start, ten + back-to-back operation changes, then terminal completion. The default FIFO + depth preserves all twelve records even if the worker gets no turn until + after the caller has scheduled the entire burst. + """ + exporter = _BufferedExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + plugin.on_invocation_start(_start({})) + operations: dict[str, OperationInfo] = {} + for index in range(10): + operation = _step(f"step-{index}", str(index)) + operations[operation.operation_id] = operation + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, + updated_operations=_ops(operation), + operations=dict(operations), + ) + ) + plugin.on_invocation_end(_end(operations)) + + assert [record["status"] for record in exporter.published] == [ + *("RUNNING" for _ in range(11)), + "SUCCEEDED", + ] + + # -- invocation-end drain is bounded by export_timeout_seconds ---------------- From 0805e0012f4defdeca6cc7e5105f674a172e62ba Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 22:27:22 +0000 Subject: [PATCH 02/13] fix(insight): preserve scheduler ordering --- .../_export_scheduler.py | 146 +++++++++++------- .../tests/test_export_scheduler.py | 51 +++++- 2 files changed, 134 insertions(+), 63 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 54afcb54..26447e0d 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 @@ -23,7 +23,8 @@ import logging import threading import time -from collections import OrderedDict, deque +from collections import deque +from dataclasses import dataclass from typing import Any from aws_durable_execution_sdk_python_insight.truncation import truncate_record @@ -53,6 +54,15 @@ _RECORD = "record" _FLUSH = "flush" +_RecordToken = tuple[str, int] + + +@dataclass(slots=True) +class _PendingRecord: + sequence: int + generation: int + value: dict[str, Any] + class _FlushBarrier: """A one-shot flush marker the invocation-end thread waits on. @@ -106,12 +116,15 @@ 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, (arn, generation)) or + # (_FLUSH, barrier). A generation changes whenever a flush is queued, so + # one ARN can have independent tokens on both sides of a barrier. self._queue: deque[tuple[str, Any]] = deque() - # arn -> bounded FIFO of pending records. Insertion order is the ARN - # fairness order; scheduling an existing ARN moves its queue token to - # the back, and the worker requeues an ARN that has more records. - self._pending: OrderedDict[str, deque[dict[str, Any]]] = OrderedDict() + # arn -> FIFO ordered by record sequence. Queue-token order, not this + # mapping, controls round-robin fairness. + self._pending: dict[str, deque[_PendingRecord]] = {} + self._generation = 0 + self._next_sequence = 0 self._stop_when_idle = False self._worker: threading.Thread | None = None @@ -120,23 +133,30 @@ def __init__( 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: - pending = self._pending[execution_arn] - if len(pending) >= self._max_pending_per_execution: - pending.popleft() - _logger.warning( - "workflow-insight: pending export FIFO for %s on %s is " - "full (cap=%d); dropping oldest pending record", - execution_arn, - type(self._exporter).__name__, - self._max_pending_per_execution, - ) - pending.append(record) - self._pending.move_to_end(execution_arn) - self._move_record_token_to_back(execution_arn) + pending = self._pending.get(execution_arn) + if pending is None: + pending = deque() + self._pending[execution_arn] = pending + elif len(pending) >= self._max_pending_per_execution: + self._drop_oldest_pending_record(execution_arn) + pending = self._pending.setdefault(execution_arn, deque()) + _logger.warning( + "workflow-insight: pending export FIFO for %s on %s is " + "full (cap=%d); dropping oldest pending record", + execution_arn, + type(self._exporter).__name__, + self._max_pending_per_execution, + ) + + generation = self._generation + has_generation = bool(pending and pending[-1].generation == generation) + pending.append(_PendingRecord(self._next_sequence, generation, record)) + self._next_sequence += 1 + token = (execution_arn, generation) + if has_generation: + self._move_record_token_to_back(token) else: - self._pending[execution_arn] = deque([record]) - self._queue.append((_RECORD, execution_arn)) + self._queue.append((_RECORD, token)) self._enforce_pending_execution_cap() self._enforce_pending_record_cap() self._ensure_worker_locked() @@ -146,6 +166,7 @@ def enqueue_flush(self) -> _FlushBarrier: barrier = _FlushBarrier() with self._cond: self._queue.append((_FLUSH, barrier)) + self._generation += 1 self._ensure_worker_locked() self._cond.notify() return barrier @@ -181,34 +202,31 @@ def cancel_flush(self, barrier: _FlushBarrier) -> None: # -- queue bookkeeping (must hold ``_cond``) ------------------------------ - def _move_record_token_to_back(self, execution_arn: str) -> None: + def _move_record_token_to_back(self, token: _RecordToken) -> None: for index, (kind, payload) in enumerate(self._queue): - if kind == _RECORD and payload == execution_arn: + if kind == _RECORD and payload == token: del self._queue[index] - self._queue.append((_RECORD, execution_arn)) + self._queue.append((_RECORD, token)) return - # No token means the ARN's last queued record is currently in flight. - # The next schedule sees it absent from ``_pending`` and appends a fresh - # FIFO plus token, preserving the in-flight record before new work. - - def _requeue_record_before_flush(self, execution_arn: str) -> None: - """Requeue an ARN behind peer records but before its drain barrier. - A record token represents the ARN's whole pending FIFO at the moment the - barrier is enqueued. Consuming one record must not move the remaining - pre-barrier records behind that barrier, or invocation end could flush - and return while part of its FIFO is still waiting. - """ + def _requeue_record_before_flush(self, token: _RecordToken) -> None: + """Requeue one generation behind peers but before its flush barrier.""" for index, (kind, _) in enumerate(self._queue): if kind == _FLUSH: - self._queue.insert(index, (_RECORD, execution_arn)) + self._queue.insert(index, (_RECORD, token)) return - self._queue.append((_RECORD, execution_arn)) + self._queue.append((_RECORD, token)) + + def _oldest_pending_arn(self) -> str: + return min( + self._pending, + key=lambda arn: self._pending[arn][0].sequence, + ) def _enforce_pending_execution_cap(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._oldest_pending_arn() + self._drop_pending_execution(old_arn) _logger.warning( "workflow-insight: export lane for %s is full " "(cap=%d); dropping pending records for %s", @@ -221,9 +239,8 @@ def _enforce_pending_record_cap(self) -> None: while sum(len(records) for records in self._pending.values()) > ( self._max_pending_records ): - old_arn = next(iter(self._pending)) - records = self._pending[old_arn] - records.popleft() + old_arn = self._oldest_pending_arn() + self._drop_oldest_pending_record(old_arn) _logger.warning( "workflow-insight: export lane for %s reached its pending " "record cap (%d); dropping oldest pending record for %s", @@ -231,13 +248,26 @@ def _enforce_pending_record_cap(self) -> None: self._max_pending_records, old_arn, ) - if not records: - del self._pending[old_arn] - self._remove_record_token(old_arn) - def _remove_record_token(self, execution_arn: str) -> None: + def _drop_oldest_pending_record(self, execution_arn: str) -> None: + records = self._pending[execution_arn] + dropped = records.popleft() + token = (execution_arn, dropped.generation) + if not records or records[0].generation != dropped.generation: + self._remove_record_token(token) + if not records: + del self._pending[execution_arn] + + def _drop_pending_execution(self, execution_arn: str) -> None: + del self._pending[execution_arn] + for index in range(len(self._queue) - 1, -1, -1): + kind, payload = self._queue[index] + if kind == _RECORD and payload[0] == execution_arn: + del self._queue[index] + + def _remove_record_token(self, token: _RecordToken) -> None: for index, (kind, payload) in enumerate(self._queue): - if kind == _RECORD and payload == execution_arn: + if kind == _RECORD and payload == token: del self._queue[index] return @@ -271,17 +301,17 @@ def _run_worker(self) -> None: kind, payload = self._queue.popleft() record: dict[str, Any] | None = None if kind == _RECORD: - pending = self._pending.get(payload) - if not pending: + token: _RecordToken = payload + execution_arn, generation = token + pending = self._pending.get(execution_arn) + if not pending or pending[0].generation != generation: continue - record = pending.popleft() - if pending: - # Round-robin across ARNs: one record per turn, then the - # ARN goes behind all queue entries already waiting. - self._pending.move_to_end(payload) - self._requeue_record_before_flush(payload) - else: - del self._pending[payload] + record = pending.popleft().value + if pending and pending[0].generation == generation: + # One record per ARN turn within this barrier generation. + self._requeue_record_before_flush(token) + elif not pending: + del self._pending[execution_arn] if kind == _RECORD and record is not None: self._export_one(record) 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 26dbcbc8..54a7c7dc 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 @@ -420,6 +420,29 @@ def test_flush_barrier_waits_for_entire_pending_fifo(): ] +def test_flush_barrier_splits_same_execution_fifo_by_generation(): + exporter = FirstExportBlockingRecorder() + scheduler = _ExportScheduler([exporter]) + lane = scheduler._lanes[0] + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + barrier = lane.enqueue_flush() + scheduler.schedule(ARN_A, _rec(ARN_A, "a3")) + exporter.release() + + assert barrier.wait(5.0) + lane.request_stop_when_idle() + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.calls == [ + ("export", "a1"), + ("export", "a2"), + ("flush", None), + ("export", "a3"), + ] + + def test_record_scheduled_after_barrier_waits_for_later_flush(): exporter = FirstExportBlockingRecorder() scheduler = _ExportScheduler([exporter]) @@ -658,6 +681,23 @@ def fail_sizing(*args, **kwargs): scheduler.end_invocation(5.0) +def test_pending_record_cap_evicts_true_oldest_across_arns(): + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter], max_pending_records=2) + scheduler.schedule(ARN_C, _rec(ARN_C, "inflight")) + assert _wait_until(exporter.started.is_set) + + scheduler.schedule(ARN_A, _rec(ARN_A, "a1")) + scheduler.schedule(ARN_B, _rec(ARN_B, "b1", status="SUCCEEDED")) + scheduler.schedule(ARN_A, _rec(ARN_A, "a2")) + + exporter.release() + scheduler.end_invocation(5.0) + # A1 is globally oldest. B1 remains even though scheduling A2 moved A's + # fairness token behind B's token. + assert exporter.exported_values() == ["inflight", "b1", "a2"] + + def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) @@ -698,13 +738,14 @@ def test_repeated_timeouts_behind_blocked_exporter_stay_bounded(): # 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 one token for the ARN; records are bounded inside its FIFO. - assert lane._queue_len() <= 1 + # Each cancelled barrier can leave one generation token, but both + # records and tokens remain bounded by the per-execution FIFO depth. + assert lane._queue_len() <= 16 assert lane._pending_record_count() <= 16 - # Bounded state: one in-flight ARN with a bounded pending FIFO, and no - # growing pile of barriers. - assert lane._queue_len() <= 1 + # Bounded state: one in-flight ARN with a bounded pending FIFO and bounded + # generation tokens, with no growing pile of barriers. + assert lane._queue_len() <= 16 assert lane._pending_count() <= 1 assert lane._pending_record_count() <= 16 assert lane._queued_flush_count() == 0 From 912f8ef72d63947f4ae995af3dc626571433d9e1 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 9 Sep 2026 23:36:40 +0000 Subject: [PATCH 03/13] fix(insight): preserve canceled generation order --- .../README.md | 14 ++-- .../_export_scheduler.py | 80 ++++++++++++++++--- .../tests/test_export_scheduler.py | 52 +++++++++++- 3 files changed, 124 insertions(+), 22 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 38f9f02a..a48e9849 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -59,14 +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. Each lane keeps a -> bounded FIFO of up to 16 pending snapshots per execution and at most 1,024 -> pending snapshots across the lane, preserving ordinary bursts without relying -> on daemon-thread scheduling. If an exporter remains slower than the producer -> and either bound fills, the oldest pending snapshot is dropped so the newest progress and terminal snapshots are retained. At +> 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 are fine. +> Each lane keeps up to 16 pending snapshots per execution, 1,024 records total, +> and 16 MB of estimated canonical JSON. When a bound fills, it drops the oldest +> pending snapshot so recent progress and terminal snapshots are retained. 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 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 26447e0d..5ed25084 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 @@ -10,8 +10,8 @@ Each exporter lane: * Keeps a bounded FIFO per execution ARN and processes ARNs round-robin. -* Drops the oldest pending snapshots when the per-execution or lane-wide limit - is reached, favoring the newest progress and terminal snapshots. +* Drops the oldest pending snapshots when the per-execution, lane-wide record, + or byte limit is reached, favoring recent progress and terminal snapshots. * 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 dataclasses import dataclass 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 @@ -50,6 +53,10 @@ # exporter. The in-flight record is not included in this count. _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION = 16 +# Canonical JSON-byte estimate retained by one blocked lane. This keeps the +# instrumentation backlog well below Lambda's 128 MiB memory floor. +_DEFAULT_MAX_PENDING_BYTES = 16_000_000 + # Queue entry kinds. _RECORD = "record" _FLUSH = "flush" @@ -62,6 +69,7 @@ class _PendingRecord: sequence: int generation: int value: dict[str, Any] + size: int class _FlushBarrier: @@ -104,11 +112,14 @@ def __init__( max_pending_records_per_execution: int = ( _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION ), + max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, ) -> None: self._exporter = exporter self._max_pending = max(1, max_pending_executions) self._max_pending_records = max(1, max_pending_records) self._max_pending_per_execution = max(1, max_pending_records_per_execution) + 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 @@ -130,9 +141,19 @@ def __init__( # -- 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) + ) pending = self._pending.get(execution_arn) if pending is None: pending = deque() @@ -150,7 +171,10 @@ def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: generation = self._generation has_generation = bool(pending and pending[-1].generation == generation) - pending.append(_PendingRecord(self._next_sequence, generation, record)) + pending.append( + _PendingRecord(self._next_sequence, generation, record, size) + ) + self._pending_bytes += size self._next_sequence += 1 token = (execution_arn, generation) if has_generation: @@ -159,6 +183,7 @@ def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: self._queue.append((_RECORD, token)) self._enforce_pending_execution_cap() self._enforce_pending_record_cap() + self._enforce_pending_byte_cap() self._ensure_worker_locked() self._cond.notify() @@ -210,9 +235,15 @@ def _move_record_token_to_back(self, token: _RecordToken) -> None: return def _requeue_record_before_flush(self, token: _RecordToken) -> None: - """Requeue one generation behind peers but before its flush barrier.""" - for index, (kind, _) in enumerate(self._queue): - if kind == _FLUSH: + """Requeue behind peers without crossing this ARN's next generation.""" + execution_arn, generation = token + for index, (kind, payload) in enumerate(self._queue): + later_same_arn = ( + kind == _RECORD + and payload[0] == execution_arn + and payload[1] > generation + ) + if kind == _FLUSH or later_same_arn: self._queue.insert(index, (_RECORD, token)) return self._queue.append((_RECORD, token)) @@ -249,17 +280,33 @@ def _enforce_pending_record_cap(self) -> None: old_arn, ) - def _drop_oldest_pending_record(self, execution_arn: str) -> None: + def _enforce_pending_byte_cap(self) -> None: + while self._pending_bytes > self._max_pending_bytes and self._pending: + old_arn = self._oldest_pending_arn() + dropped_size = self._drop_oldest_pending_record(old_arn) + _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_record(self, execution_arn: str) -> int: records = self._pending[execution_arn] dropped = records.popleft() + self._pending_bytes -= dropped.size token = (execution_arn, dropped.generation) if not records or records[0].generation != dropped.generation: self._remove_record_token(token) if not records: del self._pending[execution_arn] + return dropped.size def _drop_pending_execution(self, execution_arn: str) -> None: - del self._pending[execution_arn] + records = self._pending.pop(execution_arn) + self._pending_bytes -= sum(record.size for record in records) for index in range(len(self._queue) - 1, -1, -1): kind, payload = self._queue[index] if kind == _RECORD and payload[0] == execution_arn: @@ -306,7 +353,9 @@ def _run_worker(self) -> None: pending = self._pending.get(execution_arn) if not pending or pending[0].generation != generation: continue - record = pending.popleft().value + pending_record = pending.popleft() + self._pending_bytes -= pending_record.size + record = pending_record.value if pending and pending[0].generation == generation: # One record per ARN turn within this barrier generation. self._requeue_record_before_flush(token) @@ -384,6 +433,10 @@ def _pending_record_count(self) -> int: with self._cond: return sum(len(records) for records in self._pending.values()) + 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) @@ -405,6 +458,7 @@ def __init__( max_pending_records_per_execution: int = ( _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION ), + max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES, ) -> None: self._lanes = [ _ExporterLane( @@ -412,14 +466,16 @@ def __init__( max_pending_executions=max_pending_executions, max_pending_records=max_pending_records, max_pending_records_per_execution=max_pending_records_per_execution, + 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/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 54a7c7dc..440215a4 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 @@ -173,8 +173,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") @@ -698,6 +701,51 @@ def test_pending_record_cap_evicts_true_oldest_across_arns(): assert exporter.exported_values() == ["inflight", "b1", "a2"] +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_record_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_preserves_generation_order_and_terminal(): + 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_A, _rec(ARN_A, "pre-1")) + scheduler.schedule(ARN_A, _rec(ARN_A, "pre-2")) + barrier = lane.enqueue_flush() + scheduler.schedule(ARN_A, _rec(ARN_A, "terminal", status="SUCCEEDED")) + lane.cancel_flush(barrier) + lane.request_stop_when_idle() + exporter.release() + + assert _wait_until(lambda: not lane._worker_alive()) + assert exporter.exported_values() == [ + "inflight", + "pre-1", + "pre-2", + "terminal", + ] + assert lane._pending_record_count() == 0 + assert lane._queue_len() == 0 + + def test_cancelled_barrier_is_cleaned_up_and_worker_exits(): exporter = BlockingExporter() scheduler = _ExportScheduler([exporter]) From 64a68e310ac3cbd9bc7736f9eae5b9ca16ca47cc Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 00:02:19 +0000 Subject: [PATCH 04/13] fix(insight): isolate rejected FIFO records --- .../_export_scheduler.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 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 5ed25084..831a52ed 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 @@ -149,11 +149,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 pending = self._pending.get(execution_arn) if pending is None: pending = deque() From 6d818abc25f77a7d231f2043603fa5d41d3efab1 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 00:23:39 +0000 Subject: [PATCH 05/13] test(insight): cover real fan-in bursts --- .../tests/e2e/parallel_fan_in_int_test.py | 124 ++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-insight/tests/e2e/parallel_fan_in_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/e2e/parallel_fan_in_int_test.py b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/parallel_fan_in_int_test.py new file mode 100644 index 00000000..950f58e8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/parallel_fan_in_int_test.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end test for on-change export bursts during parallel fan-in.""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) + +from aws_durable_execution_sdk_python_insight import ( + WorkflowInsightConfig, + workflow_insight, +) +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) + + +_BRANCH_COUNT = 6 + + +class _BlockingCaptureExporter: + """Blocks the first export so the real hook burst queues deterministically.""" + + def __init__(self) -> None: + self.max_record_size_bytes: int | None = None + self.started = threading.Event() + self.release = threading.Event() + self.records: list[dict[str, Any]] = [] + self._lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + if not self.started.is_set(): + self.started.set() + self.release.wait(10.0) + with self._lock: + self.records.append(record) + + def flush(self) -> None: + return None + + def snapshots(self) -> list[dict[str, Any]]: + with self._lock: + return list(self.records) + + +def _branch(index: int): + def run(context: DurableContext) -> int: + return context.step(lambda _step_context: index, name=f"step-{index}") + + return run + + +def _parallel_handler(event: Any, context: DurableContext) -> list[int]: # noqa: ARG001 + return context.parallel( + [_branch(index) for index in range(_BRANCH_COUNT)], + name="fan-in", + ).get_results() + + +def _wait_until(predicate, timeout: float = 10.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def test_parallel_fan_in_preserves_every_on_change_snapshot() -> None: + capture = _BlockingCaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig( + exporters=[capture], + emit_mode="on-change", + operation_detail="full-tree", + ) + ) + handler = durable_execution(_parallel_handler, plugins=[plugin]) + results: list[DurableFunctionTestResult] = [] + + with DurableFunctionTestRunner(handler=handler, execution_timeout=15) as runner: + run_thread = threading.Thread( + target=lambda: results.append(runner.run(input="{}")), + daemon=True, + ) + run_thread.start() + try: + assert capture.started.wait(5.0) + lane = plugin._scheduler._lanes[0] + # Invocation start is in flight. Three real PluginExecutor changes + # and the terminal record must queue behind it before release. + assert _wait_until(lambda: lane._pending_record_count() == 4) + finally: + capture.release.set() + run_thread.join(10.0) + + assert not run_thread.is_alive() + assert len(results) == 1 + assert results[0].status is InvocationStatus.SUCCEEDED + + records = capture.snapshots() + assert [record["status"] for record in records] == [ + "RUNNING", + "RUNNING", + "RUNNING", + "RUNNING", + "SUCCEEDED", + ] + final_names = {operation["name"] for operation in records[-1]["operations"]} + assert "fan-in" in final_names + assert {f"step-{index}" for index in range(_BRANCH_COUNT)} <= final_names From 7e8d506fd392cd5dbd5684d81af4e2fafc7e3f9f Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 00:52:31 +0000 Subject: [PATCH 06/13] fix(insight): preserve shaped record admission --- .../README.md | 2 +- .../_export_scheduler.py | 89 +++++++++++++++---- .../tests/test_export_scheduler.py | 60 +++++++++---- 3 files changed, 115 insertions(+), 36 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index a48e9849..0febae3a 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 are fine. > Each lane keeps up to 16 pending snapshots per execution, 1,024 records total, -> and 16 MB of estimated canonical JSON. When a bound fills, it drops the oldest +> and 16 MB of estimated retained memory. When a bound fills, it drops the oldest > pending snapshot so recent progress and terminal snapshots are retained. At > invocation end the plugin drains and flushes the touched exporters under a > single shared deadline 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 831a52ed..d8f06e0b 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,16 +21,14 @@ import copy import logging +import sys import threading import time from collections import deque from dataclasses import dataclass 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 @@ -53,10 +51,40 @@ # exporter. The in-flight record is not included in this count. _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION = 16 -# Canonical JSON-byte estimate retained by one blocked lane. This keeps the +# Estimated Python object memory retained by one blocked lane. This keeps the # instrumentation backlog well below Lambda's 128 MiB memory floor. _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" @@ -80,11 +108,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() @@ -138,6 +167,7 @@ def __init__( self._next_sequence = 0 self._stop_when_idle = False self._worker: threading.Thread | None = None + self._disabled = False # -- producer API (checkpoint / invocation-end threads) ------------------- @@ -145,18 +175,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( @@ -204,6 +228,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._generation += 1 self._ensure_worker_locked() @@ -332,11 +361,32 @@ def _remove_record_token(self, token: _RecordToken) -> 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, @@ -344,7 +394,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) --------------------------------------- @@ -487,7 +540,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) @@ -510,6 +563,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 440215a4..ef68160f 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 @@ -183,6 +182,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 -------------------------- @@ -220,6 +226,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() @@ -629,7 +654,7 @@ def test_pending_record_cap_preserves_original_lane_memory_bound(): assert exporter.exported_values() == ["a1", "c1", "d1", "a2"] -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] @@ -638,19 +663,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) @@ -659,7 +684,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() @@ -667,21 +692,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_pending_record_cap_evicts_true_oldest_across_arns(): @@ -703,7 +727,7 @@ def test_pending_record_cap_evicts_true_oldest_across_arns(): 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) @@ -712,7 +736,7 @@ def test_pending_byte_budget_evicts_oldest_large_record(): scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 1_500)) assert lane._pending_record_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() From 1d95785590e07612665ee8cf940785fe18e6601e Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 01:24:08 +0000 Subject: [PATCH 07/13] fix(insight): preserve bounded FIFO records --- .../_export_scheduler.py | 73 +++++++--- .../tests/test_export_scheduler.py | 125 ++++++++++++------ 2 files changed, 137 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 d8f06e0b..f1c59bf7 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 @@ -71,18 +71,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. @@ -183,13 +219,17 @@ def schedule( self._stop_when_idle = False size = max(0, record_size) if size > self._max_pending_bytes: + superseded = execution_arn in self._pending + if superseded: + self._drop_pending_execution(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 FIFO" if superseded else "", ) return pending = self._pending.get(execution_arn) @@ -439,18 +479,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 ef68160f..8675b060 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 @@ -182,6 +181,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.""" @@ -318,58 +329,45 @@ def test_pending_fifo_cap_drops_oldest_and_retains_terminal(): # -- 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 ----------------------- @@ -692,6 +690,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 59699c8a05611e49eff86e69663187076da8865a Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 01:56:22 +0000 Subject: [PATCH 08/13] fix(insight): preserve timed-out FIFO exports --- .../_export_scheduler.py | 48 +++--- .../tests/test_export_scheduler.py | 157 +++++++++++++----- 2 files changed, 137 insertions(+), 68 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 f1c59bf7..b7685de7 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 @@ -56,7 +56,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() @@ -71,8 +71,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)): @@ -285,26 +289,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 @@ -407,7 +404,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 @@ -472,10 +469,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 @@ -564,20 +561,21 @@ def __init__( ), 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_records=max_pending_records, max_pending_records_per_execution=max_pending_records_per_execution, - 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 8675b060..62a2f9e5 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 @@ -151,6 +151,35 @@ def release(self) -> None: self._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.""" @@ -200,6 +229,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 -------------------------- @@ -731,6 +770,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]) @@ -809,97 +884,93 @@ def test_cancelled_barrier_preserves_generation_order_and_terminal(): assert lane._queue_len() == 0 -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 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 - # Each cancelled barrier can leave one generation token, but both - # records and tokens remain bounded by the per-execution FIFO depth. - assert lane._queue_len() <= 16 + assert scheduler.end_invocation(0.02) is False + # At most 16 generation tokens plus one detached eventual flush. + assert lane._queued_flush_count() == 1 + assert lane._queue_len() <= 17 assert lane._pending_record_count() <= 16 - # Bounded state: one in-flight ARN with a bounded pending FIFO and bounded - # generation tokens, with no growing pile of barriers. - assert lane._queue_len() <= 16 + assert lane._queue_len() <= 17 assert lane._pending_count() <= 1 assert lane._pending_record_count() <= 16 - 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 newest bounded window, 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) <= 17 # a1 plus at most 16 pending FIFO records - assert exported[-1] == "a51" # newest snapshot was retained - # 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) <= 17 + assert exported[-1] == "a51" + 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 b0deb13feaf7e698b21660e1787dfa8153e91ae3 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:17:16 +0000 Subject: [PATCH 09/13] fix(insight): bound opaque FIFO graphs --- .../_export_scheduler.py | 51 +++++++++---------- .../tests/test_export_scheduler.py | 41 ++++++++++++++- 2 files changed, 62 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 b7685de7..e10d1a50 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 @@ -20,10 +20,12 @@ from __future__ import annotations import copy +import gc import logging import sys import threading import time +import types from collections import deque from dataclasses import dataclass from typing import Any @@ -56,8 +58,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] @@ -74,33 +85,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/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 62a2f9e5..6c2dc5ae 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 @@ -712,7 +713,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) @@ -721,7 +722,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() @@ -884,6 +885,42 @@ def test_cancelled_barrier_preserves_generation_order_and_terminal(): assert lane._queue_len() == 0 +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 e6701e655a0a4135f1ddd8391383866196523fbe Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:39:47 +0000 Subject: [PATCH 10/13] fix(insight): bound FIFO traversal work --- .../_export_scheduler.py | 148 +++++++++++++++--- .../tests/test_export_scheduler.py | 32 ++++ 2 files changed, 159 insertions(+), 21 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 e10d1a50..e8d8cb45 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 @@ -20,7 +20,8 @@ from __future__ import annotations import copy -import gc +import functools +import itertools import logging import sys import threading @@ -58,44 +59,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/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 6c2dc5ae..f7ff94a3 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 ( @@ -235,6 +236,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__() @@ -921,6 +925,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 6f9986bccc9924e527d255fdbd77ddd45bd43f6b Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 02:53:06 +0000 Subject: [PATCH 11/13] fix(insight): preserve accepted FIFO snapshots --- .../_export_scheduler.py | 6 +----- .../tests/test_export_scheduler.py | 9 +++++---- 2 files changed, 6 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 e8d8cb45..6f4b49eb 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 @@ -324,17 +324,13 @@ def schedule( self._stop_when_idle = False size = max(0, record_size) if size > self._max_pending_bytes: - superseded = execution_arn in self._pending - if superseded: - self._drop_pending_execution(execution_arn) _logger.warning( "workflow-insight: pending record for %s on %s exceeds the " - "byte budget (%d > %d); dropping this record%s", + "byte budget (%d > %d); dropping this record", execution_arn, type(self._exporter).__name__, size, self._max_pending_bytes, - " and its superseded pending FIFO" if superseded else "", ) return pending = self._pending.get(execution_arn) 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 f7ff94a3..ddd972d9 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 @@ -734,24 +734,25 @@ 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(): +def test_over_budget_record_preserves_accepted_same_arn_fifo(): 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_A, _rec(ARN_A, "accepted-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_count() == 2 + assert lane._pending_record_count() == 2 assert lane._pending_bytes_count() <= 3_500 exporter.release() scheduler.end_invocation(5.0) - assert exporter.exported_values() == ["inflight", "unrelated"] + assert exporter.exported_values() == ["inflight", "accepted-running", "unrelated"] def test_retained_size_traverses_slots_after_shallow_size_failure(): From 6dfa457304d78d52b8b0a6616e690956d61a4972 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 03:00:02 +0000 Subject: [PATCH 12/13] fix(insight): preserve bounded FIFO 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 6f4b49eb..7b88ca23 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 @@ -20,6 +20,8 @@ from __future__ import annotations import copy +import datetime +import decimal import functools import itertools import logging @@ -27,6 +29,7 @@ import threading import time import types +import uuid from collections import deque from dataclasses import dataclass from typing import Any @@ -79,6 +82,16 @@ ) +_SAFE_OPAQUE_RETAINED_TYPES = ( + datetime.date, + datetime.datetime, + datetime.time, + datetime.timedelta, + decimal.Decimal, + uuid.UUID, +) + + class _RetainedChildren: __slots__ = ("iterator",) @@ -123,9 +136,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 (): @@ -133,19 +146,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 @@ -672,7 +695,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 ddd972d9..bcb27d73 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, ) @@ -900,12 +903,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: @@ -954,6 +963,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 d876e48a82b3b04bac18689b965fcf561f007b74 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Thu, 10 Sep 2026 03:22:40 +0000 Subject: [PATCH 13/13] fix(insight): bound FIFO in-flight retention --- .../_export_scheduler.py | 67 +++++++++------ .../tests/test_export_scheduler.py | 82 ++++++++++++++++++- 2 files changed, 123 insertions(+), 26 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 7b88ca23..f1de78b5 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 @@ -25,7 +25,6 @@ import functools import itertools import logging -import sys import threading import time import types @@ -73,7 +72,6 @@ bool, type(None), range, - slice, type, types.ModuleType, types.CodeType, @@ -135,6 +133,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): @@ -175,26 +175,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: @@ -313,6 +312,7 @@ def __init__( self._max_pending_per_execution = max(1, max_pending_records_per_execution) 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 @@ -412,6 +412,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. @@ -481,7 +487,10 @@ def _enforce_pending_record_cap(self) -> None: ) def _enforce_pending_byte_cap(self) -> None: - 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 = self._oldest_pending_arn() dropped_size = self._drop_oldest_pending_record(old_arn) _logger.warning( @@ -523,6 +532,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 @@ -571,6 +581,7 @@ def _run_worker(self) -> None: return kind, payload = self._queue.popleft() record: dict[str, Any] | None = None + record_size = 0 if kind == _RECORD: token: _RecordToken = payload execution_arn, generation = token @@ -578,7 +589,9 @@ def _run_worker(self) -> None: if not pending or pending[0].generation != generation: continue pending_record = pending.popleft() - self._pending_bytes -= pending_record.size + record_size = pending_record.size + self._pending_bytes -= record_size + self._inflight_bytes += record_size record = pending_record.value if pending and pending[0].generation == generation: # One record per ARN turn within this barrier generation. @@ -587,7 +600,11 @@ def _run_worker(self) -> None: del self._pending[execution_arn] 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() @@ -658,6 +675,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 bcb27d73..1ae5b5cd 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 @@ -720,16 +720,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() @@ -1114,3 +1114,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