Carry each frame's timestamp across the shmem ring

src/ezmsg/tools/plot/shmem_sweep.py

@@ -208,7 +208,7 @@ def _close_figure(self) -> None:
                 logger.exception("closing the sweep figure raised; continuing teardown")
 
     def _on_tick(self) -> None:
-        samples, _overflow = self._mirror.auto_view()
+        samples, times, _overflow = self._mirror.timed_view()
 
         shape = describe_mirror(self._mirror, label_fields=self._label_fields)
         if shape is None or shape.srate <= 0:
@@ -232,7 +232,9 @@ def _on_tick(self) -> None:
         self._apply_shape(shape)
 
         if samples is not None and samples.size:
-            self._sweep.push_data(np.ascontiguousarray(flatten_for_plot(samples, shape), dtype=np.float32))
+            self._sweep.push_data(
+                np.ascontiguousarray(flatten_for_plot(samples, shape), dtype=np.float32), timestamps=times
+            )
 
         self.on_frame(shape)
 

src/ezmsg/tools/shmem/shmem.py

@@ -20,6 +20,10 @@
 The other half must monitor the metadata shared memory to see if it changes, and if it does then it must recreate
 the data shared memory buffer reader at the new location.
 
+The data segment holds two arrays: first one float64 timestamp per frame (the buffered axis's value for that frame),
+then the frames themselves. Both are written before `write_index` advances, so a reader that takes frames from the
+ring takes their times from the same indices and the two always agree. See `ring_views`.
+
 Finally, there is a third piece of shared memory carrying everything about the AxisArray that does not fit in the
 fixed-size metadata header: the non-buffered coordinate axes (e.g. a `ch` axis holding per-channel bank/elec/label),
 axis units, and the message `attrs`. It lives at shorten_shmem_name(f"{shmem_name}/meta{meta_generation}") and is
@@ -83,15 +87,18 @@ def shorten_shmem_name(long_name: str) -> str:
 # from a header written by a build old enough to predate this check.
 SHMEM_META_MAGIC = 0x455A4D53
 
-# Bumped on any change to ShmemArrMeta._fields_ or to the .aux_meta wire format.
+# Bumped on any change to ShmemArrMeta._fields_, to the data segment's layout
+# (see ring_views), or to the .aux_meta wire format.
+#
+# 2: the data segment leads with one float64 timestamp per frame.
 #
 # The two halves of a shmem link must be the same version -- there is no
 # compatibility shim, by choice: the layouts are an internal detail between two
 # processes we deploy together, and carrying forward every past field shape would
 # cost more than it is worth. What we do owe is a loud failure rather than a
 # quiet one, so the reader validates the magic and version up front and raises
 # instead of misreading a header it does not understand.
-SHMEM_META_STRUCT_VERSION = 1
+SHMEM_META_STRUCT_VERSION = 2
 
 
 class ShmemVersionError(RuntimeError):
@@ -153,6 +160,23 @@ def key(self, value: str) -> None:
         ctypes.memmove(self._key_bytes, key_bytes[: self._key_len], self._key_len)
 
 
+def ring_views(buf: memoryview, meta: ShmemArrMeta) -> typing.Tuple[npt.NDArray, npt.NDArray]:
+    """The data segment's two arrays: per-frame timestamps (float64), then the frames.
+
+    Times lead so both arrays are aligned for any frame dtype. Writer and reader
+    both call this, so the layout is stated once.
+    """
+    n_frames = meta.shape[0]
+    times = np.ndarray((n_frames,), dtype=np.float64, buffer=buf[: n_frames * 8])
+    frames = np.ndarray(meta.shape[: meta.ndim], dtype=np.dtype(meta.dtype), buffer=buf[n_frames * 8 :])
+    return times, frames
+
+
+def ring_nbytes(n_frames: int, frame_shape: typing.Tuple[int, ...], itemsize: int) -> int:
+    """Size of the data segment described by :func:`ring_views`."""
+    return int(n_frames * 8 + n_frames * np.prod(frame_shape) * itemsize)
+
+
 class ShMemCircBuffSettings(ez.Settings):
     shmem_name: typing.Optional[str]
     buf_dur: float
@@ -180,6 +204,8 @@ class ShMemCircBuffState(ez.State):
     meta_struct: typing.Optional[ShmemArrMeta] = None
     buffer_shmem: typing.Optional[SharedMemory] = None
     buffer_arr: typing.Optional[npt.NDArray] = None
+    # One timestamp per frame of buffer_arr, at the same ring indices.
+    times_arr: typing.Optional[npt.NDArray] = None
     meta_hash: int = -1
     # The dimension currently buffered along; see ShMemCircBuffSettings.axis.
     buff_axis: typing.Optional[str] = None
@@ -331,6 +357,7 @@ def _cleanup_buffer(self):
 
         # Destroy the buffer
         self.STATE.buffer_arr = None
+        self.STATE.times_arr = None
 
         # Destroy the shared memory object
         if self.STATE.buffer_shmem is not None:
@@ -547,7 +574,7 @@ def _reset_buffer(self, msg: AxisArray) -> None:
         """
         self.STATE.meta_struct.buffer_generation += 1
         msg_dtype, msg_srate, n_frames, frame_shape = self._get_msg_meta(msg)
-        buff_size = int(n_frames * np.prod(frame_shape) * msg.data.itemsize)
+        buff_size = ring_nbytes(n_frames, frame_shape, msg.data.itemsize)
         buff_shm_name = self.SETTINGS.shmem_name + "/buffer" + str(self.STATE.meta_struct.buffer_generation)
         short_name = shorten_shmem_name(buff_shm_name)
         self.STATE.buffer_shmem = _persist_create_shmem(
@@ -556,11 +583,7 @@ def _reset_buffer(self, msg: AxisArray) -> None:
             purpose=f"data ring gen {self.STATE.meta_struct.buffer_generation} "
             f"({'x'.join(str(d) for d in (n_frames,) + frame_shape)})",
         )
-        self.STATE.buffer_arr = np.ndarray(
-            self.STATE.meta_struct.shape[: self.STATE.meta_struct.ndim],
-            dtype=np.dtype(self.STATE.meta_struct.dtype.decode("utf8")),
-            buffer=self.STATE.buffer_shmem.buf[:],
-        )
+        self.STATE.times_arr, self.STATE.buffer_arr = ring_views(self.STATE.buffer_shmem.buf, self.STATE.meta_struct)
         self.STATE.meta_struct.write_index = 0
         self.STATE.meta_struct.wrap_counter = 0
         self.STATE.meta_struct.bvalid = True
@@ -612,14 +635,23 @@ async def on_message(self, msg: AxisArray):
         self._update_aux_if_needed(msg)
 
         n_samples = data.shape[0]
-        write_stop = self.STATE.meta_struct.write_index + n_samples
+        axis = msg.axes[buff_axis]
+        times = axis.data if hasattr(axis, "data") else axis.offset + np.arange(n_samples) * axis.gain
+        write_index = self.STATE.meta_struct.write_index
+        write_stop = write_index + n_samples
 
+        # Frames and times both land before write_index moves, so a reader
+        # never sees one without the other.
         if write_stop > self.STATE.buffer_arr.shape[0]:
             overflow = write_stop - self.STATE.buffer_arr.shape[0]
-            self.STATE.buffer_arr[self.STATE.meta_struct.write_index :] = data[: n_samples - overflow]
-            self.STATE.buffer_arr[:overflow] = data[n_samples - overflow :]
+            split = n_samples - overflow
+            self.STATE.buffer_arr[write_index:] = data[:split]
+            self.STATE.buffer_arr[:overflow] = data[split:]
+            self.STATE.times_arr[write_index:] = times[:split]
+            self.STATE.times_arr[:overflow] = times[split:]
             self.STATE.meta_struct.write_index = overflow
             self.STATE.meta_struct.wrap_counter += 1
         else:
-            self.STATE.buffer_arr[self.STATE.meta_struct.write_index : write_stop] = data[:]
+            self.STATE.buffer_arr[write_index:write_stop] = data[:]
+            self.STATE.times_arr[write_index:write_stop] = times
             self.STATE.meta_struct.write_index = write_stop

src/ezmsg/tools/shmem/shmem_mirror.py

@@ -24,6 +24,7 @@
     ShmemArrMeta,
     ShMemCircBuffState,
     ShmemVersionError,
+    ring_views,
     shorten_shmem_name,
 )
 
@@ -231,6 +232,7 @@ def _cleanup_buffer(self):
         if self._mirror_state.buffer_arr is not None:
             del self._mirror_state.buffer_arr
         self._mirror_state.buffer_arr = None
+        self._mirror_state.times_arr = None
 
         if self._mirror_state.buffer_shmem is not None:
             # Note: Uncommenting the following does not eliminate the resource_tracker warnings.
@@ -308,21 +310,21 @@ def _reset_buffer(self) -> bool:
             buff_name = self._shmem_name + "/buffer" + str(self._mirror_state.meta_struct.buffer_generation)
             short_name = shorten_shmem_name(buff_name)
             self._mirror_state.buffer_shmem = SharedMemory(short_name, create=False)
-            self._mirror_state.buffer_arr = np.ndarray(
-                self._mirror_state.meta_struct.shape[: self._mirror_state.meta_struct.ndim],
-                dtype=np.dtype(self._mirror_state.meta_struct.dtype),
-                buffer=self._mirror_state.buffer_shmem.buf[:],
+            self._mirror_state.times_arr, self._mirror_state.buffer_arr = ring_views(
+                self._mirror_state.buffer_shmem.buf, self._mirror_state.meta_struct
             )
             self._last_meta = self.meta  # Copy
             if self._change_callback is not None:
                 self._change_callback()
             return True
         except FileNotFoundError:
             self._mirror_state.buffer_arr = None
+            self._mirror_state.times_arr = None
             self._mirror_state.buffer_shmem = None
         except TypeError:
             # buffer is too small for requested array
             self._mirror_state.buffer_arr = None
+            self._mirror_state.times_arr = None
             self._mirror_state.buffer_shmem = None
             print("DEBUG!")
         return False
@@ -351,6 +353,17 @@ def connect(self, name: str) -> None:
         self._last_connect_try = time.time()
 
     def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, bool]:
+        """The frames written since the last read, and whether the ring overran them."""
+        frames, _times, overflow = self.timed_view(n)
+        return frames, overflow
+
+    def timed_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, npt.NDArray, bool]:
+        """Like :meth:`auto_view`, plus each returned frame's timestamp.
+
+        The times are the buffered axis's values for those frames, as the
+        writer received them: the time axis's clock for a time-buffered stream.
+        """
+        no_times = np.empty(0, dtype=np.float64)
         if self._mirror_state.meta_struct is None:
             self.connect(self._shmem_name)
 
@@ -361,7 +374,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray,
         if self._mirror_state.meta_struct is None or not self._mirror_state.meta_struct.bvalid:
             # Still not connected
             #  or we are connected but the buffer data is invalid.
-            return np.array([[]]), False
+            return np.array([[]]), no_times, False
 
         b_connected = True
         # Determine if we need to reset the buffer
@@ -374,7 +387,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray,
 
         if not b_connected:
             # We STILL aren't connected.
-            return np.array([[]]), False
+            return np.array([[]]), no_times, False
 
         # -- From here, we should know we have a good connection to a valid buffer -- #
 
@@ -405,7 +418,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray,
         if n_available <= 1 or (n is not None and n_available < n):
             # Not enough samples available.
             # Return a null-slice of the buffer. This provides correct dimensions.
-            return self._mirror_state.buffer_arr[:0], b_overflow
+            return self._mirror_state.buffer_arr[:0], no_times, b_overflow
 
         # We have enough samples.
         if n is None:
@@ -415,6 +428,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray,
             # Return a contiguous chunk
             t_slice = np.s_[max(0, self._read_index) : self._read_index + n]
             result = self._mirror_state.buffer_arr[t_slice, :]
+            times = self._mirror_state.times_arr[t_slice].copy()
         else:
             # Split read into two chunks
             n_after_wrap = n - (self._mirror_state.meta_struct.shape[0] - self._read_index)
@@ -425,8 +439,11 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray,
                 ),
                 axis=0,
             )
+            times = np.concatenate(
+                (self._mirror_state.times_arr[self._read_index :], self._mirror_state.times_arr[:n_after_wrap])
+            )
 
         self._read_index = (self._read_index + n) % self._mirror_state.meta_struct.shape[0]
         self._last_meta.wrap_counter = self._mirror_state.meta_struct.wrap_counter
 
-        return result, b_overflow
+        return result, times, b_overflow

tests/test_shmem_mirror.py

@@ -182,3 +182,40 @@ def test_shmem_mirror_switch_buffer():
         print(f"\tFrom {dt0} to {switch_dt} after {np.sum(chunk_lens[:switch_idx])} samples")
     assert dt0 == np.float64
     assert switch_dt == np.float16
+
+
+def test_each_frame_carries_its_timestamp_across_a_wrap():
+    """The reader gets each frame's time from the same ring slots as the frame.
+
+    Driven in-process (no ez.run) so it is deterministic: 100 Hz into a 1 s
+    ring, in chunks that straddle the wrap, read back after each push.
+    """
+    import asyncio
+
+    from ezmsg.tools.shmem.shmem import ShMemCircBuffSettings
+
+    name = f"times{os.getpid()}"
+    sink = ShMemCircBuff(ShMemCircBuffSettings(shmem_name=name, buf_dur=1.0, axis="time"))
+    sink.STATE = ShMemCircBuff.STATE()
+    sink._reset_meta()
+    mirror = EZShmMirror(name)
+    try:
+        got_frames, got_times = [], []
+        t0, n = 1000.0, 0
+        for chunk in (40, 40, 40, 40):  # 160 frames through a 100-frame ring
+            data = np.arange(n, n + chunk, dtype=np.float32)[:, None] * np.ones((1, 2), np.float32)
+            msg = AxisArray(data, dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=100.0, offset=t0 + n / 100)})
+            asyncio.run(sink.on_message(msg))
+            n += chunk
+            frames, times, overflow = mirror.timed_view()
+            assert not overflow
+            got_frames.append(frames[:, 0].copy())
+            got_times.append(times)
+        frames, times = np.concatenate(got_frames), np.concatenate(got_times)
+        # Whatever frames came back, each one's time is the time it was sent with.
+        assert len(frames) > 100
+        np.testing.assert_allclose(times, t0 + frames / 100)
+    finally:
+        mirror.disconnect()
+        sink._cleanup_buffer()
+        sink._cleanup_meta()