Hold the hash witness's axes by weak reference

src/ezmsg/baseproc/stateful.py

@@ -3,6 +3,7 @@
 import pickle
 import typing
 import warnings
+import weakref
 from abc import ABC, abstractmethod
 
 from ezmsg.util.messages.axisarray import AxisArray
@@ -61,6 +62,24 @@ def _axis_value(axis: typing.Any) -> typing.Any:
     return None if gain is None else (gain, axis.offset)
 
 
+def _dead_ref() -> None:
+    return None
+
+
+def _weak(axis: typing.Any) -> typing.Callable[[], typing.Any] | None:
+    """A weak reference to *axis*; None for a dimension with no axis.
+
+    An axis type that cannot be weakly referenced gets a reference that is
+    always dead, so it is compared by value every time -- never held strongly.
+    """
+    if axis is None:
+        return None
+    try:
+        return weakref.ref(axis)
+    except TypeError:
+        return _dead_ref
+
+
 def _build_witness(
     message: typing.Any,
     dims: list[str],
@@ -91,7 +110,14 @@ def _build_witness(
     # the axis *value*, for when it cannot -- most importantly on the far side of
     # a process boundary, where unpickling hands out a new object per message but
     # the fingerprint rides along already computed.
-    kept = tuple((dim, axes.get(dim), _axis_value(axes.get(dim))) for dim in dims if dim not in exclude)
+    #
+    # The object is held by weak reference. A strong one would keep the first
+    # message's axes alive for as long as the hash holds, and on the far side of
+    # a process boundary their data are zero-copy views into the channel's
+    # shared memory: pinning them stops the channel from releasing that memory
+    # when the publisher grows it (ezmsg-org/ezmsg#272). A reference that has
+    # died answers None and the check falls through to the value.
+    kept = tuple((dim, _weak(axes.get(dim)), _axis_value(axes.get(dim))) for dim in dims if dim not in exclude)
     # The stream axis is a new object every message on any path -- its offset
     # advances -- so it is compared by value always.
     streamed = tuple((dim, getattr(axes.get(dim), "gain", None)) for dim in dims if dim in exclude)
@@ -107,7 +133,7 @@ def _build_witness(
         def validate(
             msg: typing.Any,
             _kd: str = kept_dim,
-            _ka: typing.Any = kept_axis,
+            _ka: typing.Callable[[], typing.Any] = kept_axis,
             _kv: typing.Any = kept_value,
             _sd: str = stream_dim,
             _sg: float = stream_gain,
@@ -121,7 +147,8 @@ def validate(
             axes = msg.axes
             try:
                 axis = axes[_kd]
-                if axis is not _ka and _axis_value(axis) != _kv:
+                # `axis is None` guards a dead reference, which also answers None.
+                if (axis is None or axis is not _ka()) and _axis_value(axis) != _kv:
                     return False
                 return (
                     axes[_sd].gain == _sg
@@ -152,9 +179,13 @@ def validate(
             _check_key: bool = include_key,
         ) -> bool:
             axes = msg.axes
-            for dim, axis, value in _kept:
+            for dim, ref, value in _kept:
                 incoming = axes.get(dim)
-                if incoming is not axis and (value is None or _axis_value(incoming) != value):
+                if ref is None:
+                    # The dimension had no axis; it matches only while it still has none.
+                    if incoming is not None:
+                        return False
+                elif (incoming is None or incoming is not ref()) and (value is None or _axis_value(incoming) != value):
                     return False
             for dim, gain in _streamed:
                 # No `is not None` shortcut on the axis: an excluded dimension

tests/test_hash_witness.py

@@ -196,6 +196,70 @@ def test_it_survives_nothing_it_should_not(self):
         assert probe._message_hash(m) == recomputed(m)
 
 
+def msg3(feat_axis: CoordinateAxis | None, offset: float = 0.0) -> AxisArray:
+    """A `(time, ch, feat)` message, which takes the generic validator."""
+    axes = {
+        "time": AxisArray.TimeAxis(fs=100.0, offset=offset),
+        "ch": CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"]),
+    }
+    if feat_axis is not None:
+        axes["feat"] = feat_axis
+    return AxisArray(np.zeros((4, 2, 3), np.float32), dims=["time", "ch", "feat"], axes=axes, stream_dim="time")
+
+
+def feat() -> CoordinateAxis:
+    return CoordinateAxis(data=np.array(["x", "y", "z"]), dims=["feat"])
+
+
+class TestTheWitnessHoldsNothing:
+    """The witness must not keep a message's axes alive. Across a process
+    boundary their data are zero-copy views into the channel's shared memory,
+    and a pinned view stops the channel from releasing it (ezmsg#272)."""
+
+    @pytest.mark.parametrize(
+        "build",
+        [
+            pytest.param(lambda: msg(["a", "b"]), id="specialised"),
+            pytest.param(lambda: msg3(feat()), id="generic"),
+        ],
+    )
+    def test_the_first_messages_axes_are_released(self, build):
+        import gc
+        import weakref
+
+        probe = Probe()
+        m = build()
+        refs = [weakref.ref(ax) for dim, ax in m.axes.items() if dim != "time"]
+        probe._hash = probe._message_hash(m)
+        assert probe._hash_witness is not None
+        del m
+        gc.collect()
+        assert all(r() is None for r in refs)
+
+    def test_after_the_axis_dies_equal_content_still_hits(self):
+        probe = Probe()
+        probe._hash = first = probe._message_hash(msg(["a", "b"]))
+        gc_collect()
+        m = msg(["a", "b"], offset=0.5)
+        assert probe._hash_witness[0](m)  # answered by value, not by a reset
+        assert probe._message_hash(m) == recomputed(m) == first
+
+    def test_after_the_axis_dies_losing_it_is_not_masked(self):
+        """A dead reference answers None, like a missing axis does; the two must
+        not be mistaken for a match."""
+        probe = Probe()
+        probe._hash = first = probe._message_hash(msg3(feat()))
+        gc_collect()
+        m = msg3(None, offset=0.5)
+        assert probe._message_hash(m) == recomputed(m) != first
+
+
+def gc_collect() -> None:
+    import gc
+
+    gc.collect()
+
+
 DIMSETS = [
     (["time", "ch"], "time"),
     (["win", "time", "ch"], "win"),