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) ----------------------------