From 8a595391e1a848c613adbef2d89617f8807d5aa1 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 15 Sep 2026 00:16:51 +0000 Subject: [PATCH] fix(insight): keep one pending export snapshot per execution The export scheduler held a single pending slot for the whole plugin. When two executions were in flight on one plugin instance, execution B's RUNNING snapshot could replace execution A's SUCCEEDED snapshot before the worker claimed it. A's drain() then returned once B's record was flushed, and A's terminal record was never exported. Lambda runs one invocation per environment, but the local test runner and the conformance suite drive independent executions concurrently through one shared plugin, where this lost 30-37% of terminal records in a two-thread probe. Key the pending map by executionArn. A snapshot only replaces its own execution's entry, the worker exports entries in first-arrival order, and a flush request is honored only once every pending record is exported, so drain() means everything scheduled by any execution has been delivered. A RUNNING snapshot never replaces a pending terminal snapshot for the same execution. Close the execution state in on_invocation_end before emitting and draining. An operation-change hook from a checkpoint completing during the drain, or arriving after the state was discarded, is dropped instead of recreating state and emitting a trailing RUNNING record. When the worker thread cannot be started, drain() now exports pending records inline on the calling thread instead of disabling export for the plugin's lifetime and dropping the record. --- .../README.md | 12 +- .../_export_scheduler.py | 161 ++++++++++++------ .../plugin.py | 33 +++- .../tests/test_export_scheduler.py | 89 +++++++++- .../tests/test_plugin.py | 130 ++++++++++++++ 5 files changed, 358 insertions(+), 67 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index be1027df..8aafc469 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -57,10 +57,14 @@ Behavior is validated cross-SDK by the `insight` conformance suite > **Note (asynchronous export).** Export rendering, truncation, `export()`, and > `flush()` run on one lazy background worker per plugin. Checkpoint hooks only -> replace the latest pending snapshot and wake the worker. Consecutive -> `on-change` snapshots may coalesce while an export is in flight. An invocation -> that emits a record drains the latest snapshot and flushes exporters before it -> returns; invocations that emit nothing do not start or flush the worker. +> replace the latest pending snapshot **for their own execution** and wake the +> worker; executions in flight at the same time (as under the local test runner) +> never displace each other's snapshots, and a terminal snapshot is never +> replaced by a later `RUNNING` one. Consecutive `on-change` snapshots of one +> execution may coalesce while an export is in flight. An invocation that emits +> a record drains every pending snapshot and flushes exporters before it +> returns; invocations that emit nothing do not start or flush the worker. If +> the worker thread cannot be started, that drain exports inline instead. ## 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 ef22ea54..8e6f24fe 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -1,7 +1,17 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -"""Latest-pending asynchronous export scheduling for Workflow Insight.""" +"""Latest-pending asynchronous export scheduling for Workflow Insight. + +One pending slot is kept **per execution** (keyed by ``executionArn``). Each +record is a complete snapshot of its execution, so a newer snapshot for the same +execution supersedes an older one that has not been exported yet, while records +for different executions never displace each other. The plugin already tracks +state per execution, and the local test runner drives independent executions +concurrently through one shared plugin instance, so a single plugin-wide slot +would silently drop one execution's terminal record whenever another execution +scheduled a snapshot first. +""" from __future__ import annotations @@ -15,70 +25,81 @@ _logger = logging.getLogger("aws_durable_execution_sdk_python_insight") +_TERMINAL_STATUSES = frozenset({"SUCCEEDED", "FAILED"}) + + +def _is_terminal(record: dict[str, Any]) -> bool: + return record.get("status") in _TERMINAL_STATUSES + class _ExportScheduler: - """Run all exporters on one lazy worker with one latest pending record.""" + """Run all exporters on one lazy worker with one pending record per execution.""" def __init__(self, exporters: list[InsightExporter]) -> None: self._exporters = exporters self._condition = threading.Condition(threading.Lock()) - self._pending: dict[str, Any] | None = None + # executionArn -> latest pending snapshot for that execution. Insertion + # ordered, so the worker exports executions in first-arrival order; + # replacing an entry keeps its position. + self._pending: dict[str, dict[str, Any]] = {} self._flush_requested = False self._flush_event: threading.Event | None = None self._worker: threading.Thread | None = None - self._disabled = False + self._start_failure_logged = False def schedule(self, record: dict[str, Any]) -> None: - """Replace the pending snapshot and return without running exporters.""" + """Replace this execution's pending snapshot and return without exporting.""" + key = str(record.get("executionArn", "")) displaced: dict[str, Any] | None = None - failed_pending: dict[str, Any] | None = None start_error: Exception | None = None with self._condition: - if self._disabled: + displaced = self._pending.get(key) + # A terminal snapshot is final. A RUNNING snapshot for the same + # execution that arrives after it (an operation-change hook from a + # checkpoint completing during the end-of-invocation drain) must not + # replace it, or the execution would be reported as still running. + if ( + displaced is not None + and _is_terminal(displaced) + and not _is_terminal(record) + ): return - displaced = self._pending - self._pending = record - failed_pending, start_error = self._ensure_worker_locked() + self._pending[key] = record + start_error = self._ensure_worker_locked() self._condition.notify() - # Releasing either record may run custom finalizers, so do it unlocked. - del displaced, failed_pending - if start_error is not None: - _logger.warning( - "workflow-insight: could not start export worker; disabling " - "asynchronous export: %s", - start_error, - ) + # Releasing the displaced record may run custom finalizers, so do it unlocked. + del displaced + self._log_start_failure(start_error) def drain(self) -> None: - """Wait until the latest pending record is exported and exporters flush.""" - failed_pending: dict[str, Any] | None = None + """Wait until every pending record is exported and exporters flush. + + Records scheduled by any execution are exported before the flush, so a + drain issued at one execution's invocation end also delivers snapshots + that a concurrently running execution scheduled earlier. + """ start_error: Exception | None = None with self._condition: - if self._disabled: - return if not self._flush_requested: self._flush_requested = True self._flush_event = threading.Event() flush_event = self._flush_event assert flush_event is not None - failed_pending, start_error = self._ensure_worker_locked() - started = not self._disabled + start_error = self._ensure_worker_locked() + worker_running = self._worker is not None self._condition.notify() - del failed_pending - if start_error is not None: - _logger.warning( - "workflow-insight: could not start export worker; disabling " - "asynchronous export: %s", - start_error, - ) - if started: + self._log_start_failure(start_error) + if worker_running: flush_event.wait() + return + # No worker could be started. Export and flush on the calling thread so + # nothing scheduled is dropped; this is the invocation-end path, which + # already waits for delivery. + self._pump(flush_event) - def _ensure_worker_locked( - self, - ) -> tuple[dict[str, Any] | None, Exception | None]: + def _ensure_worker_locked(self) -> Exception | None: if self._worker is not None and self._worker.is_alive(): - return None, None + return None worker = threading.Thread( target=self._run, name=f"workflow-insight-export-{id(self)}", @@ -88,32 +109,45 @@ def _ensure_worker_locked( try: worker.start() except Exception as exc: # noqa: BLE001 - instrumentation must not escape hooks - self._disabled = True + # Leave the pending records in place: drain() exports them inline, + # and a later schedule() retries starting a worker. self._worker = None - failed_pending = self._pending - self._pending = None - failed_event = self._flush_event - self._flush_event = None - self._flush_requested = False - if failed_event is not None: - failed_event.set() - return failed_pending, exc - return None, None + return exc + return None + + def _log_start_failure(self, start_error: Exception | None) -> None: + if start_error is None or self._start_failure_logged: + return + self._start_failure_logged = True + _logger.warning( + "workflow-insight: could not start export worker; records are " + "exported inline at invocation end instead: %s", + start_error, + ) + + def _pop_pending_locked(self) -> dict[str, Any] | None: + if not self._pending: + return None + key = next(iter(self._pending)) + return self._pending.pop(key) + + def _take_flush_locked(self) -> threading.Event | None: + flush_event = self._flush_event + self._flush_event = None + self._flush_requested = False + return flush_event def _run(self) -> None: while True: record: dict[str, Any] | None = None flush_event: threading.Event | None = None with self._condition: - while self._pending is None and not self._flush_requested: + while not self._pending and not self._flush_requested: self._condition.wait() - if self._pending is not None: - record = self._pending - self._pending = None - else: - flush_event = self._flush_event - self._flush_event = None - self._flush_requested = False + record = self._pop_pending_locked() + if record is None: + # Every pending record is exported: honor the flush request. + flush_event = self._take_flush_locked() if record is not None: self._export(record) @@ -123,10 +157,25 @@ def _run(self) -> None: if flush_event is not None: flush_event.set() with self._condition: - if self._pending is None and not self._flush_requested: + if not self._pending and not self._flush_requested: self._worker = None return + def _pump(self, flush_event: threading.Event) -> None: + """Export every pending record, then flush, on the calling thread.""" + while True: + with self._condition: + record = self._pop_pending_locked() + if record is None: + # Another inline drain may already have taken the request; + # flushing twice is harmless, losing a record is not. + self._take_flush_locked() + if record is None: + break + self._export(record) + self._flush() + flush_event.set() + def _export(self, record: dict[str, Any]) -> None: for exporter in self._exporters: try: @@ -159,4 +208,4 @@ def _worker_alive(self) -> bool: def _pending_count(self) -> int: with self._condition: - return int(self._pending is not None) + return len(self._pending) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index f19796a0..458b119a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -160,7 +160,7 @@ def _apply_result_override( class _ExecutionState: - __slots__ = ("start_time", "parsed_arn", "cached_input", "operations") + __slots__ = ("start_time", "parsed_arn", "cached_input", "operations", "closed") def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: self.start_time = start_time @@ -169,6 +169,10 @@ def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: # operation_id -> OperationInfo, adopted verbatim from the SDK's # authoritative snapshot (invocation start/end and operation-change). self.operations: dict[str, OperationInfo] = {} + # Set by on_invocation_end before it emits. An operation-change hook + # from a checkpoint that completes during the end-of-invocation drain + # is dropped, so no RUNNING snapshot can follow the final record. + self.closed = False class WorkflowInsightPlugin(DurableInstrumentationPlugin): @@ -230,6 +234,23 @@ def _discard_state(self, execution_arn: str) -> None: with self._lock: self._state.pop(execution_arn, None) + def _open_state(self, execution_arn: str) -> _ExecutionState | None: + """Return the execution's state only while its invocation is open. + + Unlike ``_ensure_state`` this never creates state: an operation-change + hook always follows an invocation start, so missing state means the + invocation already ended and its state was discarded. + """ + with self._lock: + state = self._state.get(execution_arn) + if state is None or state.closed: + return None + return state + + def _close_state(self, state: _ExecutionState) -> None: + with self._lock: + state.closed = True + def _adopt_operations( self, state: _ExecutionState, operations: dict[str, OperationInfo] ) -> None: @@ -270,7 +291,11 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: arn = info.execution_arn if not arn or not self._sampled_in(arn): return - state = self._ensure_state(arn) + state = self._open_state(arn) + if state is None: + # The invocation already ended (or is draining its final record); + # a snapshot from a checkpoint that completed late is stale. + return # Replace state with the full operations snapshot carried by the hook. self._adopt_operations(state, info.operations) # on-change mode exports an updated RUNNING record on each change so @@ -297,6 +322,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # Refresh from the fresh end-of-invocation snapshot before emitting so # the terminal record reflects the final operation map. self._adopt_operations(state, info.operations) + # Close before emitting: an operation-change hook arriving from a + # checkpoint that completes during the drain below is rejected, so no + # RUNNING snapshot can follow (or replace) the final record. + self._close_state(state) status = _STATUS_MAP.get(info.status, "RUNNING") is_terminal = status in ("SUCCEEDED", "FAILED") is_failure = status == "FAILED" diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 34879059..7e2bbbe8 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -13,8 +13,10 @@ ) -def _record(value: str) -> dict[str, Any]: - return {"status": "RUNNING", "value": value, "operations": []} +def _record( + value: str, *, arn: str = "exec-a", status: str = "RUNNING" +) -> dict[str, Any]: + return {"executionArn": arn, "status": status, "value": value, "operations": []} def _wait_until(predicate, timeout: float = 5.0) -> bool: @@ -105,17 +107,94 @@ def test_drain_flushes_after_export() -> None: assert capture.calls == [("export", "terminal"), ("flush", None)] -def test_worker_start_failure_never_escapes_hook(monkeypatch) -> None: +def test_worker_start_failure_exports_inline_on_drain(monkeypatch) -> None: def fail_start(self) -> None: # noqa: ARG001 raise RuntimeError("cannot start") monkeypatch.setattr(threading.Thread, "start", fail_start) - scheduler = _ExportScheduler([CaptureExporter()]) + capture = CaptureExporter() + scheduler = _ExportScheduler([capture]) - scheduler.schedule(_record("dropped")) + scheduler.schedule(_record("kept")) + assert scheduler._pending_count() == 1 scheduler.drain() + # Nothing is dropped and nothing escapes the hook: drain() exported the + # pending record on the calling thread and flushed. + assert capture.calls == [("export", "kept"), ("flush", None)] assert scheduler._pending_count() == 0 + assert not scheduler._worker_alive() + + +def test_pending_is_keyed_per_execution() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("a-first", arn="exec-a")) + assert exporter.started.wait(5.0) + + # While the worker is blocked: two more snapshots for A (coalesce to the + # latest) and one for B (kept in its own slot, never displaced by A). + scheduler.schedule(_record("a-middle", arn="exec-a")) + scheduler.schedule(_record("b-only", arn="exec-b")) + scheduler.schedule(_record("a-latest", arn="exec-a")) + assert scheduler._pending_count() == 2 + + exporter.release.set() + scheduler.drain() + assert exporter.calls == [ + ("export", "a-first"), + ("export", "a-latest"), + ("export", "b-only"), + ("flush", None), + ] + + +def test_running_never_supersedes_pending_terminal_of_same_execution() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("x-inflight", arn="exec-x")) + assert exporter.started.wait(5.0) + + scheduler.schedule(_record("a-terminal", arn="exec-a", status="SUCCEEDED")) + # A late RUNNING snapshot for A (an operation-change hook from a checkpoint + # completing during the drain) must not replace A's terminal record. + scheduler.schedule(_record("a-late-running", arn="exec-a")) + assert scheduler._pending_count() == 1 + + exporter.release.set() + scheduler.drain() + assert exporter.calls == [ + ("export", "x-inflight"), + ("export", "a-terminal"), + ("flush", None), + ] + + +def test_concurrent_executions_each_deliver_their_terminal_record() -> None: + capture = CaptureExporter() + scheduler = _ExportScheduler([capture]) + rounds = 50 + barrier = threading.Barrier(2) + + def drive(arn: str) -> None: + for i in range(rounds): + barrier.wait() + scheduler.schedule(_record(f"{arn}-running-{i}", arn=arn)) + scheduler.schedule( + _record(f"{arn}-terminal-{i}", arn=arn, status="SUCCEEDED") + ) + scheduler.drain() + + threads = [threading.Thread(target=drive, args=(arn,)) for arn in ("a", "b")] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30.0) + assert not any(thread.is_alive() for thread in threads) + + exported = {value for kind, value in capture.calls if kind == "export"} + expected = {f"{arn}-terminal-{i}" for arn in ("a", "b") for i in range(rounds)} + assert expected <= exported, sorted(expected - exported) def test_superseded_record_finalizes_after_lane_unlock() -> None: diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 9d7638eb..bd2d0512 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -13,8 +13,12 @@ from __future__ import annotations import datetime +import threading +import time from typing import Any +import pytest + from aws_durable_execution_sdk_python.lambda_service import ( ErrorObject, OperationStatus, @@ -402,6 +406,132 @@ def test_on_change_schedules_running_and_delivers_terminal(): assert len(ids) == len(set(ids)) +# -- several executions in flight on one plugin ------------------------------ + + +class BlockingExporter(CaptureExporter): + """Blocks the first export until released, so tests can observe the + scheduler while a record is in flight.""" + + def __init__(self) -> None: + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + + def export(self, record: dict[str, Any]) -> None: + if not self.started.is_set(): + self.started.set() + self.release.wait(5.0) + super().export(record) + + +def _drive_to_success(plugin, arn: str, barrier: threading.Barrier) -> None: + op = _step("s", op_id="1") + barrier.wait() + plugin.on_invocation_start(_start(arn, operations={})) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=arn, updated_operations=_ops(op), operations=_ops(op) + ) + ) + plugin.on_invocation_end(_end(arn, operations=_ops(op))) + + +@pytest.mark.parametrize("emit_mode", ["on-complete", "on-change"]) +def test_concurrent_executions_on_one_plugin_each_deliver_terminal_record( + emit_mode, +): + # The local runner drives independent executions concurrently through one + # shared plugin. Every execution's terminal record must reach the exporter; + # one execution scheduling a snapshot must never displace another's. + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode=emit_mode) + ) + rounds = 50 + for _ in range(rounds): + barrier = threading.Barrier(2) + threads = [ + threading.Thread(target=_drive_to_success, args=(plugin, arn, barrier)) + for arn in (ARN, ARN_B) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10.0) + assert not any(thread.is_alive() for thread in threads) + + terminal = [r for r in exporter.records if r["status"] == "SUCCEEDED"] + per_arn = {arn: 0 for arn in (ARN, ARN_B)} + for record in terminal: + per_arn[record["executionArn"]] += 1 + assert per_arn == {ARN: rounds, ARN_B: rounds} + assert plugin._state == {} + + +def test_late_operation_change_during_drain_cannot_follow_terminal_record(): + # on-change mode, with the exporter blocked on the RUNNING record emitted + # at invocation start. on_invocation_end schedules SUCCEEDED and blocks in + # drain(); an operation-change hook that arrives meanwhile (a checkpoint + # completing late) must be dropped, not exported as a trailing RUNNING. + exporter = BlockingExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + assert exporter.started.wait(5.0) + + ender = threading.Thread( + target=plugin.on_invocation_end, args=(_end(operations=_ops(op)),) + ) + ender.start() + deadline = time.monotonic() + 5.0 + while plugin._scheduler._pending_count() == 0 and time.monotonic() < deadline: + time.sleep(0.005) + assert plugin._scheduler._pending_count() == 1 + + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + + exporter.release.set() + ender.join(5.0) + assert not ender.is_alive() + + assert [r["status"] for r in exporter.records] == ["RUNNING", "SUCCEEDED"] + assert plugin._state == {} + + +def test_operation_change_after_invocation_end_is_ignored(): + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin.on_invocation_start(_start(operations={})) + # Deliver the start RUNNING snapshot before the end hook schedules SUCCEEDED, + # otherwise the terminal snapshot may legitimately supersede it in the + # execution's pending slot and the record list would depend on timing. + plugin._scheduler.drain() + plugin.on_invocation_end(_end(operations=_ops(op))) + assert plugin._state == {} + + # A hook for an execution whose invocation already ended must not recreate + # state or emit anything. + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + plugin._scheduler.drain() + + assert [r["status"] for r in exporter.records] == ["RUNNING", "SUCCEEDED"] + assert plugin._state == {} + + # -- no cross-execution contamination (comment 3) ----------------------------