Copy, don't retain, message data carried across calls

src/ezmsg/sigproc/binned_aggregate.py

@@ -54,6 +54,7 @@
 )
 
 from .aggregate import AggregationFunction, aggregate_slices, needs_coordinates
+from .util.array import xp_copy
 from .util.binning import BinSchedule, BinStep
 from .util.message import is_empty_along
 
@@ -269,7 +270,10 @@ def _process(self, message: AxisArray) -> AxisArray:
 
         if step.n_bins == 0:
             # No bin completes in this chunk; grow the carry and emit nothing.
-            self._state.carry = message.data if carry is None else xp.concat((carry, message.data), axis=axis_idx)
+            # `xp_copy` is intentional.
+            self._state.carry = (
+                xp_copy(message.data) if carry is None else xp.concat((carry, message.data), axis=axis_idx)
+            )
             return self._empty_like(message, axis_idx, step)
 
         # Prepend the carried partial-bin samples so bin 0 spans carry + current.
@@ -288,9 +292,10 @@ def _process(self, message: AxisArray) -> AxisArray:
 
         # Leftover after the last completed bin becomes the next chunk's carry
         # (its length is step.carry_count, tracked by the schedule).
+        # Copied, not viewed, intentionally.
         last_work = ends_work[-1]
         self._state.carry = (
-            slice_along_axis(work, slice(last_work, None), axis=axis_idx) if step.carry_count > 0 else None
+            xp_copy(slice_along_axis(work, slice(last_work, None), axis=axis_idx)) if step.carry_count > 0 else None
         )
 
         return replace(

src/ezmsg/sigproc/diff.py

@@ -18,6 +18,8 @@
 from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
 from ezmsg.util.messages.util import replace
 
+from .util.array import xp_copy
+
 
 class DiffSettings(ez.Settings):
     axis: str | None = None
@@ -50,7 +52,10 @@ def _hash_message(self, message: AxisArray) -> int:
 
     def _reset_state(self, message) -> None:
         ax_idx = message.get_axis_idx(self.settings.axis)
-        self.state.last_dat = slice_along_axis(message.data, slice(0, 1), axis=ax_idx)
+        # Copied for the same reason as in `_process`: state must never alias the
+        # message's (possibly shared-memory-backed) buffer, even though this one
+        # happens to be overwritten before the call returns.
+        self.state.last_dat = xp_copy(slice_along_axis(message.data, slice(0, 1), axis=ax_idx))
         if self.settings.scale_by_fs:
             ax_info = message.get_axis(self.settings.axis)
             if hasattr(ax_info, "data"):
@@ -68,8 +73,8 @@ def _process(self, message: AxisArray) -> AxisArray:
             xp.concat((self.state.last_dat, message.data), axis=ax_idx),
             axis=ax_idx,
         )
-        # Prepare last_dat for next iteration
-        self.state.last_dat = slice_along_axis(message.data, slice(-1, None), axis=ax_idx)
+        # Prepare last_dat for next iteration. Copied, not viewed, intentionally.
+        self.state.last_dat = xp_copy(slice_along_axis(message.data, slice(-1, None), axis=ax_idx))
         # Scale by fs if requested. This converts the diff to a derivative. e.g., diff of position becomes velocity.
         if self.settings.scale_by_fs:
             ax_info = message.get_axis(axis)

tests/helpers/recycled_shm.py

@@ -0,0 +1,138 @@
+"""Utilities for catching transformers that retain a message's buffer.
+
+Across a cross-process link a subscriber does not own the bytes of the message
+it receives. ezmsg serializes with PEP 574 out-of-band buffers
+(``pickle.dumps(..., protocol=5, buffer_callback=...)``), so a numpy array
+deserializes as a *view* onto the publisher's shared-memory slot rather than as
+the owner of a copy, and ``Subscriber.recv_zero_copy`` says so explicitly: the
+message "should not be modified or stored beyond the context manager's scope".
+The publisher writes into slot ``msg_id % num_buffers`` of a ring and is free to
+reuse a slot as soon as the subscriber's context exits.
+
+A transformer that keeps ``message.data`` -- or a *view* of it, such as a tail
+slice carried across chunk boundaries -- therefore reads recycled bytes on its
+next call. Nothing raises: the array stays a valid object of the right shape and
+dtype and simply contains different numbers.
+
+None of this is visible in an ordinary test. In-process publishers pass the
+object by reference with no serialization at all, and even a marshalled message
+is fine as long as its slot is never overwritten. So these helpers do two
+things a normal test does not: they marshal each message through ezmsg's own
+``Marshal``, and they compress the publisher's ring to a *single* slot so that
+reuse is deterministic on the very next message instead of ``num_buffers``
+(default 32) later.
+
+Typical use -- run the same inputs both ways and require the outputs to agree::
+
+    assert_survives_buffer_recycling(
+        lambda: DiffTransformer(DiffSettings(axis="time")), messages
+    )
+
+Equal-sized messages make the strongest test: the next message lands on exactly
+the bytes a retained view points at, so a retained view reads plausible garbage
+rather than something obviously wrong.
+"""
+
+import contextlib
+import typing
+
+import numpy as np
+from ezmsg.core.messagemarshal import Marshal
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+SLOT_BYTES = 1 << 22
+"""Default slot size. Only has to fit one message; oversizing costs nothing."""
+
+
+class RecycledSlot:
+    """A publisher's shared-memory ring, compressed to a single slot.
+
+    Each :meth:`publish` overwrites the one slot and yields a message whose
+    arrays are views onto it, exactly as a subscriber would receive over a
+    cross-process link once the real 32-slot ring has wrapped around.
+    """
+
+    def __init__(self, nbytes: int = SLOT_BYTES) -> None:
+        self._slot = memoryview(bytearray(nbytes))
+        self._msg_id = 0
+
+    @contextlib.contextmanager
+    def publish(self, msg: AxisArray) -> typing.Iterator[AxisArray]:
+        """Serialize ``msg`` into the slot and yield the deserialized view of it."""
+        Marshal.to_mem(self._msg_id, msg, self._slot)
+        self._msg_id += 1
+        with Marshal.obj_from_mem(self._slot) as received:
+            yield received
+
+
+def _detach(result):
+    """Snapshot a result so later slot reuse cannot change it out from under us.
+
+    The output of a transformer may itself alias the message it came from, which
+    is legitimate -- it is handed straight downstream and not retained -- but it
+    means the collected outputs have to be copied before the next publish.
+    """
+    if result is None:
+        return None
+    return replace(result, data=np.array(result.data))
+
+
+def run_recycled(proc, messages: typing.Sequence[AxisArray], *, slot_bytes: int = SLOT_BYTES) -> list:
+    """Push ``messages`` through ``proc`` with every one aliasing a reused slot."""
+    slot = RecycledSlot(slot_bytes)
+    outputs = []
+    for msg in messages:
+        with slot.publish(msg) as received:
+            outputs.append(_detach(proc(received)))
+    return outputs
+
+
+def run_owned(proc, messages: typing.Sequence[AxisArray]) -> list:
+    """Push ``messages`` through ``proc`` as ordinary, independently-owned arrays.
+
+    Each message gets a fresh copy of its data, so nothing the transformer keeps
+    can ever be invalidated -- the reference behaviour to compare against.
+    """
+    return [_detach(proc(replace(msg, data=np.array(msg.data)))) for msg in messages]
+
+
+def assert_survives_buffer_recycling(
+    make_proc: typing.Callable[[], typing.Any],
+    messages: typing.Sequence[AxisArray],
+    *,
+    slot_bytes: int = SLOT_BYTES,
+) -> list:
+    """Assert a transformer's output does not depend on who owns the input bytes.
+
+    Runs the same ``messages`` through two fresh transformers -- one on owned
+    arrays, one on arrays aliasing a single recycled slot -- and requires the
+    outputs to match exactly. They are the same arithmetic on the same numbers,
+    so any difference at all means state was read back from recycled memory.
+
+    :param make_proc: Zero-argument factory; called once per run so the two runs
+        do not share state.
+    :param messages: Inputs, in order. Equal-sized messages exercise the failure
+        hardest (see the module docstring).
+    :return: The outputs of the owned run, for further assertions.
+    """
+    owned = run_owned(make_proc(), messages)
+    recycled = run_recycled(make_proc(), messages, slot_bytes=slot_bytes)
+
+    assert len(owned) == len(recycled)
+    for ix, (exp, got) in enumerate(zip(owned, recycled)):
+        if exp is None or got is None:
+            assert exp is got, f"message {ix}: one run returned None and the other did not"
+            continue
+        assert exp.dims == got.dims, f"message {ix}: dims {got.dims} != {exp.dims}"
+        assert exp.data.shape == got.data.shape, f"message {ix}: shape {got.data.shape} != {exp.data.shape}"
+        assert np.array_equal(exp.data, got.data), (
+            f"message {ix}: output differs when the input buffer is recycled -- "
+            f"the transformer is retaining message.data (or a view of it) in its state.\n"
+            f"  owned:    {np.asarray(exp.data).ravel()[:8]}\n"
+            f"  recycled: {np.asarray(got.data).ravel()[:8]}"
+        )
+        for name, axis in exp.axes.items():
+            assert name in got.axes, f"message {ix}: missing axis {name!r}"
+            assert axis == got.axes[name], f"message {ix}: axis {name!r} differs"
+    return owned

tests/unit/test_buffer_recycling.py

@@ -0,0 +1,192 @@
+"""Transformers must not read back a message buffer they no longer own.
+
+See :mod:`tests.helpers.recycled_shm` for why this is invisible to every other
+test in the suite: in-process links never serialize, so an aliased array and an
+owned one behave identically no matter what a transformer retains.
+"""
+
+import ezmsg.core as ez
+import numpy as np
+import pytest
+from ezmsg.baseproc import BaseStatefulTransformer, processor_state
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
+from frozendict import frozendict
+
+from ezmsg.sigproc.aggregate import AggregationFunction
+from ezmsg.sigproc.binned_aggregate import BinnedAggregateSettings, BinnedAggregateTransformer
+from ezmsg.sigproc.diff import DiffSettings, DiffTransformer
+from tests.helpers.recycled_shm import (
+    RecycledSlot,
+    assert_survives_buffer_recycling,
+)
+
+FS = 100.0
+N_CH = 2
+
+
+def _msgs(blocks: list[np.ndarray], fs: float = FS) -> list[AxisArray]:
+    msgs, offset = [], 0.0
+    for blk in blocks:
+        msgs.append(
+            AxisArray(
+                data=blk,
+                dims=["time", "ch"],
+                axes=frozendict(
+                    {
+                        "time": AxisArray.TimeAxis(fs=fs, offset=offset),
+                        "ch": AxisArray.CoordinateAxis(data=np.arange(blk.shape[1]).astype(str), dims=["ch"]),
+                    }
+                ),
+                key="test_buffer_recycling",
+            )
+        )
+        offset += blk.shape[0] / fs
+    return msgs
+
+
+def _equal_blocks(n_msgs: int, n_time: int, seed: int = 0) -> list[np.ndarray]:
+    """Equal-sized random blocks: the next message lands on the exact bytes a
+    retained view points at, so corruption is silent rather than shape-obvious."""
+    rng = np.random.default_rng(seed)
+    return [rng.standard_normal((n_time, N_CH)) for _ in range(n_msgs)]
+
+
+# -- the harness itself ------------------------------------------------------
+
+
+class _RetainSettings(ez.Settings):
+    pass
+
+
+@processor_state
+class _RetainState:
+    last: object = None
+
+
+class _RetainingTransformer(BaseStatefulTransformer[_RetainSettings, AxisArray, AxisArray, _RetainState]):
+    """Canary: deliberately does the wrong thing, to prove the harness bites.
+
+    If this ever stops failing, the harness has gone blind and the real tests
+    below are passing for the wrong reason.
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash(message.key)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.last = None
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        prev = self._state.last
+        # Retain a view of the tail -- the bug this whole module is about.
+        self._state.last = slice_along_axis(message.data, slice(-1, None), axis=0)
+        if prev is None:
+            prev = np.zeros_like(self._state.last)
+        return replace(message, data=np.concatenate((prev, message.data), axis=0)[:-1])
+
+
+def test_harness_detects_a_retained_view():
+    msgs = _msgs(_equal_blocks(4, 12, seed=7))
+    with pytest.raises(AssertionError, match="retaining message.data"):
+        assert_survives_buffer_recycling(lambda: _RetainingTransformer(_RetainSettings()), msgs)
+
+
+def test_recycled_slot_actually_recycles():
+    """The slot must alias, and be overwritten by the next publish."""
+    msgs = _msgs(_equal_blocks(2, 8, seed=3))
+    slot = RecycledSlot()
+    with slot.publish(msgs[0]) as first:
+        assert first.data.base is not None, "received array should be a view, not an owner"
+        retained = first.data
+        assert np.array_equal(retained, msgs[0].data)
+    with slot.publish(msgs[1]) as _second:
+        assert not np.array_equal(retained, msgs[0].data), "slot was not reused; harness would not catch anything"
+
+
+# -- diff --------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("scale_by_fs", [False, True])
+def test_diff_does_not_retain_message_data(scale_by_fs):
+    """`last_dat` is the previous message's final sample, used on the next call."""
+    msgs = _msgs(_equal_blocks(4, 4, seed=0))
+    assert_survives_buffer_recycling(
+        lambda: DiffTransformer(DiffSettings(axis="time", scale_by_fs=scale_by_fs)),
+        msgs,
+    )
+
+
+def test_diff_boundary_sample_is_correct_across_recycling():
+    """Pin the actual value of the cross-message diff, not just self-consistency."""
+    blocks = _equal_blocks(2, 4, seed=0)
+    msgs = _msgs(blocks)
+    proc = DiffTransformer(DiffSettings(axis="time"))
+    slot = RecycledSlot()
+    outs = []
+    for msg in msgs:
+        with slot.publish(msg) as received:
+            outs.append(np.array(proc(received).data))
+    got = np.concatenate(outs, axis=0)
+
+    a, b = blocks
+    expected_boundary = b[0] - a[-1]
+    assert np.allclose(got[a.shape[0]], expected_boundary)
+
+
+# -- binned_aggregate --------------------------------------------------------
+
+
+def test_binned_aggregate_no_bin_completed_does_not_retain():
+    """`carry is None` and no bin completes: the whole message became the carry."""
+    rng = np.random.default_rng(0)
+    # 0.1 s bins at 100 Hz = 10 samples; a 4-sample chunk completes no bin.
+    blocks = [rng.standard_normal((4, N_CH)), rng.standard_normal((8, N_CH)), rng.standard_normal((8, N_CH))]
+    assert_survives_buffer_recycling(
+        lambda: BinnedAggregateTransformer(BinnedAggregateSettings(axis="time", bin_duration=0.1)),
+        _msgs(blocks),
+    )
+
+
+def test_binned_aggregate_leftover_tail_does_not_retain():
+    """A bin completes with no prior carry, so the leftover tail was a view of
+    `message.data`. Equal-size chunks put the next message on those bytes."""
+    assert_survives_buffer_recycling(
+        lambda: BinnedAggregateTransformer(BinnedAggregateSettings(axis="time", bin_duration=0.1)),
+        _msgs(_equal_blocks(4, 12, seed=1)),
+    )
+
+
+@pytest.mark.parametrize("fractional", [True, False])
+def test_binned_aggregate_fractional_grid_does_not_retain(fractional):
+    """An off-nominal rate makes bin lengths vary, so the carry length varies too.
+
+    Blocks are longer than a bin (101.3 samples at 1013 Hz) so the very first
+    message closes a bin with no prior carry -- the branch that leaves the carry
+    as a view of the message.
+    """
+    assert_survives_buffer_recycling(
+        lambda: BinnedAggregateTransformer(
+            BinnedAggregateSettings(axis="time", bin_duration=0.1, fractional=fractional)
+        ),
+        _msgs(_equal_blocks(6, 128, seed=2), fs=1013.0),
+    )
+
+
+def test_binned_aggregate_multi_op_does_not_retain():
+    """As above, through the stacked-operation path.
+
+    MIN/MAX only differ if the substituted samples are the extreme of their bin,
+    so the seed is one where they are -- corruption of a two-sample carry is
+    otherwise easy to average away.
+    """
+    assert_survives_buffer_recycling(
+        lambda: BinnedAggregateTransformer(
+            BinnedAggregateSettings(
+                axis="time",
+                bin_duration=0.1,
+                operation=(AggregationFunction.MIN, AggregationFunction.MAX),
+            )
+        ),
+        _msgs(_equal_blocks(4, 12, seed=0)),
+    )