fix(insight): bound FIFO in-flight retention

Author: Alex Wang

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

@@ -25,7 +25,6 @@
 import functools
 import itertools
 import logging
-import sys
 import threading
 import time
 import types
@@ -73,7 +72,6 @@
     bool,
     type(None),
     range,
-    slice,
     type,
     types.ModuleType,
     types.CodeType,
@@ -135,6 +133,8 @@ def _retained_children(value: Any) -> Any:
         return itertools.chain(frozenset.__iter__(value), custom)
     if isinstance(value, deque):
         return itertools.chain(deque.__iter__(value), custom)
+    if isinstance(value, slice):
+        return itertools.chain((value.start, value.stop, value.step), custom)
     if isinstance(value, memoryview):
         return itertools.chain((value.obj,), custom)
     if isinstance(value, functools.partial):
@@ -175,26 +175,25 @@ def _retained_children(value: Any) -> Any:
 
 
 def _retained_shallow_size(value: Any) -> int:
-    try:
-        size = sys.getsizeof(value)
-    except Exception:  # noqa: BLE001 - estimation must never break a hook
-        size = 1_024
+    """Return shallow size without dispatching to user-defined ``__sizeof__``."""
     try:
         if isinstance(value, dict):
-            size = max(size, dict.__sizeof__(value))
-        elif isinstance(value, list):
-            size = max(size, list.__sizeof__(value))
-        elif isinstance(value, tuple):
-            size = max(size, tuple.__sizeof__(value))
-        elif isinstance(value, set):
-            size = max(size, set.__sizeof__(value))
-        elif isinstance(value, frozenset):
-            size = max(size, frozenset.__sizeof__(value))
-        elif isinstance(value, deque):
-            size = max(size, deque.__sizeof__(value))
-    except Exception:  # noqa: BLE001 - base sizing remains best-effort
-        pass
-    return size
+            return dict.__sizeof__(value)
+        if isinstance(value, list):
+            return list.__sizeof__(value)
+        if isinstance(value, tuple):
+            return tuple.__sizeof__(value)
+        if isinstance(value, set):
+            return set.__sizeof__(value)
+        if isinstance(value, frozenset):
+            return frozenset.__sizeof__(value)
+        if isinstance(value, deque):
+            return deque.__sizeof__(value)
+        if type(value) in _ATOMIC_RETAINED_TYPES + _SAFE_OPAQUE_RETAINED_TYPES:
+            return value.__sizeof__()
+        return object.__sizeof__(value)
+    except Exception:  # noqa: BLE001 - estimation must never break a hook
+        return 1_024
 
 
 def _estimate_retained_size(value: Any, max_size: int | None = None) -> int:
@@ -313,6 +312,7 @@ def __init__(
         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
         # 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
@@ -412,6 +412,12 @@ def cancel_flush(self, barrier: _FlushBarrier) -> None:
         """Stop waiting for a timed-out barrier while retaining one later flush."""
         with self._cond:
             barrier.canceled = True
+            if not any(
+                kind == _FLUSH and payload is barrier for kind, payload in self._queue
+            ):
+                # The worker already owns this barrier. Do not erase a detached
+                # flush installed by a later invocation while this one was in flight.
+                return
             # 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.
@@ -481,7 +487,10 @@ def _enforce_pending_record_cap(self) -> None:
             )
 
     def _enforce_pending_byte_cap(self) -> None:
-        while self._pending_bytes > self._max_pending_bytes and self._pending:
+        while (
+            self._pending_bytes + self._inflight_bytes > self._max_pending_bytes
+            and self._pending
+        ):
             old_arn = self._oldest_pending_arn()
             dropped_size = self._drop_oldest_pending_record(old_arn)
             _logger.warning(
@@ -523,6 +532,7 @@ def _disable_locked(self, exc: Exception) -> None:
         self._worker = None
         self._pending.clear()
         self._pending_bytes = 0
+        self._inflight_bytes = 0
         for kind, payload in self._queue:
             if kind == _FLUSH and payload is not None:
                 barrier: _FlushBarrier = payload
@@ -571,14 +581,17 @@ def _run_worker(self) -> None:
                     return
                 kind, payload = self._queue.popleft()
                 record: dict[str, Any] | None = None
+                record_size = 0
                 if kind == _RECORD:
                     token: _RecordToken = payload
                     execution_arn, generation = token
                     pending = self._pending.get(execution_arn)
                     if not pending or pending[0].generation != generation:
                         continue
                     pending_record = pending.popleft()
-                    self._pending_bytes -= pending_record.size
+                    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.
@@ -587,7 +600,11 @@ def _run_worker(self) -> None:
                         del self._pending[execution_arn]
 
             if kind == _RECORD and record is not None:
-                self._export_one(record)
+                try:
+                    self._export_one(record)
+                finally:
+                    with self._cond:
+                        self._inflight_bytes -= record_size
             else:  # _FLUSH
                 barrier: _FlushBarrier | None = payload
                 self._flush()
@@ -658,6 +675,10 @@ def _pending_bytes_count(self) -> int:
         with self._cond:
             return self._pending_bytes
 
+    def _retained_bytes_count(self) -> int:
+        with self._cond:
+            return self._pending_bytes + self._inflight_bytes
+
     def _queue_len(self) -> int:
         with self._cond:
             return len(self._queue)

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

@@ -720,16 +720,16 @@ def test_non_json_record_reaches_exporter_without_evicting_backlog():
 
 def test_individually_over_budget_record_does_not_evict_existing_backlog():
     exporter = BlockingExporter()
-    scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000)
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=5_000)
     lane = scheduler._lanes[0]
     scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
     assert _wait_until(exporter.started.is_set)
     scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 700))
     scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 700))
-    scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000))
+    scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 5_000))
 
     assert lane._pending_count() == 2
-    assert lane._pending_bytes_count() <= 3_000
+    assert lane._pending_bytes_count() <= 5_000
     exporter.release()
     scheduler.end_invocation(5.0)
     exported = exporter.exported_values()
@@ -1114,3 +1114,79 @@ def test_cancel_flush_after_pop_lets_worker_complete_barrier():
     assert exporter.calls.count(("flush", None)) == 1
     lane.request_stop_when_idle()
     assert _wait_until(lambda: not lane._worker_alive())
+
+
+def test_retained_size_counts_slice_referents():
+    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, "slice")
+    record["payload"] = slice(bytearray(4_000), None)
+    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_retained_size_does_not_dispatch_custom_sizeof():
+    called = threading.Event()
+
+    class CustomSized:
+        def __sizeof__(self) -> int:
+            called.set()
+            raise AssertionError("custom __sizeof__ must not run")
+
+    exporter = RecordingExporter()
+    scheduler = _ExportScheduler([exporter])
+    record = _rec(ARN_A, "custom-sized")
+    record["payload"] = CustomSized()
+
+    scheduler.schedule(ARN_A, record)
+    scheduler.end_invocation(5.0)
+
+    assert called.is_set() is False
+    assert exporter.exported_values() == ["custom-sized"]
+
+
+def test_inflight_record_remains_charged_until_export_returns():
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000)
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "a" * 1_500))
+    assert _wait_until(exporter.started.is_set)
+
+    scheduler.schedule(ARN_B, _rec(ARN_B, "b" * 1_500))
+
+    assert lane._pending_count() == 0
+    assert lane._retained_bytes_count() <= 3_000
+    exporter.release()
+    scheduler.end_invocation(5.0)
+    assert exporter.exported_values() == ["a" * 1_500]
+    assert lane._retained_bytes_count() == 0
+
+
+def test_cancel_popped_barrier_preserves_later_detached_flush():
+    exporter = BlockingFlushExporter()
+    scheduler = _ExportScheduler([exporter])
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "a1"))
+    older = lane.enqueue_flush()
+    assert _wait_until(exporter.flush_started.is_set)
+
+    later = lane.enqueue_flush()
+    lane.cancel_flush(later)
+    assert lane._queued_flush_count() == 1
+    lane.cancel_flush(older)
+
+    assert lane._queued_flush_count() == 1
+    assert later.is_done()
+    exporter.release_flush()
+    assert _wait_until(older.is_done)
+    lane.request_stop_when_idle()
+    assert _wait_until(lambda: not lane._worker_alive())
+    assert exporter.calls.count(("flush", None)) == 2