fix(insight): preserve timed-out FIFO exports

Author: Alex Wang

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

@@ -56,7 +56,7 @@
 _DEFAULT_MAX_PENDING_BYTES = 16_000_000
 
 
-def _estimate_retained_size(value: Any) -> int:
+def _estimate_retained_size(value: Any, max_size: int | None = None) -> int:
     """Estimate retained Python memory without serializing or calling render()."""
     total = 0
     seen: set[int] = set()
@@ -71,8 +71,12 @@ def _estimate_retained_size(value: Any) -> int:
             total += sys.getsizeof(item)
         except Exception:  # noqa: BLE001 - estimation must never break a hook
             total += 1_024
+        if max_size is not None and total > max_size:
+            return max_size + 1
         try:
-            if isinstance(item, dict):
+            if isinstance(item, memoryview):
+                stack.append(item.obj)
+            elif isinstance(item, dict):
                 stack.extend(item.keys())
                 stack.extend(item.values())
             elif isinstance(item, (list, tuple, set, frozenset, deque)):
@@ -285,26 +289,19 @@ def request_stop_when_idle(self) -> None:
             self._cond.notify()
 
     def cancel_flush(self, barrier: _FlushBarrier) -> None:
-        """Cancel a timed-out flush barrier so it cannot pile up behind a
-        blocked worker.
-
-        Under the lane lock: mark the barrier cancelled and, if its ``_FLUSH``
-        marker is still queued, remove that exact marker and complete the
-        barrier here. Removing it is what keeps queue/barrier state bounded
-        across many warm invocations behind a blocked exporter -- otherwise one
-        stale barrier per invocation would accumulate behind the stuck worker.
-
-        If the worker has already popped the marker (the flush is in flight or
-        about to run) the marker is no longer in the queue: we only set
-        ``canceled`` and leave completion to the worker, which skips the
-        now-pointless flush and completes the barrier itself. A synchronous
-        in-flight ``flush()`` is never interrupted.
-        """
+        """Stop waiting for a timed-out barrier while retaining one later flush."""
         with self._cond:
             barrier.canceled = True
+            # Keep at most one detached flush. Moving it to this barrier's
+            # position makes it cover all work scheduled before the latest
+            # timeout without accumulating one marker per warm invocation.
+            for index in range(len(self._queue) - 1, -1, -1):
+                kind, payload = self._queue[index]
+                if kind == _FLUSH and payload is None:
+                    del self._queue[index]
             for index, (kind, payload) in enumerate(self._queue):
                 if kind == _FLUSH and payload is barrier:
-                    del self._queue[index]
+                    self._queue[index] = (_FLUSH, None)
                     barrier.complete()
                     return
 
@@ -407,7 +404,7 @@ def _disable_locked(self, exc: Exception) -> None:
         self._pending.clear()
         self._pending_bytes = 0
         for kind, payload in self._queue:
-            if kind == _FLUSH:
+            if kind == _FLUSH and payload is not None:
                 barrier: _FlushBarrier = payload
                 barrier.canceled = True
                 barrier.failed = True
@@ -472,10 +469,10 @@ def _run_worker(self) -> None:
             if kind == _RECORD and record is not None:
                 self._export_one(record)
             else:  # _FLUSH
-                barrier: _FlushBarrier = payload
-                if not barrier.canceled:
-                    self._flush()
-                barrier.complete()
+                barrier: _FlushBarrier | None = payload
+                self._flush()
+                if barrier is not None:
+                    barrier.complete()
 
     def _export_one(self, record: dict[str, Any]) -> None:
         exporter = self._exporter
@@ -564,20 +561,21 @@ def __init__(
         ),
         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=max_pending_bytes,
+                max_pending_bytes=self._max_pending_bytes,
             )
             for exporter in exporters
         ]
 
     def schedule(self, execution_arn: str, record: dict[str, Any]) -> None:
         """Fan a canonical record out to every lane. Returns immediately."""
-        record_size = _estimate_retained_size(record)
+        record_size = _estimate_retained_size(record, self._max_pending_bytes)
         for lane in self._lanes:
             lane.schedule(execution_arn, record, record_size)
 

packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py

@@ -151,6 +151,35 @@ def release(self) -> None:
         self._release.set()
 
 
+class BlockingBufferedExporter:
+    """Blocks export and publishes buffered records only when flush runs."""
+
+    max_record_size_bytes = None
+
+    def __init__(self) -> None:
+        self.started = threading.Event()
+        self._release = threading.Event()
+        self.buffered: list[Any] = []
+        self.published: list[Any] = []
+        self.flushed = 0
+
+    def render(self, record: dict[str, Any]) -> Any:
+        return record
+
+    def export(self, record: dict[str, Any]) -> None:
+        self.started.set()
+        self._release.wait(5.0)
+        self.buffered.append(record.get("v"))
+
+    def flush(self) -> None:
+        self.flushed += 1
+        self.published.extend(self.buffered)
+        self.buffered.clear()
+
+    def release(self) -> None:
+        self._release.set()
+
+
 class FailingExporter:
     """Raises in both export and flush."""
 
@@ -200,6 +229,16 @@ def __sizeof__(self) -> int:
         raise RuntimeError("size unavailable")
 
 
+class _TrackedLargeList(list[Any]):
+    def __init__(self) -> None:
+        super().__init__([None] * 10_000)
+        self.iterated = False
+
+    def __iter__(self):
+        self.iterated = True
+        return super().__iter__()
+
+
 # -- lazy worker creation / one worker per exporter --------------------------
 
 
@@ -731,6 +770,42 @@ def test_retained_size_traverses_slots_after_shallow_size_failure():
         assert exporter.exported_values() == ["inflight"]
 
 
+def test_retained_size_saturates_before_traversing_large_shallow_container():
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500)
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
+    assert _wait_until(exporter.started.is_set)
+
+    payload = _TrackedLargeList()
+    record = _rec(ARN_B, "large-shallow")
+    record["payload"] = payload
+    scheduler.schedule(ARN_B, record)
+
+    assert payload.iterated is False
+    assert lane._pending_count() == 0
+    assert lane._pending_bytes_count() == 0
+    exporter.release()
+    scheduler.end_invocation(5.0)
+
+
+def test_retained_size_counts_memoryview_backing_buffer():
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=2_500)
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
+    assert _wait_until(exporter.started.is_set)
+
+    record = _rec(ARN_B, "memoryview")
+    record["payload"] = memoryview(bytearray(4_000))
+    scheduler.schedule(ARN_B, record)
+
+    assert lane._pending_count() == 0
+    assert lane._pending_bytes_count() == 0
+    exporter.release()
+    scheduler.end_invocation(5.0)
+
+
 def test_record_sizing_exception_does_not_escape_schedule():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
@@ -809,97 +884,93 @@ def test_cancelled_barrier_preserves_generation_order_and_terminal():
     assert lane._queue_len() == 0
 
 
-def test_cancelled_barrier_is_cleaned_up_and_worker_exits():
+def test_timed_out_barrier_flushes_eventually_and_worker_exits():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
     lane = scheduler._lanes[0]
     scheduler.schedule(ARN_A, _rec(ARN_A, "a1"))
     assert _wait_until(exporter.started.is_set)
-    ok = scheduler.end_invocation(0.1)  # times out -> barrier cancelled
+    ok = scheduler.end_invocation(0.1)
     assert ok is False
-    # Once the exporter unblocks, the worker drains the cancelled barrier
-    # (skipping the pointless flush) and exits idle -- no permanent leak.
+    # The caller returns on time, but one detached flush remains queued so a
+    # buffered exporter can publish before the worker exits idle.
+    assert lane._queued_flush_count() == 1
     exporter.release()
     assert _wait_until(lambda: not lane._worker_alive())
-    assert exporter.flushed == 0  # cancelled barrier did not flush
+    assert exporter.flushed == 1
+
+
+def test_timed_out_buffered_export_is_published_eventually():
+    exporter = BlockingBufferedExporter()
+    scheduler = _ExportScheduler([exporter])
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "terminal", status="SUCCEEDED"))
+    assert _wait_until(exporter.started.is_set)
+
+    assert scheduler.end_invocation(0.1) is False
+    assert exporter.published == []
+    exporter.release()
+
+    assert _wait_until(lambda: not lane._worker_alive())
+    assert exporter.published == ["terminal"]
+    assert exporter.flushed == 1
 
 
 def test_repeated_timeouts_behind_blocked_exporter_stay_bounded():
-    """A blocked exporter across many warm invocations must not accumulate
-    barriers or grow queue state, must keep the SAME worker (no replacement),
-    must not execute any cancelled flush, and must drain + exit after release.
-    """
+    """Warm timeouts coalesce to one eventual flush on the same worker."""
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
     lane = scheduler._lanes[0]
 
-    # First record puts the single worker into a blocked export.
     scheduler.schedule(ARN_A, _rec(ARN_A, "a1"))
     assert _wait_until(exporter.started.is_set)
     worker = lane._worker
     assert worker is not None and worker.is_alive()
 
-    # Many warm invocations. Each schedules a record for the same ARN then ends
-    # with a short timeout; the barrier always times out because the worker is
-    # still stuck in the first export.
     for i in range(50):
         scheduler.schedule(ARN_A, _rec(ARN_A, f"a{i + 2}"))
-        ok = scheduler.end_invocation(0.02)
-        assert ok is False  # degraded every time -- worker is blocked
-        # The cancelled barrier is pulled from the queue immediately, so no
-        # _FLUSH marker lingers behind the blocked worker.
-        assert lane._queued_flush_count() == 0
-        # Each cancelled barrier can leave one generation token, but both
-        # records and tokens remain bounded by the per-execution FIFO depth.
-        assert lane._queue_len() <= 16
+        assert scheduler.end_invocation(0.02) is False
+        # At most 16 generation tokens plus one detached eventual flush.
+        assert lane._queued_flush_count() == 1
+        assert lane._queue_len() <= 17
         assert lane._pending_record_count() <= 16
 
-    # Bounded state: one in-flight ARN with a bounded pending FIFO and bounded
-    # generation tokens, with no growing pile of barriers.
-    assert lane._queue_len() <= 16
+    assert lane._queue_len() <= 17
     assert lane._pending_count() <= 1
     assert lane._pending_record_count() <= 16
-    assert lane._queued_flush_count() == 0
-    # The blocked worker was never replaced.
+    assert lane._queued_flush_count() == 1
     assert lane._worker is worker
     assert worker.is_alive()
     assert _lane_worker_count(lane) == 1
-    # No cancelled flush ran while the worker was blocked.
     assert exporter.flushed == 0
 
-    # Release: the worker drains the newest bounded window, then exits idle.
     exporter.release()
     assert _wait_until(lambda: not lane._worker_alive())
     exported = exporter.exported_values()
-    assert exported[0] == "a1"  # the in-flight record delivered first
-    assert len(exported) <= 17  # a1 plus at most 16 pending FIFO records
-    assert exported[-1] == "a51"  # newest snapshot was retained
-    # Cancelled barriers never triggered a flush, and the idle-stop path does
-    # not flush either.
-    assert exporter.flushed == 0
+    assert exported[0] == "a1"
+    assert len(exported) <= 17
+    assert exported[-1] == "a51"
+    assert exporter.flushed == 1
 
 
-def test_cancel_flush_removes_queued_barrier_immediately():
-    """Queued-barrier race: while the worker is blocked the barrier is still in
-    the queue, so cancel_flush pulls it out and completes it synchronously --
-    without waiting for the worker and without ever flushing."""
+def test_cancel_flush_replaces_queued_barrier_with_detached_flush():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
     lane = scheduler._lanes[0]
     scheduler.schedule(ARN_A, _rec(ARN_A, "a1"))
-    assert _wait_until(exporter.started.is_set)  # worker blocked in export
+    assert _wait_until(exporter.started.is_set)
     barrier = lane.enqueue_flush()
     assert lane._queued_flush_count() == 1
+
     lane.cancel_flush(barrier)
-    # Removed from the queue and completed here, without the worker.
-    assert lane._queued_flush_count() == 0
+
+    assert lane._queued_flush_count() == 1
     assert barrier.canceled is True
     assert barrier.is_done()
-    # Finish the in-flight export and go idle; the pulled barrier never flushed.
     exporter.release()
     lane.request_stop_when_idle()
     assert _wait_until(lambda: not lane._worker_alive())
-    assert exporter.flushed == 0
+    assert exporter.flushed == 1
     assert exporter.exported_values() == ["a1"]