From 7d642647485d231c555ca4e9a5335e9020e1052d Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Wed, 30 Sep 2026 20:13:38 -0400 Subject: [PATCH] Hold the hash witness's axes by weak reference (#17) The witness kept the axis objects of the message that set the hash, for an identity fast path. Across a process boundary their data are zero-copy views into the channel's shared memory, so the first message's segment stayed exported for the transformer's lifetime, and a publisher growing its SHM killed the subscriber's channel (ezmsg-org/ezmsg#272). The witness now holds weakref.ref(axis): the fast path still hits when a producer reuses its axes, nothing is pinned, and a dead reference falls through to the value comparison. A dead reference is never mistaken for a dimension that lost its axis. Axis types that cannot be weakly referenced are compared by value every time rather than held strongly. --- src/ezmsg/baseproc/stateful.py | 41 +++++++++++++++++++--- tests/test_hash_witness.py | 64 ++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 5 deletions(-) diff --git a/src/ezmsg/baseproc/stateful.py b/src/ezmsg/baseproc/stateful.py index 925dd60..2b9eae7 100644 --- a/src/ezmsg/baseproc/stateful.py +++ b/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 diff --git a/tests/test_hash_witness.py b/tests/test_hash_witness.py index 81dfcda..dec85f3 100644 --- a/tests/test_hash_witness.py +++ b/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"),