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