fix(insight): preserve bounded FIFO records

Author: Alex Wang

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

@@ -71,18 +71,54 @@ def _estimate_retained_size(value: Any) -> int:
             total += sys.getsizeof(item)
         except Exception:  # noqa: BLE001 - estimation must never break a hook
             total += 1_024
+        try:
+            if isinstance(item, dict):
+                stack.extend(item.keys())
+                stack.extend(item.values())
+            elif isinstance(item, (list, tuple, set, frozenset, deque)):
+                stack.extend(item)
+            else:
+                try:
+                    stack.append(vars(item))
+                except Exception:  # noqa: BLE001 - custom objects may use slots
+                    pass
+                for cls in type(item).__mro__:
+                    slots = vars(cls).get("__slots__", ())
+                    if isinstance(slots, str):
+                        slots = (slots,)
+                    for slot in slots:
+                        if slot in {"__dict__", "__weakref__"}:
+                            continue
+                        if slot.startswith("__") and not slot.endswith("__"):
+                            slot = f"_{cls.__name__.lstrip('_')}{slot}"
+                        try:
+                            stack.append(getattr(item, slot))
+                        except Exception:  # noqa: BLE001 - unset/custom slots are best-effort
+                            pass
+        except Exception:  # noqa: BLE001 - traversal must never break a hook
+            pass
+    return total
+
+
+def _copy_record_containers(record: dict[str, Any]) -> dict[str, Any]:
+    """Copy built-in containers while treating custom values as opaque leaves."""
+    memo: dict[int, Any] = {}
+    seen: set[int] = set()
+    stack: list[Any] = [record]
+    while stack:
+        item = stack.pop()
+        identity = id(item)
+        if identity in seen:
             continue
-        if isinstance(item, dict):
+        seen.add(identity)
+        if type(item) is dict:
             stack.extend(item.keys())
             stack.extend(item.values())
-        elif isinstance(item, (list, tuple, set, frozenset, deque)):
+        elif type(item) in {list, tuple, set, frozenset, deque}:
             stack.extend(item)
         else:
-            try:
-                stack.append(vars(item))
-            except Exception:  # noqa: BLE001 - custom objects are best-effort
-                pass
-    return total
+            memo[identity] = item
+    return copy.deepcopy(record, memo)
 
 
 # Queue entry kinds.
@@ -183,13 +219,17 @@ def schedule(
             self._stop_when_idle = False
             size = max(0, record_size)
             if size > self._max_pending_bytes:
+                superseded = execution_arn in self._pending
+                if superseded:
+                    self._drop_pending_execution(execution_arn)
                 _logger.warning(
                     "workflow-insight: pending record for %s on %s exceeds the "
-                    "byte budget (%d > %d); dropping this record",
+                    "byte budget (%d > %d); dropping this record%s",
                     execution_arn,
                     type(self._exporter).__name__,
                     size,
                     self._max_pending_bytes,
+                    " and its superseded pending FIFO" if superseded else "",
                 )
                 return
             pending = self._pending.get(execution_arn)
@@ -439,18 +479,15 @@ def _run_worker(self) -> None:
 
     def _export_one(self, record: dict[str, Any]) -> None:
         exporter = self._exporter
-        # Copy for exporter isolation: every lane shares the same canonical
-        # record, and truncation/export must never mutate what another lane
-        # sees. If the copy fails we must NOT fall back to the shared record --
-        # exporting the alias would let this lane's truncation mutate the object
-        # other lanes still read, breaking workflow isolation. Treat a copy
-        # failure like a render/truncation failure: log and skip this record for
-        # this lane, then continue processing the lane's queue.
+        # Copy the record's built-in containers for lane isolation, but preserve
+        # custom values as opaque leaves for exporter-specific rendering. This
+        # keeps one lane's render/truncation mutations out of other lanes without
+        # requiring custom-renderable values to implement ``deepcopy``.
         try:
-            local = copy.deepcopy(record)
-        except Exception as exc:  # noqa: BLE001 - a non-copyable payload must not alias the shared record or break the lane
+            local = _copy_record_containers(record)
+        except Exception as exc:  # noqa: BLE001 - malformed containers must not break the lane
             _logger.warning(
-                "workflow-insight: record copy failed for exporter %s; "
+                "workflow-insight: record container copy failed for exporter %s; "
                 "skipping export for this record: %s",
                 type(exporter).__name__,
                 exc,

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

@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-import logging
 import threading
 import time
 from typing import Any
@@ -182,6 +181,18 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Any:
         raise RuntimeError("uncopyable payload")
 
 
+class _SlottedPayload:
+    __slots__ = ("payload",)
+
+    def __init__(self, payload: Any) -> None:
+        self.payload = payload
+
+
+class _UnsizedSlottedPayload(_SlottedPayload):
+    def __sizeof__(self) -> int:
+        raise RuntimeError("size unavailable")
+
+
 class _Unsized:
     """A payload whose custom ``__sizeof__`` raises."""
 
@@ -318,58 +329,45 @@ def test_pending_fifo_cap_drops_oldest_and_retains_terminal():
 # -- copy failure isolation ---------------------------------------------------
 
 
-def test_deepcopy_failure_skips_record_and_lane_continues(caplog):
-    exporter = RecordingExporter()
+def test_uncopyable_custom_value_reaches_exporter_render():
+    class CustomRenderExporter(RecordingExporter):
+        def __init__(self) -> None:
+            super().__init__(max_record_size_bytes=10_000)
+            self.rendered_values: list[str] = []
+
+        def render(self, record: dict[str, Any]) -> Any:
+            value = record["payload"]["value"]
+            self.rendered_values.append(value)
+            return {"value": value}
+
+    exporter = CustomRenderExporter()
     scheduler = _ExportScheduler([exporter])
-    # A record whose deepcopy raises must be skipped for this lane -- never
-    # exported by aliasing the shared object -- and the lane must keep draining.
-    bad = _rec(ARN_A, "bad")
-    bad["payload"] = _Uncopyable()
-    good = _rec(ARN_B, "good")
-    with caplog.at_level(
-        logging.WARNING, logger="aws_durable_execution_sdk_python_insight"
-    ):
-        scheduler.schedule(ARN_A, bad)  # queued first: copy fails -> skipped
-        scheduler.schedule(ARN_B, good)  # queued behind it: must still export
-        # The good record delivering proves the lane continued past the failure;
-        # a single-lane worker drains FIFO, so "bad" was processed (and skipped)
-        # before "good" ran.
-        assert _wait_until(lambda: exporter.exported_values() == ["good"])
-        scheduler.end_invocation(5.0)
-    # The exporter was never called for the un-copyable record.
-    assert exporter.exported_values() == ["good"]
-    # The failure was logged through the module logger.
-    assert any(
-        "record copy failed" in record.getMessage()
-        for record in caplog.records
-        if record.name == "aws_durable_execution_sdk_python_insight"
-    )
+    record = _rec(ARN_A, "before-render")
+    record["payload"] = _Uncopyable()
 
+    scheduler.schedule(ARN_A, record)
+    scheduler.end_invocation(5.0)
+
+    assert exporter.rendered_values == ["safe"]
+    assert exporter.exported_values() == ["before-render"]
 
-def test_deepcopy_failure_does_not_alias_shared_record():
-    # Before the fix a copy failure aliased the shared record and passed it to
-    # truncate_record -> render, which could mutate the canonical object other
-    # lanes still read. With the fix the record is skipped before render, so it
-    # is never aliased or mutated in place.
+
+def test_uncopyable_custom_value_does_not_alias_record_containers():
     class MutatingRenderExporter(RecordingExporter):
         def render(self, record: dict[str, Any]) -> Any:
-            record["mutated"] = True  # would corrupt an aliased shared record
+            record["mutated"] = True
             return record
 
     exporter = MutatingRenderExporter()
     scheduler = _ExportScheduler([exporter])
-    bad = _rec(ARN_A, "bad")
-    bad["payload"] = _Uncopyable()
-    scheduler.schedule(ARN_A, bad)
-    # A good record behind it lets us deterministically wait for the lane to
-    # drain past the bad one (single lane drains FIFO).
-    scheduler.schedule(ARN_B, _rec(ARN_B, "good"))
-    assert _wait_until(lambda: exporter.exported_values() == ["good"])
+    record = _rec(ARN_A, "custom")
+    record["payload"] = _Uncopyable()
+
+    scheduler.schedule(ARN_A, record)
     scheduler.end_invocation(5.0)
-    # render never ran on the un-copyable record, so the canonical object was
-    # neither aliased into export nor mutated in place.
-    assert "mutated" not in bad
-    assert exporter.exported_values() == ["good"]
+
+    assert "mutated" not in record
+    assert exporter.exported_values() == ["custom"]
 
 
 # -- non-blocking hook return / fast-vs-slow isolation -----------------------
@@ -692,6 +690,47 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog():
     assert exported[1:] == ["b" * 700, "c" * 700]
 
 
+def test_over_budget_replacement_removes_superseded_same_arn_only():
+    exporter = BlockingExporter()
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=3_500)
+    lane = scheduler._lanes[0]
+    scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
+    assert _wait_until(exporter.started.is_set)
+    scheduler.schedule(ARN_A, _rec(ARN_A, "stale-running"))
+    scheduler.schedule(ARN_B, _rec(ARN_B, "unrelated"))
+    scheduler.schedule(
+        ARN_A,
+        _rec(ARN_A, "terminal" * 500, status="SUCCEEDED"),
+    )
+
+    assert lane._pending_count() == 1
+    assert lane._pending_bytes_count() <= 3_500
+    exporter.release()
+    scheduler.end_invocation(5.0)
+    assert exporter.exported_values() == ["inflight", "unrelated"]
+
+
+def test_retained_size_traverses_slots_after_shallow_size_failure():
+    for payload in (
+        _SlottedPayload("x" * 4_000),
+        _UnsizedSlottedPayload("x" * 4_000),
+    ):
+        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, "opaque")
+        record["payload"] = payload
+        scheduler.schedule(ARN_B, record)
+
+        assert lane._pending_count() == 0
+        assert lane._pending_bytes_count() == 0
+        exporter.release()
+        scheduler.end_invocation(5.0)
+        assert exporter.exported_values() == ["inflight"]
+
+
 def test_record_sizing_exception_does_not_escape_schedule():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])