fix(insight): buffer on-change export bursts

Author: Alex Wang

packages/aws-durable-execution-sdk-python-insight/README.md

@@ -62,13 +62,23 @@ Behavior is validated cross-SDK by the `insight` conformance suite
 > 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. 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
+> 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
 

packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
-"""Asynchronous, coalescing export scheduler for the Workflow Insight plugin.
+"""Asynchronous, bounded export scheduler for the Workflow Insight plugin.
 
 The plugin builds one canonical ``WorkflowInsight`` record on the SDK checkpoint
 thread and hands it to :class:`_ExportScheduler`. The scheduler keeps all
@@ -16,12 +16,12 @@
 * One lazy daemon worker per exporter lane; never more than one live worker per
   lane, and a blocked worker is retained -- never replaced -- so threads cannot
   grow without bound.
-* Per lane, at most one in-flight record and one latest *pending* record per
-  execution ARN. Records are cumulative snapshots, so a newer pending record for
-  an ARN replaces the older one (coalescing); an in-flight record is never
-  cancelled. Updating a pending ARN moves it to the back of the queue for
-  fairness across ARNs. Pending ARNs are capped; the oldest is evicted when the
-  cap is exceeded (only reachable behind a blocked/slow exporter).
+* Per lane, each execution ARN has a small bounded FIFO of pending snapshots.
+  Ordinary bursts remain observable even when Python does not schedule the
+  worker between checkpoint hooks. Once that FIFO is full, its oldest pending
+  snapshot is dropped so the newest progress and terminal snapshots are
+  retained. ARNs are served round-robin for fairness, and pending ARNs are also
+  capped; the oldest execution is evicted when that cap is exceeded.
 * Invocation end enqueues one flush barrier per touched lane after the latest
   record and waits for all barriers under a single shared timeout deadline. On
   timeout the workflow response is returned, degradation is logged, and each
@@ -54,6 +54,17 @@
 # bounded regardless of how long a worker stays blocked.
 _DEFAULT_MAX_PENDING_EXECUTIONS = 1024
 
+# 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"
 _FLUSH = "flush"
@@ -95,9 +106,15 @@ def __init__(
         exporter: InsightExporter,
         *,
         max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS,
+        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_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
@@ -107,9 +124,10 @@ 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 (coalesced). Insertion order is the
-        # fairness order; updating an arn moves it to the back.
-        self._pending: OrderedDict[str, dict[str, Any]] = 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
 
@@ -119,15 +137,24 @@ 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:
-                # Coalesce: replace the pending record and move it to the back so
-                # a busy execution cannot starve the others.
-                self._pending[execution_arn] = record
+                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
+                self._pending[execution_arn] = deque([record])
                 self._queue.append((_RECORD, execution_arn))
-                self._enforce_pending_cap()
+                self._enforce_pending_execution_cap()
+            self._enforce_pending_record_cap()
             self._ensure_worker_locked()
             self._cond.notify()
 
@@ -176,22 +203,54 @@ 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_cap(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._pending.popitem(last=False)
             self._remove_record_token(old_arn)
             _logger.warning(
                 "workflow-insight: export lane for %s is full "
-                "(cap=%d); dropping pending record for %s",
+                "(cap=%d); dropping pending records for %s",
                 type(self._exporter).__name__,
                 self._max_pending,
                 old_arn,
             )
 
+    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()
+            _logger.warning(
+                "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_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:
         for index, (kind, payload) in enumerate(self._queue):
             if kind == _RECORD and payload == execution_arn:
@@ -228,9 +287,17 @@ def _run_worker(self) -> None:
                 kind, payload = self._queue.popleft()
                 record: dict[str, Any] | None = None
                 if kind == _RECORD:
-                    record = self._pending.pop(payload, None)
-                    if record is None:
+                    pending = self._pending.get(payload)
+                    if not pending:
                         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]
 
             if kind == _RECORD and record is not None:
                 self._export_one(record)
@@ -299,6 +366,10 @@ def _pending_count(self) -> int:
         with self._cond:
             return len(self._pending)
 
+    def _pending_record_count(self) -> int:
+        with self._cond:
+            return sum(len(records) for records in self._pending.values())
+
     def _queue_len(self) -> int:
         with self._cond:
             return len(self._queue)
@@ -316,9 +387,18 @@ def __init__(
         exporters: list[InsightExporter],
         *,
         max_pending_executions: int = _DEFAULT_MAX_PENDING_EXECUTIONS,
+        max_pending_records: int = _DEFAULT_MAX_PENDING_RECORDS,
+        max_pending_records_per_execution: int = (
+            _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION
+        ),
     ) -> None:
         self._lanes = [
-            _ExporterLane(exporter, max_pending_executions=max_pending_executions)
+            _ExporterLane(
+                exporter,
+                max_pending_executions=max_pending_executions,
+                max_pending_records=max_pending_records,
+                max_pending_records_per_execution=max_pending_records_per_execution,
+            )
             for exporter in exporters
         ]