From 14fca8bbbd2bc2c7d30bf115317682cce3c21711 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 15 Sep 2026 02:05:50 +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 15-18% 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 so far was delivered. Memory stays bounded by executions in flight, the bound the plugin's _state already has. Coalescing within one execution is unchanged and now pinned by tests instead of hidden: a drain-between-hooks test asserts every on-change record arrives when nothing coalesces, and a blocked-exporter test asserts a burst collapses to first + latest. The README states the literal coalescing condition. Addresses issues 1 and 2 raised in the #719 review. --- .../README.md | 12 ++- .../_export_scheduler.py | 45 +++++--- .../tests/test_export_scheduler.py | 100 +++++++++++++++++- .../tests/test_plugin.py | 61 ++++++++++- 4 files changed, 191 insertions(+), 27 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..06a83acf 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 that execution's latest pending snapshot and wake the worker. +> Consecutive `on-change` snapshots for one execution coalesce whenever the +> worker has not yet claimed the previous one, including when no export is +> running, so under fast in-memory checkpoints (the local runner) a burst may +> deliver only its first and latest snapshots. Snapshots from different +> executions never coalesce with each other. 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. ## 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..07488781 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 @@ -17,27 +17,38 @@ 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. + + The pending map is keyed by ``executionArn`` so a snapshot only ever + replaces its own execution's pending snapshot. Several executions may be in + flight on one plugin instance (the local runner drives independent + executions concurrently through a shared ``workflow_insight(...)``), and a + single shared slot let execution B's ``RUNNING`` snapshot displace + execution A's terminal snapshot before the worker claimed it. + """ 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, in first-arrival order. + 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 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.""" displaced: dict[str, Any] | None = None - failed_pending: dict[str, Any] | None = None + failed_pending: dict[str, dict[str, Any]] | None = None start_error: Exception | None = None + arn = str(record.get("executionArn", "")) with self._condition: if self._disabled: return - displaced = self._pending - self._pending = record + displaced = self._pending.get(arn) + # Assigning an existing key keeps its first-arrival position. + self._pending[arn] = record failed_pending, start_error = self._ensure_worker_locked() self._condition.notify() # Releasing either record may run custom finalizers, so do it unlocked. @@ -50,8 +61,8 @@ def schedule(self, record: dict[str, Any]) -> None: ) 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.""" + failed_pending: dict[str, dict[str, Any]] | None = None start_error: Exception | None = None with self._condition: if self._disabled: @@ -76,7 +87,7 @@ def drain(self) -> None: def _ensure_worker_locked( self, - ) -> tuple[dict[str, Any] | None, Exception | None]: + ) -> tuple[dict[str, dict[str, Any]] | None, Exception | None]: if self._worker is not None and self._worker.is_alive(): return None, None worker = threading.Thread( @@ -91,7 +102,7 @@ def _ensure_worker_locked( self._disabled = True self._worker = None failed_pending = self._pending - self._pending = None + self._pending = {} failed_event = self._flush_event self._flush_event = None self._flush_requested = False @@ -105,11 +116,13 @@ def _run(self) -> None: 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 + if self._pending: + # Export every pending execution before honoring a flush, so + # drain() means everything scheduled so far was delivered. + arn = next(iter(self._pending)) + record = self._pending.pop(arn) else: flush_event = self._flush_event self._flush_event = None @@ -123,7 +136,7 @@ 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 @@ -159,4 +172,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/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 34879059..768d9b78 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,13 @@ ) -def _record(value: str) -> dict[str, Any]: - return {"status": "RUNNING", "value": value, "operations": []} +def _record(value: str, arn: str = "exec-a") -> dict[str, Any]: + return { + "executionArn": arn, + "status": "RUNNING", + "value": value, + "operations": [], + } def _wait_until(predicate, timeout: float = 5.0) -> bool: @@ -157,3 +162,94 @@ def test_drain_waits_for_blocked_exporter() -> None: assert not drain_thread.is_alive() assert exporter.calls == [("export", "terminal"), ("flush", None)] + + +# -- pending map is keyed per execution (#719 review, issue 1) ---------------- + + +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) + + # B's snapshot must not displace A's pending terminal snapshot. + scheduler.schedule(_record("a-terminal", arn="exec-a")) + scheduler.schedule(_record("b-running", arn="exec-b")) + assert scheduler._pending_count() == 2 + + exporter.release.set() + scheduler.drain() + assert exporter.calls == [ + ("export", "a-first"), + ("export", "a-terminal"), + ("export", "b-running"), + ("flush", None), + ] + + +def test_drain_delivers_every_pending_execution_before_flush() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("a", arn="exec-a")) + assert exporter.started.wait(5.0) + scheduler.schedule(_record("b", arn="exec-b")) + scheduler.schedule(_record("c", arn="exec-c")) + + exporter.release.set() + scheduler.drain() + + exported = [value for kind, value in exporter.calls if kind == "export"] + assert exported == ["a", "b", "c"] + assert exporter.calls[-1] == ("flush", None) + + +def test_concurrent_executions_each_deliver_terminal_record() -> None: + rounds = 50 + exporter = CaptureExporter() + scheduler = _ExportScheduler([exporter]) + + def drive(arn: str) -> None: + for i in range(rounds): + scheduler.schedule(_record(f"{arn}-running-{i}", arn=arn)) + scheduler.schedule(_record(f"{arn}-terminal-{i}", arn=arn)) + 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 exporter.calls + if kind == "export" and value is not None + ] + for arn in ("a", "b"): + terminals = [v for v in exported if v.startswith(f"{arn}-terminal-")] + assert terminals == [f"{arn}-terminal-{i}" for i in range(rounds)] + + +# -- coalescing contract (#719 review, issue 2) ------------------------------- + + +def test_same_execution_coalesces_to_first_and_latest_while_blocked() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("first")) + assert exporter.started.wait(5.0) + + scheduler.schedule(_record("second")) + scheduler.schedule(_record("third")) + scheduler.schedule(_record("latest")) + assert scheduler._pending_count() == 1 + + exporter.release.set() + scheduler.drain() + assert exporter.calls == [ + ("export", "first"), + ("export", "latest"), + ("flush", 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..99371e7b 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,10 @@ from __future__ import annotations import datetime +import threading from typing import Any +import pytest from aws_durable_execution_sdk_python.lambda_service import ( ErrorObject, OperationStatus, @@ -368,10 +370,14 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): assert rec["durationMs"] is not None and rec["durationMs"] >= 0 -# -- on-change schedules progress and delivers terminal state ---------------- +# -- on-change delivers every record when nothing coalesces ------------------ -def test_on_change_schedules_running_and_delivers_terminal(): +def test_on_change_delivers_every_record_when_drained_between_hooks(): + # Draining after each hook removes the worker race, so this pins the + # "no back-pressure means no loss" contract: one RUNNING per start/change + # plus the terminal record, in order. (The coalescing contract itself is + # pinned in test_export_scheduler.py with a blocked exporter.) exporter = CaptureExporter() plugin = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") @@ -380,28 +386,73 @@ def test_on_change_schedules_running_and_delivers_terminal(): op2 = _step("s2", op_id="2") plugin.on_invocation_start(_start(operations={})) + plugin._scheduler.drain() plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1) ) ) + plugin._scheduler.drain() plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op2), operations=_ops(op1, op2) ) ) + plugin._scheduler.drain() plugin.on_invocation_end(_end(operations=_ops(op1, op2))) statuses = [record["status"] for record in exporter.records] - assert statuses - 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"] ids = [op["id"] for op in final["operations"]] assert len(ids) == len(set(ids)) +# -- concurrent executions on one plugin keep their terminal records --------- + + +@pytest.mark.parametrize("emit_mode", ["on-complete", "on-change"]) +def test_concurrent_executions_on_one_plugin_each_deliver_terminal_record( + emit_mode, +): + # Two executions in flight on one plugin (as the local runner does). Each + # invocation end must deliver its own terminal record; a shared pending + # slot let one execution's RUNNING snapshot displace the other's terminal. + rounds = 50 + exporter = CaptureExporter() + plugin = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode=emit_mode) + ) + + def drive(arn: str) -> None: + op = _step(f"{arn}-step", op_id=f"{arn}-1") + for _ in range(rounds): + 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))) + + threads = [threading.Thread(target=drive, args=(arn,)) for arn in (ARN, ARN_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) + + for arn in (ARN, ARN_B): + terminal = [ + r + for r in exporter.records + if r["executionArn"] == arn and r["status"] == "SUCCEEDED" + ] + assert len(terminal) == rounds + assert plugin._state == {} + + # -- no cross-execution contamination (comment 3) ----------------------------