fix(insight): preserve shaped record admission

Author: Alex Wang

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

@@ -63,7 +63,7 @@ Behavior is validated cross-SDK by the `insight` conformance suite
 > `WorkflowInsightPlugin`: listing it twice or sharing it across plugin instances
 > 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 canonical JSON. When a bound fills, it drops the oldest
+> 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

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

@@ -21,16 +21,14 @@
 
 import copy
 import logging
+import sys
 import threading
 import time
 from collections import deque
 from dataclasses import dataclass
 from typing import Any
 
-from aws_durable_execution_sdk_python_insight.truncation import (
-    json_byte_size,
-    truncate_record,
-)
+from aws_durable_execution_sdk_python_insight.truncation import truncate_record
 from aws_durable_execution_sdk_python_insight.types import InsightExporter
 
 
@@ -53,10 +51,40 @@
 # exporter. The in-flight record is not included in this count.
 _DEFAULT_MAX_PENDING_RECORDS_PER_EXECUTION = 16
 
-# Canonical JSON-byte estimate retained by one blocked lane. This keeps the
+# 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
 
+
+def _estimate_retained_size(value: Any) -> int:
+    """Estimate retained Python memory without serializing or calling render()."""
+    total = 0
+    seen: set[int] = set()
+    stack: list[Any] = [value]
+    while stack:
+        item = stack.pop()
+        identity = id(item)
+        if identity in seen:
+            continue
+        seen.add(identity)
+        try:
+            total += sys.getsizeof(item)
+        except Exception:  # noqa: BLE001 - estimation must never break a hook
+            total += 1_024
+            continue
+        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 are best-effort
+                pass
+    return total
+
+
 # Queue entry kinds.
 _RECORD = "record"
 _FLUSH = "flush"
@@ -80,11 +108,12 @@ class _FlushBarrier:
     a later, still-blocked worker skips the now-pointless flush.
     """
 
-    __slots__ = ("_event", "canceled")
+    __slots__ = ("_event", "canceled", "failed")
 
     def __init__(self) -> None:
         self._event = threading.Event()
         self.canceled = False
+        self.failed = False
 
     def complete(self) -> None:
         self._event.set()
@@ -138,25 +167,20 @@ def __init__(
         self._next_sequence = 0
         self._stop_when_idle = False
         self._worker: threading.Thread | None = None
+        self._disabled = False
 
     # -- producer API (checkpoint / invocation-end threads) -------------------
 
     def schedule(
         self,
         execution_arn: str,
         record: dict[str, Any],
-        record_size: int | None,
+        record_size: int,
     ) -> None:
         with self._cond:
-            self._stop_when_idle = False
-            if record_size is None:
-                _logger.warning(
-                    "workflow-insight: cannot measure pending record for %s on "
-                    "%s; dropping this record",
-                    execution_arn,
-                    type(self._exporter).__name__,
-                )
+            if self._disabled:
                 return
+            self._stop_when_idle = False
             size = max(0, record_size)
             if size > self._max_pending_bytes:
                 _logger.warning(
@@ -204,6 +228,11 @@ def schedule(
     def enqueue_flush(self) -> _FlushBarrier:
         barrier = _FlushBarrier()
         with self._cond:
+            if self._disabled:
+                barrier.canceled = True
+                barrier.failed = True
+                barrier.complete()
+                return barrier
             self._queue.append((_FLUSH, barrier))
             self._generation += 1
             self._ensure_worker_locked()
@@ -332,19 +361,43 @@ def _remove_record_token(self, token: _RecordToken) -> None:
                 del self._queue[index]
                 return
 
+    def _disable_locked(self, exc: Exception) -> None:
+        self._disabled = True
+        self._worker = None
+        self._pending.clear()
+        self._pending_bytes = 0
+        for kind, payload in self._queue:
+            if kind == _FLUSH:
+                barrier: _FlushBarrier = payload
+                barrier.canceled = True
+                barrier.failed = True
+                barrier.complete()
+        self._queue.clear()
+        _logger.warning(
+            "workflow-insight: could not start worker for exporter %s; "
+            "disabling this lane: %s",
+            type(self._exporter).__name__,
+            exc,
+        )
+
     def _ensure_worker_locked(self) -> None:
         # Never create a replacement while a prior worker is alive (a blocked
         # worker keeps ``_worker`` non-None). A worker that exits cleanly nulls
         # ``_worker`` under the lock before returning, so this check is a
         # race-free "start iff there is no live worker".
+        if self._disabled:
+            return
         if self._worker is None or not self._worker.is_alive():
             worker = threading.Thread(
                 target=self._run_worker,
                 name=f"workflow-insight-export-{id(self)}",
                 daemon=True,
             )
             self._worker = worker
-            worker.start()
+            try:
+                worker.start()
+            except Exception as exc:  # noqa: BLE001 - instrumentation must not break hooks
+                self._disable_locked(exc)
 
     # -- worker (single daemon thread) ---------------------------------------
 
@@ -487,7 +540,7 @@ def __init__(
 
     def schedule(self, execution_arn: str, record: dict[str, Any]) -> None:
         """Fan a canonical record out to every lane. Returns immediately."""
-        record_size = json_byte_size(record)
+        record_size = _estimate_retained_size(record)
         for lane in self._lanes:
             lane.schedule(execution_arn, record, record_size)
 
@@ -510,6 +563,8 @@ def end_invocation(self, timeout_seconds: float) -> bool:
                 # accumulate behind a blocked worker.
                 lane.cancel_flush(barrier)
                 degraded = True
+            elif barrier.failed:
+                degraded = True
         for lane in self._lanes:
             lane.request_stop_when_idle()
         if degraded:

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

@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import threading
 import time
@@ -183,6 +182,13 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Any:
         raise RuntimeError("uncopyable payload")
 
 
+class _Unsized:
+    """A payload whose custom ``__sizeof__`` raises."""
+
+    def __sizeof__(self) -> int:
+        raise RuntimeError("size unavailable")
+
+
 # -- lazy worker creation / one worker per exporter --------------------------
 
 
@@ -220,6 +226,25 @@ def test_one_worker_per_exporter():
     )
 
 
+def test_worker_start_failure_disables_lane_and_fails_barrier(monkeypatch):
+    exporter = RecordingExporter()
+    scheduler = _ExportScheduler([exporter])
+    lane = scheduler._lanes[0]
+
+    def fail_start(self):
+        raise RuntimeError("cannot start new thread")
+
+    monkeypatch.setattr(threading.Thread, "start", fail_start)
+    scheduler.schedule(ARN_A, _rec(ARN_A, "a1"))
+
+    assert lane._disabled is True
+    assert lane._pending_count() == 0
+    assert lane._pending_bytes_count() == 0
+    assert lane._queue_len() == 0
+    assert scheduler.end_invocation(0.1) is False
+    assert exporter.exported_values() == []
+
+
 def test_repeated_scheduling_does_not_grow_threads():
     base = _insight_thread_count()
     exporter = RecordingExporter()
@@ -629,7 +654,7 @@ def test_pending_record_cap_preserves_original_lane_memory_bound():
     assert exporter.exported_values() == ["a1", "c1", "d1", "a2"]
 
 
-def test_unmeasurable_record_does_not_evict_existing_backlog():
+def test_non_json_record_reaches_exporter_without_evicting_backlog():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
     lane = scheduler._lanes[0]
@@ -638,19 +663,19 @@ def test_unmeasurable_record_does_not_evict_existing_backlog():
     scheduler.schedule(ARN_B, _rec(ARN_B, "b1"))
     scheduler.schedule(ARN_C, _rec(ARN_C, "c1"))
 
-    unmeasurable = _rec(ARN_D, "bad")
-    unmeasurable["payload"] = {"not-json"}
-    scheduler.schedule(ARN_D, unmeasurable)
+    non_json = _rec(ARN_D, "custom")
+    non_json["payload"] = {"not-json"}
+    scheduler.schedule(ARN_D, non_json)
 
-    assert lane._pending_count() == 2
+    assert lane._pending_count() == 3
     exporter.release()
     scheduler.end_invocation(5.0)
-    assert exporter.exported_values() == ["inflight", "b1", "c1"]
+    assert exporter.exported_values() == ["inflight", "b1", "c1", "custom"]
 
 
 def test_individually_over_budget_record_does_not_evict_existing_backlog():
     exporter = BlockingExporter()
-    scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000)
+    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)
@@ -659,29 +684,28 @@ def test_individually_over_budget_record_does_not_evict_existing_backlog():
     scheduler.schedule(ARN_D, _rec(ARN_D, "d" * 3_000))
 
     assert lane._pending_count() == 2
-    assert lane._pending_bytes_count() <= 2_000
+    assert lane._pending_bytes_count() <= 3_500
     exporter.release()
     scheduler.end_invocation(5.0)
     exported = exporter.exported_values()
     assert exported[0] == "inflight"
     assert exported[1:] == ["b" * 700, "c" * 700]
 
 
-def test_record_sizing_exception_does_not_escape_schedule(monkeypatch):
+def test_record_sizing_exception_does_not_escape_schedule():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])
     scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
     assert _wait_until(exporter.started.is_set)
 
-    def fail_sizing(*args, **kwargs):
-        raise RecursionError("record nesting is too deep")
-
-    monkeypatch.setattr(json, "dumps", fail_sizing)
-    scheduler.schedule(ARN_B, _rec(ARN_B, "too-deep"))
+    record = _rec(ARN_B, "custom-sized")
+    record["payload"] = _Unsized()
+    scheduler.schedule(ARN_B, record)
 
-    assert scheduler._lanes[0]._pending_count() == 0
+    assert scheduler._lanes[0]._pending_count() == 1
     exporter.release()
     scheduler.end_invocation(5.0)
+    assert exporter.exported_values() == ["inflight", "custom-sized"]
 
 
 def test_pending_record_cap_evicts_true_oldest_across_arns():
@@ -703,7 +727,7 @@ def test_pending_record_cap_evicts_true_oldest_across_arns():
 
 def test_pending_byte_budget_evicts_oldest_large_record():
     exporter = BlockingExporter()
-    scheduler = _ExportScheduler([exporter], max_pending_bytes=2_000)
+    scheduler = _ExportScheduler([exporter], max_pending_bytes=3_000)
     lane = scheduler._lanes[0]
     scheduler.schedule(ARN_A, _rec(ARN_A, "inflight"))
     assert _wait_until(exporter.started.is_set)
@@ -712,7 +736,7 @@ def test_pending_byte_budget_evicts_oldest_large_record():
     scheduler.schedule(ARN_C, _rec(ARN_C, "c" * 1_500))
 
     assert lane._pending_record_count() == 1
-    assert lane._pending_bytes_count() <= 2_000
+    assert lane._pending_bytes_count() <= 3_000
     exporter.release()
     scheduler.end_invocation(5.0)
     exported = exporter.exported_values()