fix(insight): buffer on-change export bursts

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

@@ -61,16 +61,22 @@ Behavior is validated cross-SDK by the `insight` conformance suite
 > workflow progress. Because each configured exporter is driven by its own
 > single background worker, each exporter object may belong to only one live
 > `WorkflowInsightPlugin`: listing it twice or sharing it across plugin instances
-> raises `ValueError`. Separate instances of the same exporter class (e.g. two
-> `S3Exporter`s for different buckets) are fine. A blocked lane retains at most
-> 1,024 pending executions and 16 MB of estimated retained memory; it drops the
-> oldest pending snapshot when either bound is reached. 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
+> raises `ValueError`. Separate instances of the same exporter class are fine.
+> Each lane keeps up to 16 pending snapshots per execution, 1,024 records total,
+> and 16 MB of estimated retained memory. When a bound fills, it drops the oldest
+> pending snapshot so recent 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

@@ -9,10 +9,9 @@
 
 Each exporter lane:
 
-* Keeps the latest pending snapshot per execution ARN and processes ARNs
-  round-robin.
-* Drops the oldest pending snapshot when the execution-count or byte budget is
-  reached, keeping memory bounded.
+* Keeps a bounded FIFO per execution ARN and processes ARNs round-robin.
+* Drops the oldest pending snapshots when the per-execution, lane-wide record,
+  or byte limit is reached, favoring recent progress and terminal snapshots.
 * Uses a lane-wide flush barrier at invocation end. All barriers share one
   timeout; timed-out barriers are removed without replacing a blocked worker.
 * Stops its worker after an invocation has drained and the lane becomes idle.
@@ -30,7 +29,8 @@
 import time
 import types
 import uuid
-from collections import OrderedDict, deque
+from collections import deque
+from dataclasses import dataclass
 from typing import Any
 
 from aws_durable_execution_sdk_python_insight.truncation import truncate_record
@@ -45,9 +45,19 @@
 # bounded regardless of how long a worker stays blocked.
 _DEFAULT_MAX_PENDING_EXECUTIONS = 1024
 
-# Estimated Python object memory retained by one blocked lane. Lambda functions
-# can be configured with 128 MiB, so keep the instrumentation backlog well below
-# that floor. This is a conservative budget signal, not an exact heap measurement.
+# 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
+
+# Estimated Python object memory retained by one blocked lane. This keeps the
+# instrumentation backlog well below Lambda's 128 MiB memory floor.
 _DEFAULT_MAX_PENDING_BYTES = 16_000_000
 
 
@@ -242,6 +252,16 @@ def _copy_record_containers(record: dict[str, Any]) -> dict[str, Any]:
 _RECORD = "record"
 _FLUSH = "flush"
 
+_RecordToken = tuple[str, int]
+
+
+@dataclass(slots=True)
+class _PendingRecord:
+    sequence: int
+    generation: int
+    value: dict[str, Any]
+    size: int
+
 
 class _FlushBarrier:
     """A one-shot flush marker the invocation-end thread waits on.
@@ -280,10 +300,16 @@ 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
+        ),
         max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES,
     ) -> 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)
         self._max_pending_bytes = max(1, max_pending_bytes)
         self._pending_bytes = 0
         self._inflight_bytes = 0
@@ -294,12 +320,15 @@ def __init__(
         # makes any accidental recursive acquisition fail loudly instead of
         # silently succeeding.
         self._cond = threading.Condition(threading.Lock())
-        # Ordered work list: entries are (_RECORD, arn) or (_FLUSH, barrier).
+        # Ordered work list: entries are (_RECORD, (arn, generation)) or
+        # (_FLUSH, barrier). A generation changes whenever a flush is queued, so
+        # one ARN can have independent tokens on both sides of a barrier.
         self._queue: deque[tuple[str, Any]] = deque()
-        # arn -> (latest pending record, retained-memory estimate). Insertion
-        # order is both record age and fairness order because replacing an ARN
-        # moves it to the back.
-        self._pending: OrderedDict[str, tuple[dict[str, Any], int]] = OrderedDict()
+        # arn -> FIFO ordered by record sequence. Queue-token order, not this
+        # mapping, controls round-robin fairness.
+        self._pending: dict[str, deque[_PendingRecord]] = {}
+        self._generation = 0
+        self._next_sequence = 0
         self._stop_when_idle = False
         self._worker: threading.Thread | None = None
         self._disabled = False
@@ -318,37 +347,45 @@ def schedule(
             self._stop_when_idle = False
             size = max(0, record_size)
             if size > self._max_pending_bytes:
-                superseded = self._pending.pop(execution_arn, None)
-                if superseded is not None:
-                    _, superseded_size = superseded
-                    self._pending_bytes -= superseded_size
-                    self._remove_record_token(execution_arn)
                 _logger.warning(
                     "workflow-insight: pending record for %s on %s exceeds the "
-                    "byte budget (%d > %d); dropping this record%s",
+                    "byte budget (%d > %d); dropping this record",
                     execution_arn,
                     type(self._exporter).__name__,
                     size,
                     self._max_pending_bytes,
-                    " and its superseded pending snapshot"
-                    if superseded is not None
-                    else "",
                 )
                 return
-            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.
-                _, old_size = self._pending[execution_arn]
-                self._pending_bytes -= old_size
-                self._pending[execution_arn] = (record, size)
-                self._pending_bytes += size
-                self._pending.move_to_end(execution_arn)
-                self._move_record_token_to_back(execution_arn)
+            pending = self._pending.get(execution_arn)
+            if pending is None:
+                pending = deque()
+                self._pending[execution_arn] = pending
+            elif len(pending) >= self._max_pending_per_execution:
+                self._drop_oldest_pending_record(execution_arn)
+                pending = self._pending.setdefault(execution_arn, deque())
+                _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,
+                )
+
+            generation = self._generation
+            has_generation = bool(pending and pending[-1].generation == generation)
+            pending.append(
+                _PendingRecord(self._next_sequence, generation, record, size)
+            )
+            self._pending_bytes += size
+            self._next_sequence += 1
+            token = (execution_arn, generation)
+            if has_generation:
+                self._move_record_token_to_back(token)
             else:
-                self._pending[execution_arn] = (record, size)
-                self._pending_bytes += size
-                self._queue.append((_RECORD, execution_arn))
-            self._enforce_pending_caps()
+                self._queue.append((_RECORD, token))
+                self._enforce_pending_execution_cap()
+            self._enforce_pending_record_cap()
+            self._enforce_pending_byte_cap()
             self._ensure_worker_locked()
             self._cond.notify()
 
@@ -361,6 +398,7 @@ def enqueue_flush(self) -> _FlushBarrier:
                 barrier.complete()
                 return barrier
             self._queue.append((_FLUSH, barrier))
+            self._generation += 1
             self._ensure_worker_locked()
             self._cond.notify()
         return barrier
@@ -395,31 +433,66 @@ def cancel_flush(self, barrier: _FlushBarrier) -> None:
 
     # -- queue bookkeeping (must hold ``_cond``) ------------------------------
 
-    def _move_record_token_to_back(self, execution_arn: str) -> None:
+    def _move_record_token_to_back(self, token: _RecordToken) -> None:
         for index, (kind, payload) in enumerate(self._queue):
-            if kind == _RECORD and payload == execution_arn:
+            if kind == _RECORD and payload == token:
                 del self._queue[index]
-                self._queue.append((_RECORD, execution_arn))
+                self._queue.append((_RECORD, token))
                 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.
 
-    def _enforce_pending_caps(self) -> None:
+    def _requeue_record_before_flush(self, token: _RecordToken) -> None:
+        """Requeue behind peers without crossing this ARN's next generation."""
+        execution_arn, generation = token
+        for index, (kind, payload) in enumerate(self._queue):
+            later_same_arn = (
+                kind == _RECORD
+                and payload[0] == execution_arn
+                and payload[1] > generation
+            )
+            if kind == _FLUSH or later_same_arn:
+                self._queue.insert(index, (_RECORD, token))
+                return
+        self._queue.append((_RECORD, token))
+
+    def _oldest_pending_arn(self) -> str:
+        return min(
+            self._pending,
+            key=lambda arn: self._pending[arn][0].sequence,
+        )
+
+    def _enforce_pending_execution_cap(self) -> None:
         while len(self._pending) > self._max_pending:
-            old_arn, _ = self._drop_oldest_pending()
+            old_arn = self._oldest_pending_arn()
+            self._drop_pending_execution(old_arn)
             _logger.warning(
-                "workflow-insight: export lane for %s reached its execution cap "
-                "(%d); dropping pending record for %s",
+                "workflow-insight: export lane for %s is full "
+                "(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 = self._oldest_pending_arn()
+            self._drop_oldest_pending_record(old_arn)
+            _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,
+            )
+
+    def _enforce_pending_byte_cap(self) -> None:
         while (
             self._pending_bytes + self._inflight_bytes > self._max_pending_bytes
             and self._pending
         ):
-            old_arn, dropped_size = self._drop_oldest_pending()
+            old_arn = self._oldest_pending_arn()
+            dropped_size = self._drop_oldest_pending_record(old_arn)
             _logger.warning(
                 "workflow-insight: export lane for %s reached its pending byte "
                 "budget (%d); dropping %d-byte pending record for %s",
@@ -429,15 +502,28 @@ def _enforce_pending_caps(self) -> None:
                 old_arn,
             )
 
-    def _drop_oldest_pending(self) -> tuple[str, int]:
-        old_arn, (_, old_size) = self._pending.popitem(last=False)
-        self._pending_bytes -= old_size
-        self._remove_record_token(old_arn)
-        return old_arn, old_size
+    def _drop_oldest_pending_record(self, execution_arn: str) -> int:
+        records = self._pending[execution_arn]
+        dropped = records.popleft()
+        self._pending_bytes -= dropped.size
+        token = (execution_arn, dropped.generation)
+        if not records or records[0].generation != dropped.generation:
+            self._remove_record_token(token)
+        if not records:
+            del self._pending[execution_arn]
+        return dropped.size
+
+    def _drop_pending_execution(self, execution_arn: str) -> None:
+        records = self._pending.pop(execution_arn)
+        self._pending_bytes -= sum(record.size for record in records)
+        for index in range(len(self._queue) - 1, -1, -1):
+            kind, payload = self._queue[index]
+            if kind == _RECORD and payload[0] == execution_arn:
+                del self._queue[index]
 
-    def _remove_record_token(self, execution_arn: str) -> None:
+    def _remove_record_token(self, token: _RecordToken) -> None:
         for index, (kind, payload) in enumerate(self._queue):
-            if kind == _RECORD and payload == execution_arn:
+            if kind == _RECORD and payload == token:
                 del self._queue[index]
                 return
 
@@ -497,12 +583,21 @@ def _run_worker(self) -> None:
                 record: dict[str, Any] | None = None
                 record_size = 0
                 if kind == _RECORD:
-                    pending = self._pending.pop(payload, None)
-                    if pending is None:
+                    token: _RecordToken = payload
+                    execution_arn, generation = token
+                    pending = self._pending.get(execution_arn)
+                    if not pending or pending[0].generation != generation:
                         continue
-                    record, record_size = pending
+                    pending_record = pending.popleft()
+                    record_size = pending_record.size
                     self._pending_bytes -= record_size
                     self._inflight_bytes += record_size
+                    record = pending_record.value
+                    if pending and pending[0].generation == generation:
+                        # One record per ARN turn within this barrier generation.
+                        self._requeue_record_before_flush(token)
+                    elif not pending:
+                        del self._pending[execution_arn]
 
             if kind == _RECORD and record is not None:
                 try:
@@ -572,6 +667,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 _pending_bytes_count(self) -> int:
         with self._cond:
             return self._pending_bytes
@@ -597,13 +696,19 @@ 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
+        ),
         max_pending_bytes: int = _DEFAULT_MAX_PENDING_BYTES,
     ) -> None:
         self._max_pending_bytes = max(1, max_pending_bytes)
         self._lanes = [
             _ExporterLane(
                 exporter,
                 max_pending_executions=max_pending_executions,
+                max_pending_records=max_pending_records,
+                max_pending_records_per_execution=max_pending_records_per_execution,
                 max_pending_bytes=self._max_pending_bytes,
             )
             for exporter in exporters