fix(insight): bound retained traversal work

Author: Alex Wang

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

@@ -21,7 +21,8 @@
 from __future__ import annotations
 
 import copy
-import gc
+import functools
+import itertools
 import logging
 import sys
 import threading
@@ -48,44 +49,149 @@
 _DEFAULT_MAX_PENDING_BYTES = 16_000_000
 
 
-_RETAINED_GRAPH_BOUNDARIES = (
+_UNSUPPORTED_RETAINED_GRAPH = object()
+_ATOMIC_RETAINED_TYPES = (
+    str,
+    bytes,
+    bytearray,
+    int,
+    float,
+    complex,
+    bool,
+    type(None),
+    range,
+    slice,
     type,
     types.ModuleType,
-    types.FunctionType,
-    types.BuiltinFunctionType,
     types.CodeType,
+    types.WrapperDescriptorType,
+    types.MethodDescriptorType,
 )
 
 
+class _RetainedChildren:
+    __slots__ = ("iterator",)
+
+    def __init__(self, items: Any) -> None:
+        self.iterator = iter(items)
+
+
+def _custom_retained_children(value: Any) -> list[Any]:
+    children: list[Any] = []
+    try:
+        children.append(object.__getattribute__(value, "__dict__"))
+    except Exception:  # noqa: BLE001 - custom objects may use slots only
+        pass
+    for cls in type(value).__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:
+                children.append(object.__getattribute__(value, slot))
+            except Exception:  # noqa: BLE001 - unset/custom slots are best-effort
+                pass
+    return children
+
+
+def _retained_children(value: Any) -> Any:
+    custom = _custom_retained_children(value)
+    if isinstance(value, dict):
+        return itertools.chain(dict.__iter__(value), dict.values(value), custom)
+    if isinstance(value, list):
+        return itertools.chain(list.__iter__(value), custom)
+    if isinstance(value, tuple):
+        return itertools.chain(tuple.__iter__(value), custom)
+    if isinstance(value, set):
+        return itertools.chain(set.__iter__(value), custom)
+    if isinstance(value, frozenset):
+        return itertools.chain(frozenset.__iter__(value), custom)
+    if isinstance(value, deque):
+        return itertools.chain(deque.__iter__(value), custom)
+    if isinstance(value, memoryview):
+        return (value.obj,)
+    if isinstance(value, functools.partial):
+        return (value.func, value.args, value.keywords)
+    if isinstance(value, types.FunctionType):
+        closure = []
+        for cell in value.__closure__ or ():
+            try:
+                closure.append(cell.cell_contents)
+            except ValueError:
+                pass
+        return itertools.chain(closure, (value.__defaults__, value.__kwdefaults__))
+    if isinstance(value, types.MethodType):
+        return (value.__self__, value.__func__)
+    if isinstance(value, types.BuiltinFunctionType):
+        owner = value.__self__
+        return () if owner is None or isinstance(owner, types.ModuleType) else (owner,)
+    if isinstance(value, types.MethodWrapperType):
+        return (value.__self__,)
+    if isinstance(value, types.GeneratorType):
+        frame = value.gi_frame
+        return () if frame is None else (frame.f_locals, value.gi_yieldfrom)
+    if isinstance(value, _ATOMIC_RETAINED_TYPES):
+        return ()
+    if custom:
+        return custom
+    return _UNSUPPORTED_RETAINED_GRAPH
+
+
+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
+    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
+
+
 def _estimate_retained_size(value: Any, max_size: int | None = None) -> int:
-    """Estimate a bounded retained graph without serializing or calling render()."""
+    """Estimate retained memory with bounded, non-overridable traversal."""
     total = 0
     seen: set[int] = set()
     stack: list[Any] = [value]
     while stack:
         item = stack.pop()
+        if isinstance(item, _RetainedChildren):
+            try:
+                child = next(item.iterator)
+            except StopIteration:
+                continue
+            except Exception:  # noqa: BLE001 - fail closed on malformed iterators
+                return max_size + 1 if max_size is not None else total + 1_024
+            stack.append(item)
+            stack.append(child)
+            continue
         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
+        total += _retained_shallow_size(item)
         if max_size is not None and total > max_size:
             return max_size + 1
-        try:
-            referents = gc.get_referents(item)
-        except Exception:  # noqa: BLE001 - estimation must never break a hook
-            continue
-        for referent in referents:
-            # Type/module/function/code objects lead into process-global graphs,
-            # not memory retained specifically by this record. Bound methods,
-            # partial args, generator iterators, slots, buffers, and container
-            # subclasses remain traversable through their other referents.
-            if isinstance(referent, _RETAINED_GRAPH_BOUNDARIES):
-                continue
-            stack.append(referent)
+        children = _retained_children(item)
+        if children is _UNSUPPORTED_RETAINED_GRAPH:
+            return max_size + 1 if max_size is not None else total + 1_024
+        stack.append(_RetainedChildren(children))
     return total
 
 

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

@@ -79,12 +79,6 @@
 def _claim_exporter_lanes(lanes: list[Any]) -> None:
     """Give each exporter object to at most one live scheduler lane."""
 
-    def release(lane_ref: weakref.ReferenceType[Any]) -> None:
-        with _exporter_owner_lock:
-            _exporter_owners[:] = [
-                existing for existing in _exporter_owners if existing is not lane_ref
-            ]
-
     with _exporter_owner_lock:
         _exporter_owners[:] = [
             lane_ref for lane_ref in _exporter_owners if lane_ref() is not None
@@ -98,7 +92,9 @@ def release(lane_ref: weakref.ReferenceType[Any]) -> None:
                         "the same exporter instance cannot be shared across "
                         "Workflow Insight plugin instances"
                     )
-        _exporter_owners.extend(weakref.ref(lane, release) for lane in lanes)
+        # Callback-free weakrefs avoid lock re-entry during synchronous finalization.
+        # Dead entries are pruned at the start of every subsequent claim.
+        _exporter_owners.extend(weakref.ref(lane) for lane in lanes)
 
 
 def _parse_execution_arn(execution_arn: str) -> dict[str, str]:

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

@@ -18,6 +18,7 @@
 
 import pytest
 
+import aws_durable_execution_sdk_python_insight.plugin as insight_plugin_module
 from aws_durable_execution_sdk_python_insight import (
     EmitMode,
     OperationDetail,
@@ -278,6 +279,10 @@ def test_exporter_plugin_cycle_is_not_rooted_by_ownership_registry():
     exporter_ref = weakref.ref(exporter)
     plugin_ref = weakref.ref(plugin)
     lane_ref = weakref.ref(lane)
+    assert all(
+        owner_ref.__callback__ is None
+        for owner_ref in insight_plugin_module._exporter_owners
+    )
 
     del lane
     del plugin

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

@@ -14,6 +14,7 @@
 import functools
 import threading
 import time
+import types
 from typing import Any
 
 from aws_durable_execution_sdk_python_insight._export_scheduler import (
@@ -217,6 +218,9 @@ def __init__(self) -> None:
         super().__init__([None] * 10_000)
         self.iterated = False
 
+    def __sizeof__(self) -> int:
+        return 1
+
     def __iter__(self):
         self.iterated = True
         return super().__iter__()
@@ -743,6 +747,34 @@ def __iter__(self):
     assert hidden.iterated is False
 
 
+def test_retained_size_counts_closures_and_bound_builtin_owners():
+    closure_buffer = bytearray(4_000)
+
+    def closure() -> bytearray:
+        return closure_buffer
+
+    bound_owner = [bytearray(4_000)]
+    wrapper_owner = [bytearray(4_000)]
+    method_wrapper = wrapper_owner.__str__
+    assert isinstance(method_wrapper, types.MethodWrapperType)
+    payloads = [closure, bound_owner.append, method_wrapper]
+
+    for payload in payloads:
+        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, "retained-callable")
+        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)
+
+
 def test_timed_out_barrier_flushes_eventually_and_worker_exits():
     exporter = BlockingExporter()
     scheduler = _ExportScheduler([exporter])