From 77909eb0109373f2b1eaabb6f12b22f1c02baeb1 Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Wed, 1 Jul 2026 15:49:31 -0400 Subject: [PATCH 1/7] fix for channel failure on shm resize race --- src/ezmsg/core/messagechannel.py | 99 ++++++++++++++++++++++++-------- tests/test_channel.py | 71 ++++++++++++++++++++++- 2 files changed, 145 insertions(+), 25 deletions(-) diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index 130895f1..c7382fa0 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -4,7 +4,7 @@ import logging from uuid import UUID -from contextlib import contextmanager, suppress +from contextlib import contextmanager from .shm import SHMContext from .messagemarshal import MessageMarshal @@ -26,6 +26,10 @@ logger = logging.getLogger("ezmsg") +class ChannelError(RuntimeError): + """Raised when a channel cannot safely process publisher traffic.""" + + class LeakyQueue(asyncio.Queue[typing.Tuple[UUID, int]]): """ An asyncio.Queue that drops oldest items when full. @@ -266,34 +270,38 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: channel_kind = ProfileChannelType.SHM shm_name = await read_str(reader) - if self.shm is not None and self.shm.name != shm_name: - shm_entries = self.cache.keys() - self.cache.clear() - self.shm.close() - await self.shm.wait_closed() - - try: - self.shm = await GraphService( - self._graph_address - ).attach_shm(shm_name) - except ValueError: - logger.info( - "Invalid SHM received from publisher; may be dead" - ) - raise - - for id in shm_entries: - self.cache.put_from_mem(self.shm[id % self.num_buffers]) - - assert self.shm is not None - assert MessageMarshal.msg_id(self.shm[buf_idx]) == msg_id - self.cache.put_from_mem(self.shm[buf_idx]) + if self.shm is None or self.shm.name != shm_name: + await self._reattach_shm(shm_name) + + if self.shm is None: + raise ChannelError( + f"channel {self.id} has no SHM attached for {shm_name}" + ) + + try: + shm_buf = self.shm[buf_idx] + except BufferError as exc: + raise ChannelError( + f"channel {self.id} lost SHM {shm_name} while reading {msg_id=}" + ) from exc + + if MessageMarshal.msg_id(shm_buf) != msg_id: + raise ChannelError( + f"channel {self.id} saw mismatched SHM contents in {shm_name}: " + f"expected {msg_id}, got {MessageMarshal.msg_id(shm_buf)}" + ) + + self.cache.put_from_mem(shm_buf) elif msg == Command.TX_TCP.value: channel_kind = ProfileChannelType.TCP buf_size = await read_int(reader) obj_bytes = await reader.readexactly(buf_size) - assert MessageMarshal.msg_id(obj_bytes) == msg_id + if MessageMarshal.msg_id(obj_bytes) != msg_id: + raise ChannelError( + f"channel {self.id} saw mismatched TCP contents: " + f"expected {msg_id}, got {MessageMarshal.msg_id(obj_bytes)}" + ) self.cache.put_from_mem(memoryview(obj_bytes).toreadonly()) else: @@ -308,6 +316,12 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError): logger.debug(f"connection fail: channel:{self.id} - pub:{self.pub_id}") + except (ChannelError, FileNotFoundError): + logger.exception( + "Publisher connection failed for channel %s from publisher %s", + self.id, + self.pub_id, + ) finally: self.cache.clear() @@ -318,6 +332,43 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: logger.debug(f"disconnected: channel:{self.id} -> pub:{self.pub_id}") + async def _reattach_shm(self, shm_name: str) -> None: + cached_msg_ids = self.cache.keys() + self.cache.clear() + + prior_shm = self.shm + self.shm = None + if prior_shm is not None: + prior_shm.close() + await prior_shm.wait_closed() + + try: + self.shm = await GraphService(self._graph_address).attach_shm(shm_name) + except (ValueError, FileNotFoundError) as exc: + raise ChannelError( + f"channel {self.id} failed to attach SHM {shm_name}" + ) from exc + + for cached_msg_id in cached_msg_ids: + cached_buf_idx = cached_msg_id % self.num_buffers + try: + cached_mem = self.shm[cached_buf_idx] + except BufferError as exc: + raise ChannelError( + f"channel {self.id} lost SHM {shm_name} while restoring cache" + ) from exc + + if MessageMarshal.msg_id(cached_mem) != cached_msg_id: + logger.warning( + "Dropping stale cached message %s during SHM switch to %s on channel %s", + cached_msg_id, + shm_name, + self.id, + ) + continue + + self.cache.put_from_mem(cached_mem) + def _set_channel_kind(self, kind: ProfileChannelType) -> None: if self._channel_kind == ProfileChannelType.UNKNOWN: self._channel_kind = kind diff --git a/tests/test_channel.py b/tests/test_channel.py index e45f7192..77c64fe9 100644 --- a/tests/test_channel.py +++ b/tests/test_channel.py @@ -1,12 +1,14 @@ import asyncio from uuid import uuid4 +from unittest.mock import patch import pytest -from ezmsg.core.messagechannel import Channel +from ezmsg.core.messagechannel import Channel, ChannelError from ezmsg.core.messagecache import CacheMiss from ezmsg.core.netprotocol import Command, uint64_to_bytes from ezmsg.core.backpressure import Backpressure +from ezmsg.core.messagemarshal import MessageMarshal class DummyWriter: @@ -91,3 +93,70 @@ def test_channel_put_local_requires_local_backpressure(): channel = Channel(uuid4(), uuid4(), 1, None, None, Channel._SENTINEL) with pytest.raises(ValueError): channel.put_local(1, "no pub") + + +class FakeSHM: + def __init__(self, name: str, buffers: dict[int, memoryview] | None = None): + self.name = name + self._buffers = buffers or {} + self.closed = False + self.waited = False + + def __getitem__(self, idx: int) -> memoryview: + return self._buffers[idx] + + def close(self) -> None: + self.closed = True + + async def wait_closed(self) -> None: + self.waited = True + + +def _marshal_message(msg_id: int, obj: object) -> memoryview: + with MessageMarshal.serialize(msg_id, obj) as (size, header, buffers): + raw = bytearray(size + 1) + mem = memoryview(raw) + MessageMarshal._write(mem, header, buffers) + return mem.toreadonly() + + +@pytest.mark.asyncio +async def test_channel_reattach_shm_drops_stale_cached_messages(): + old_shm = FakeSHM("old") + channel = Channel(uuid4(), uuid4(), 2, old_shm, ("127.0.0.1", 0), Channel._SENTINEL) + + channel.cache.put_local("cached", 1) + new_shm = FakeSHM("new", {1: _marshal_message(3, "fresh")}) + + class FakeGraphService: + def __init__(self, address): + self.address = address + + async def attach_shm(self, shm_name: str): + assert shm_name == "new" + return new_shm + + with patch("ezmsg.core.messagechannel.GraphService", FakeGraphService): + await channel._reattach_shm("new") + + assert channel.shm is new_shm + assert old_shm.closed is True + assert old_shm.waited is True + with pytest.raises(CacheMiss): + _ = channel.cache[1] + + +@pytest.mark.asyncio +async def test_channel_reattach_shm_wraps_missing_segment(): + channel = Channel(uuid4(), uuid4(), 2, None, ("127.0.0.1", 0), Channel._SENTINEL) + + class FakeGraphService: + def __init__(self, address): + self.address = address + + async def attach_shm(self, shm_name: str): + raise FileNotFoundError(shm_name) + + with patch("ezmsg.core.messagechannel.GraphService", FakeGraphService): + with pytest.raises(ChannelError): + await channel._reattach_shm("missing") From 1ff1185eb1b6466666ac3f1dfdf3b5e53ad3e1e8 Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Wed, 1 Jul 2026 16:47:16 -0400 Subject: [PATCH 2/7] added shm race failure to test harness --- tests/shm_resize_race_runner.py | 73 +++++++++++++++++++++++++++++++++ tests/test_clean_shutdown.py | 18 ++++++++ 2 files changed, 91 insertions(+) create mode 100644 tests/shm_resize_race_runner.py diff --git a/tests/shm_resize_race_runner.py b/tests/shm_resize_race_runner.py new file mode 100644 index 00000000..1106f789 --- /dev/null +++ b/tests/shm_resize_race_runner.py @@ -0,0 +1,73 @@ +import asyncio +import os +import time +from dataclasses import dataclass + +import ezmsg.core as ez + +INITIAL_SHM_SIZE = 64 +NUM_MSGS = 4 +SUBSCRIBER_DELAY_S = 0.001 +READY_TOKEN = "READY" +DONE_TOKEN = "DONE" + + +@dataclass +class BurstMessage: + seq: int + payload: bytes + created_at: float + + +class BurstyPublisher(ez.Unit): + OUTPUT = ez.OutputStream( + BurstMessage, + num_buffers=2, + buf_size=INITIAL_SHM_SIZE, + force_tcp=False, + allow_local=False, + ) + + @ez.publisher(OUTPUT) + async def pump(self): + cur_size = INITIAL_SHM_SIZE + print(READY_TOKEN, flush=True) + for itr in range(NUM_MSGS): + cur_size *= 3 + yield self.OUTPUT, BurstMessage( + seq=itr, + payload=bytes(cur_size), + created_at=time.time(), + ) + + +class SubscriberState(ez.State): + cur_msg: int = 0 + + +class Subscriber(ez.Unit): + INPUT = ez.InputStream(BurstMessage) + STATE = SubscriberState + + @ez.subscriber(INPUT) + async def on_message(self, msg: BurstMessage) -> None: + await asyncio.sleep(SUBSCRIBER_DELAY_S) + self.STATE.cur_msg += 1 + if self.STATE.cur_msg == NUM_MSGS: + print(DONE_TOKEN, flush=True) + raise ez.NormalTermination + + +class ReproSystem(ez.Collection): + PUB = BurstyPublisher() + SUB = Subscriber() + + def network(self) -> ez.NetworkDefinition: + return ((self.PUB.OUTPUT, self.SUB.INPUT),) + + def process_components(self) -> list[ez.Component]: + return [self.PUB, self.SUB] + + +if __name__ == "__main__": + ez.run(SYSTEM=ReproSystem()) diff --git a/tests/test_clean_shutdown.py b/tests/test_clean_shutdown.py index cec2df90..566319e0 100644 --- a/tests/test_clean_shutdown.py +++ b/tests/test_clean_shutdown.py @@ -12,6 +12,7 @@ ROOT = Path(__file__).resolve().parents[1] RUNNER = Path(__file__).with_name("shutdown_runner.py") EXAMPLE_RUNNER = Path(__file__).with_name("clean_shutdown_examples_runner.py") +SHM_RACE_RUNNER = Path(__file__).with_name("shm_resize_race_runner.py") def _run_process( @@ -187,6 +188,19 @@ def _run_example_case( ) +def _run_shm_resize_race_case(*, timeout: float = 1.0) -> None: + env = os.environ.copy() + env.pop("EZMSG_STRICT_SHUTDOWN", None) + _run_process( + [sys.executable, "-u", str(SHM_RACE_RUNNER)], + env=env, + signals=0, + ready_token="READY", + timeout=timeout, + allowed_returncodes={0}, + ) + + def _sigint_returncodes() -> set[int]: if os.name == "nt": return {1, 3221225786} @@ -223,3 +237,7 @@ def test_infinite_requires_sigint(start_method: str) -> None: signals=1, allowed_returncodes={0}, ) + + +def test_shm_resize_race_repro_completes() -> None: + _run_shm_resize_race_case(timeout=10.0) From 53222320c05575c8f7b705ccae120b43d0520072 Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Wed, 1 Jul 2026 16:48:32 -0400 Subject: [PATCH 3/7] Revert "fix for channel failure on shm resize race" This reverts commit 77909eb0109373f2b1eaabb6f12b22f1c02baeb1. --- src/ezmsg/core/messagechannel.py | 99 ++++++++------------------------ tests/test_channel.py | 71 +---------------------- 2 files changed, 25 insertions(+), 145 deletions(-) diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index c7382fa0..130895f1 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -4,7 +4,7 @@ import logging from uuid import UUID -from contextlib import contextmanager +from contextlib import contextmanager, suppress from .shm import SHMContext from .messagemarshal import MessageMarshal @@ -26,10 +26,6 @@ logger = logging.getLogger("ezmsg") -class ChannelError(RuntimeError): - """Raised when a channel cannot safely process publisher traffic.""" - - class LeakyQueue(asyncio.Queue[typing.Tuple[UUID, int]]): """ An asyncio.Queue that drops oldest items when full. @@ -270,38 +266,34 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: channel_kind = ProfileChannelType.SHM shm_name = await read_str(reader) - if self.shm is None or self.shm.name != shm_name: - await self._reattach_shm(shm_name) - - if self.shm is None: - raise ChannelError( - f"channel {self.id} has no SHM attached for {shm_name}" - ) - - try: - shm_buf = self.shm[buf_idx] - except BufferError as exc: - raise ChannelError( - f"channel {self.id} lost SHM {shm_name} while reading {msg_id=}" - ) from exc - - if MessageMarshal.msg_id(shm_buf) != msg_id: - raise ChannelError( - f"channel {self.id} saw mismatched SHM contents in {shm_name}: " - f"expected {msg_id}, got {MessageMarshal.msg_id(shm_buf)}" - ) - - self.cache.put_from_mem(shm_buf) + if self.shm is not None and self.shm.name != shm_name: + shm_entries = self.cache.keys() + self.cache.clear() + self.shm.close() + await self.shm.wait_closed() + + try: + self.shm = await GraphService( + self._graph_address + ).attach_shm(shm_name) + except ValueError: + logger.info( + "Invalid SHM received from publisher; may be dead" + ) + raise + + for id in shm_entries: + self.cache.put_from_mem(self.shm[id % self.num_buffers]) + + assert self.shm is not None + assert MessageMarshal.msg_id(self.shm[buf_idx]) == msg_id + self.cache.put_from_mem(self.shm[buf_idx]) elif msg == Command.TX_TCP.value: channel_kind = ProfileChannelType.TCP buf_size = await read_int(reader) obj_bytes = await reader.readexactly(buf_size) - if MessageMarshal.msg_id(obj_bytes) != msg_id: - raise ChannelError( - f"channel {self.id} saw mismatched TCP contents: " - f"expected {msg_id}, got {MessageMarshal.msg_id(obj_bytes)}" - ) + assert MessageMarshal.msg_id(obj_bytes) == msg_id self.cache.put_from_mem(memoryview(obj_bytes).toreadonly()) else: @@ -316,12 +308,6 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError): logger.debug(f"connection fail: channel:{self.id} - pub:{self.pub_id}") - except (ChannelError, FileNotFoundError): - logger.exception( - "Publisher connection failed for channel %s from publisher %s", - self.id, - self.pub_id, - ) finally: self.cache.clear() @@ -332,43 +318,6 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: logger.debug(f"disconnected: channel:{self.id} -> pub:{self.pub_id}") - async def _reattach_shm(self, shm_name: str) -> None: - cached_msg_ids = self.cache.keys() - self.cache.clear() - - prior_shm = self.shm - self.shm = None - if prior_shm is not None: - prior_shm.close() - await prior_shm.wait_closed() - - try: - self.shm = await GraphService(self._graph_address).attach_shm(shm_name) - except (ValueError, FileNotFoundError) as exc: - raise ChannelError( - f"channel {self.id} failed to attach SHM {shm_name}" - ) from exc - - for cached_msg_id in cached_msg_ids: - cached_buf_idx = cached_msg_id % self.num_buffers - try: - cached_mem = self.shm[cached_buf_idx] - except BufferError as exc: - raise ChannelError( - f"channel {self.id} lost SHM {shm_name} while restoring cache" - ) from exc - - if MessageMarshal.msg_id(cached_mem) != cached_msg_id: - logger.warning( - "Dropping stale cached message %s during SHM switch to %s on channel %s", - cached_msg_id, - shm_name, - self.id, - ) - continue - - self.cache.put_from_mem(cached_mem) - def _set_channel_kind(self, kind: ProfileChannelType) -> None: if self._channel_kind == ProfileChannelType.UNKNOWN: self._channel_kind = kind diff --git a/tests/test_channel.py b/tests/test_channel.py index 77c64fe9..e45f7192 100644 --- a/tests/test_channel.py +++ b/tests/test_channel.py @@ -1,14 +1,12 @@ import asyncio from uuid import uuid4 -from unittest.mock import patch import pytest -from ezmsg.core.messagechannel import Channel, ChannelError +from ezmsg.core.messagechannel import Channel from ezmsg.core.messagecache import CacheMiss from ezmsg.core.netprotocol import Command, uint64_to_bytes from ezmsg.core.backpressure import Backpressure -from ezmsg.core.messagemarshal import MessageMarshal class DummyWriter: @@ -93,70 +91,3 @@ def test_channel_put_local_requires_local_backpressure(): channel = Channel(uuid4(), uuid4(), 1, None, None, Channel._SENTINEL) with pytest.raises(ValueError): channel.put_local(1, "no pub") - - -class FakeSHM: - def __init__(self, name: str, buffers: dict[int, memoryview] | None = None): - self.name = name - self._buffers = buffers or {} - self.closed = False - self.waited = False - - def __getitem__(self, idx: int) -> memoryview: - return self._buffers[idx] - - def close(self) -> None: - self.closed = True - - async def wait_closed(self) -> None: - self.waited = True - - -def _marshal_message(msg_id: int, obj: object) -> memoryview: - with MessageMarshal.serialize(msg_id, obj) as (size, header, buffers): - raw = bytearray(size + 1) - mem = memoryview(raw) - MessageMarshal._write(mem, header, buffers) - return mem.toreadonly() - - -@pytest.mark.asyncio -async def test_channel_reattach_shm_drops_stale_cached_messages(): - old_shm = FakeSHM("old") - channel = Channel(uuid4(), uuid4(), 2, old_shm, ("127.0.0.1", 0), Channel._SENTINEL) - - channel.cache.put_local("cached", 1) - new_shm = FakeSHM("new", {1: _marshal_message(3, "fresh")}) - - class FakeGraphService: - def __init__(self, address): - self.address = address - - async def attach_shm(self, shm_name: str): - assert shm_name == "new" - return new_shm - - with patch("ezmsg.core.messagechannel.GraphService", FakeGraphService): - await channel._reattach_shm("new") - - assert channel.shm is new_shm - assert old_shm.closed is True - assert old_shm.waited is True - with pytest.raises(CacheMiss): - _ = channel.cache[1] - - -@pytest.mark.asyncio -async def test_channel_reattach_shm_wraps_missing_segment(): - channel = Channel(uuid4(), uuid4(), 2, None, ("127.0.0.1", 0), Channel._SENTINEL) - - class FakeGraphService: - def __init__(self, address): - self.address = address - - async def attach_shm(self, shm_name: str): - raise FileNotFoundError(shm_name) - - with patch("ezmsg.core.messagechannel.GraphService", FakeGraphService): - with pytest.raises(ChannelError): - await channel._reattach_shm("missing") From a4c5b8f5ac65d58e308e4c29f5ce612c6790afdd Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Wed, 1 Jul 2026 17:22:28 -0400 Subject: [PATCH 4/7] made shm churn less fragile --- src/ezmsg/core/messagechannel.py | 57 +++++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 12 deletions(-) diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index 130895f1..5dd7b991 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -9,7 +9,7 @@ from .shm import SHMContext from .messagemarshal import MessageMarshal from .backpressure import Backpressure -from .messagecache import MessageCache +from .messagecache import MessageCache, CacheMiss from .graphserver import GraphService from .netprotocol import ( Command, @@ -277,17 +277,42 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: self._graph_address ).attach_shm(shm_name) except ValueError: - logger.info( - "Invalid SHM received from publisher; may be dead" + logger.warning( + "Channel %s received stale SHM %s for publisher %s; waiting for next valid SHM", + self.id, + shm_name, + self.pub_id, ) - raise - - for id in shm_entries: - self.cache.put_from_mem(self.shm[id % self.num_buffers]) - - assert self.shm is not None - assert MessageMarshal.msg_id(self.shm[buf_idx]) == msg_id - self.cache.put_from_mem(self.shm[buf_idx]) + self.shm = None + + if self.shm is not None: + for id in shm_entries: + shm_buf = self.shm[id % self.num_buffers] + if MessageMarshal.msg_id(shm_buf) == id: + self.cache.put_from_mem(shm_buf) + + if self.shm is None: + logger.warning( + "Channel %s dropping message %s from publisher %s because its SHM generation is stale", + self.id, + msg_id, + self.pub_id, + ) + self._release_backpressure(msg_id, self.id) + continue + + shm_buf = self.shm[buf_idx] + if MessageMarshal.msg_id(shm_buf) != msg_id: + logger.warning( + "Channel %s skipping stale SHM contents for message %s from publisher %s; will use next valid SHM generation", + self.id, + msg_id, + self.pub_id, + ) + self._release_backpressure(msg_id, self.id) + continue + + self.cache.put_from_mem(shm_buf) elif msg == Command.TX_TCP.value: channel_kind = ProfileChannelType.TCP @@ -407,7 +432,15 @@ def _release_backpressure(self, msg_id: int, client_id: UUID) -> None: buf_idx = msg_id % self.num_buffers self.backpressure.free(client_id, buf_idx) if self.backpressure.buffers[buf_idx].is_empty: - self.cache.release(msg_id) + try: + self.cache.release(msg_id) + except CacheMiss: + logger.debug( + "Channel %s observed cache miss while releasing msg_id=%s from publisher %s; continuing backpressure release", + self.id, + msg_id, + self.pub_id, + ) # If pub is in same process as this channel, avoid TCP if self._local_backpressure is not None: From 5ae385c76fdc2ff3e9df5a2b5624e28ef2a2547d Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sun, 19 Jul 2026 20:11:48 -0400 Subject: [PATCH 5/7] Guard msg_id() against UninitializedMemory + add cross-process grow test On top of #254: the two MessageMarshal.msg_id() call sites in the resize path (repopulate loop and the buffer-id check) can hit a slot that is uninitialized after a mid-stream resize. msg_id() calls _assert_initialized() and raises UninitializedMemory there, which would escape the receive loop and kill the channel task -- the exact failure #254 is closing. Guard both sites: treat uninitialized as skip/mismatch (drop + release backpressure). Add tests/test_shm_grow.py: a parametrized real cross-process pub/sub with small buf_size and mid-stream oversized messages (single grow, two successive grows, grow-on-first-message), asserting every message is delivered. Complements the existing shm_resize_race_runner repro. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/ezmsg/core/messagechannel.py | 23 ++++- tests/test_shm_grow.py | 160 +++++++++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 4 deletions(-) create mode 100644 tests/test_shm_grow.py diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index 5dd7b991..93251f4f 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -7,7 +7,7 @@ from contextlib import contextmanager, suppress from .shm import SHMContext -from .messagemarshal import MessageMarshal +from .messagemarshal import MessageMarshal, UninitializedMemory from .backpressure import Backpressure from .messagecache import MessageCache, CacheMiss from .graphserver import GraphService @@ -288,8 +288,15 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: if self.shm is not None: for id in shm_entries: shm_buf = self.shm[id % self.num_buffers] - if MessageMarshal.msg_id(shm_buf) == id: - self.cache.put_from_mem(shm_buf) + # A carried-over slot may not have survived the + # grow/copy; msg_id() raises UninitializedMemory + # on an uninitialized slot. Skip it rather than + # let that escape and kill the channel task. + try: + if MessageMarshal.msg_id(shm_buf) == id: + self.cache.put_from_mem(shm_buf) + except UninitializedMemory: + pass if self.shm is None: logger.warning( @@ -302,7 +309,15 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: continue shm_buf = self.shm[buf_idx] - if MessageMarshal.msg_id(shm_buf) != msg_id: + # The slot for this msg_id may be uninitialized after a + # mid-stream resize; msg_id() raises UninitializedMemory in + # that case. Treat it as a mismatch (drop + release) instead + # of letting it escape and silently kill the channel task. + try: + slot_msg_id = MessageMarshal.msg_id(shm_buf) + except UninitializedMemory: + slot_msg_id = None + if slot_msg_id != msg_id: logger.warning( "Channel %s skipping stale SHM contents for message %s from publisher %s; will use next valid SHM generation", self.id, diff --git a/tests/test_shm_grow.py b/tests/test_shm_grow.py new file mode 100644 index 00000000..6216d79c --- /dev/null +++ b/tests/test_shm_grow.py @@ -0,0 +1,160 @@ +"""Cross-process SHM grow / reattach coverage. + +When a published message exceeds the publisher's shared-memory ``buf_size``, +``PublisherClient.broadcast`` grows the segment (allocate ``total_size*2``, copy +buffers, swap) and hands subscribers a new SHM name. Each subscriber must detach +from the old segment and reattach to the new one mid-stream +(``MessageChannel.monitor``). Historically that reattach path had bare +``assert``s that, on any transient skew, escaped the receive loop and silently +killed the channel task -> the subscriber stalled forever. + +These tests drive a real cross-process publisher/subscriber with a deliberately +small ``buf_size`` and a mid-stream oversized message, and assert every message +is still delivered (the grow does not wedge the stream). Small messages before +the grow populate the cache so the reattach repopulation loop is exercised too. +""" + +from dataclasses import dataclass +from collections.abc import AsyncGenerator + +import json + +import pytest + +import ezmsg.core as ez + +from ez_test_utils import get_test_fn + + +# A payload whose serialized size we can control via the bytes field. +@dataclass +class BlobMessage: + seq: int + payload: bytes + + +class BlobGeneratorSettings(ez.Settings): + sizes: tuple[int, ...] + """Per-message payload byte counts; a value above ``buf_size`` forces a grow.""" + + buf_size: int + num_buffers: int + + +class BlobGenerator(ez.Unit): + SETTINGS = BlobGeneratorSettings + + # buf_size is intentionally small so a large payload triggers a grow. + OUTPUT = ez.OutputStream( + BlobMessage, + num_buffers=4, + buf_size=4096, + allow_local=False, # keep cross-process on the SHM path, not local fast path + ) + + async def initialize(self) -> None: + # Apply the parametrized transport sizing on the stream instance. + self.OUTPUT.buf_size = self.SETTINGS.buf_size + self.OUTPUT.num_buffers = self.SETTINGS.num_buffers + + @ez.publisher(OUTPUT) + async def spawn(self) -> AsyncGenerator: + for seq, size in enumerate(self.SETTINGS.sizes): + yield self.OUTPUT, BlobMessage(seq=seq, payload=b"x" * size) + raise ez.Complete + + +class BlobReceiverSettings(ez.Settings): + num_msgs: int + output_fn: str + + +class BlobReceiverState(ez.State): + num_received: int = 0 + + +class BlobReceiver(ez.Unit): + STATE = BlobReceiverState + SETTINGS = BlobReceiverSettings + + INPUT = ez.InputStream(BlobMessage) + + @ez.subscriber(INPUT) + async def on_message(self, msg: BlobMessage) -> None: + self.STATE.num_received += 1 + with open(self.SETTINGS.output_fn, "a") as output_file: + output_file.write( + json.dumps({"seq": msg.seq, "len": len(msg.payload)}) + "\n" + ) + if self.STATE.num_received == self.SETTINGS.num_msgs: + raise ez.Complete + + +class GrowSystemSettings(ez.Settings): + sizes: tuple[int, ...] + buf_size: int + num_buffers: int + output_fn: str + + +class GrowSystem(ez.Collection): + SETTINGS = GrowSystemSettings + + PUB = BlobGenerator() + SUB = BlobReceiver() + + def configure(self) -> None: + self.PUB.apply_settings( + BlobGeneratorSettings( + sizes=self.SETTINGS.sizes, + buf_size=self.SETTINGS.buf_size, + num_buffers=self.SETTINGS.num_buffers, + ) + ) + self.SUB.apply_settings( + BlobReceiverSettings( + num_msgs=len(self.SETTINGS.sizes), + output_fn=self.SETTINGS.output_fn, + ) + ) + + def network(self) -> ez.NetworkDefinition: + return ((self.PUB.OUTPUT, self.SUB.INPUT),) + + def process_components(self): + # Force PUB and SUB into separate processes so the boundary uses the + # cross-process SHM transport (and thus the grow/reattach path). + return (self.PUB, self.SUB) + + +@pytest.mark.parametrize( + "sizes", + [ + # small, small, GROW (exceeds buf_size), small, small + (8, 8, 16384, 8, 8), + # two successive grows of increasing size -> two reattaches + (8, 16384, 8, 65536, 8), + # grow on the very first message (empty cache to repopulate) + (16384, 8, 8), + ], +) +def test_cross_process_grow_delivers_all(sizes): + with get_test_fn() as test_filename: + system = GrowSystem( + GrowSystemSettings( + sizes=sizes, + buf_size=4096, + num_buffers=4, + output_fn=str(test_filename), + ) + ) + ez.run(SYSTEM=system) + + results = [] + with open(test_filename, "r") as file: + for line in file: + results.append(json.loads(line)) + + # Every message survived the mid-stream grow: none dropped, in order. + assert [r["seq"] for r in results] == list(range(len(sizes))) + assert [r["len"] for r in results] == list(sizes) From d462dbec852416a86aa72d7d7ecb2315dc4d18a9 Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Mon, 10 Aug 2026 12:49:52 -0400 Subject: [PATCH 6/7] address pickling failure in test --- src/ezmsg/core/shm_grow_test_support.py | 100 ++++++++++++++++++ tests/test_shm_grow.py | 128 ++---------------------- 2 files changed, 109 insertions(+), 119 deletions(-) create mode 100644 src/ezmsg/core/shm_grow_test_support.py diff --git a/src/ezmsg/core/shm_grow_test_support.py b/src/ezmsg/core/shm_grow_test_support.py new file mode 100644 index 00000000..25c0b5f6 --- /dev/null +++ b/src/ezmsg/core/shm_grow_test_support.py @@ -0,0 +1,100 @@ +from collections.abc import AsyncGenerator +from dataclasses import dataclass + +import json + +import ezmsg.core as ez + + +@dataclass +class BlobMessage: + seq: int + payload: bytes + + +class BlobGeneratorSettings(ez.Settings): + sizes: tuple[int, ...] + buf_size: int + num_buffers: int + + +class BlobGenerator(ez.Unit): + SETTINGS = BlobGeneratorSettings + + OUTPUT = ez.OutputStream( + BlobMessage, + num_buffers=4, + buf_size=4096, + allow_local=False, + ) + + async def initialize(self) -> None: + self.OUTPUT.buf_size = self.SETTINGS.buf_size + self.OUTPUT.num_buffers = self.SETTINGS.num_buffers + + @ez.publisher(OUTPUT) + async def spawn(self) -> AsyncGenerator: + for seq, size in enumerate(self.SETTINGS.sizes): + yield self.OUTPUT, BlobMessage(seq=seq, payload=b"x" * size) + raise ez.Complete + + +class BlobReceiverSettings(ez.Settings): + num_msgs: int + output_fn: str + + +class BlobReceiverState(ez.State): + num_received: int = 0 + + +class BlobReceiver(ez.Unit): + STATE = BlobReceiverState + SETTINGS = BlobReceiverSettings + + INPUT = ez.InputStream(BlobMessage) + + @ez.subscriber(INPUT) + async def on_message(self, msg: BlobMessage) -> None: + self.STATE.num_received += 1 + with open(self.SETTINGS.output_fn, "a") as output_file: + output_file.write( + json.dumps({"seq": msg.seq, "len": len(msg.payload)}) + "\n" + ) + if self.STATE.num_received == self.SETTINGS.num_msgs: + raise ez.Complete + + +class GrowSystemSettings(ez.Settings): + sizes: tuple[int, ...] + buf_size: int + num_buffers: int + output_fn: str + + +class GrowSystem(ez.Collection): + SETTINGS = GrowSystemSettings + + PUB = BlobGenerator() + SUB = BlobReceiver() + + def configure(self) -> None: + self.PUB.apply_settings( + BlobGeneratorSettings( + sizes=self.SETTINGS.sizes, + buf_size=self.SETTINGS.buf_size, + num_buffers=self.SETTINGS.num_buffers, + ) + ) + self.SUB.apply_settings( + BlobReceiverSettings( + num_msgs=len(self.SETTINGS.sizes), + output_fn=self.SETTINGS.output_fn, + ) + ) + + def network(self) -> ez.NetworkDefinition: + return ((self.PUB.OUTPUT, self.SUB.INPUT),) + + def process_components(self): + return (self.PUB, self.SUB) diff --git a/tests/test_shm_grow.py b/tests/test_shm_grow.py index 6216d79c..f4d76e8a 100644 --- a/tests/test_shm_grow.py +++ b/tests/test_shm_grow.py @@ -1,22 +1,16 @@ """Cross-process SHM grow / reattach coverage. When a published message exceeds the publisher's shared-memory ``buf_size``, -``PublisherClient.broadcast`` grows the segment (allocate ``total_size*2``, copy -buffers, swap) and hands subscribers a new SHM name. Each subscriber must detach -from the old segment and reattach to the new one mid-stream -(``MessageChannel.monitor``). Historically that reattach path had bare -``assert``s that, on any transient skew, escaped the receive loop and silently -killed the channel task -> the subscriber stalled forever. - -These tests drive a real cross-process publisher/subscriber with a deliberately -small ``buf_size`` and a mid-stream oversized message, and assert every message -is still delivered (the grow does not wedge the stream). Small messages before -the grow populate the cache so the reattach repopulation loop is exercised too. +``Publisher.broadcast`` grows the segment (allocate ``total_size*2``, copy +buffers, swap) and hands subscribers a new SHM name. Each subscriber must +detach from the old segment and reattach to the new one mid-stream. These tests +assert that the subscriber stays alive and every message is still delivered in +order across those grows. + +The spawned publisher/subscriber units live in an importable package module so +multiprocessing ``spawn`` can unpickle them reliably under pytest. """ -from dataclasses import dataclass -from collections.abc import AsyncGenerator - import json import pytest @@ -24,117 +18,14 @@ import ezmsg.core as ez from ez_test_utils import get_test_fn - - -# A payload whose serialized size we can control via the bytes field. -@dataclass -class BlobMessage: - seq: int - payload: bytes - - -class BlobGeneratorSettings(ez.Settings): - sizes: tuple[int, ...] - """Per-message payload byte counts; a value above ``buf_size`` forces a grow.""" - - buf_size: int - num_buffers: int - - -class BlobGenerator(ez.Unit): - SETTINGS = BlobGeneratorSettings - - # buf_size is intentionally small so a large payload triggers a grow. - OUTPUT = ez.OutputStream( - BlobMessage, - num_buffers=4, - buf_size=4096, - allow_local=False, # keep cross-process on the SHM path, not local fast path - ) - - async def initialize(self) -> None: - # Apply the parametrized transport sizing on the stream instance. - self.OUTPUT.buf_size = self.SETTINGS.buf_size - self.OUTPUT.num_buffers = self.SETTINGS.num_buffers - - @ez.publisher(OUTPUT) - async def spawn(self) -> AsyncGenerator: - for seq, size in enumerate(self.SETTINGS.sizes): - yield self.OUTPUT, BlobMessage(seq=seq, payload=b"x" * size) - raise ez.Complete - - -class BlobReceiverSettings(ez.Settings): - num_msgs: int - output_fn: str - - -class BlobReceiverState(ez.State): - num_received: int = 0 - - -class BlobReceiver(ez.Unit): - STATE = BlobReceiverState - SETTINGS = BlobReceiverSettings - - INPUT = ez.InputStream(BlobMessage) - - @ez.subscriber(INPUT) - async def on_message(self, msg: BlobMessage) -> None: - self.STATE.num_received += 1 - with open(self.SETTINGS.output_fn, "a") as output_file: - output_file.write( - json.dumps({"seq": msg.seq, "len": len(msg.payload)}) + "\n" - ) - if self.STATE.num_received == self.SETTINGS.num_msgs: - raise ez.Complete - - -class GrowSystemSettings(ez.Settings): - sizes: tuple[int, ...] - buf_size: int - num_buffers: int - output_fn: str - - -class GrowSystem(ez.Collection): - SETTINGS = GrowSystemSettings - - PUB = BlobGenerator() - SUB = BlobReceiver() - - def configure(self) -> None: - self.PUB.apply_settings( - BlobGeneratorSettings( - sizes=self.SETTINGS.sizes, - buf_size=self.SETTINGS.buf_size, - num_buffers=self.SETTINGS.num_buffers, - ) - ) - self.SUB.apply_settings( - BlobReceiverSettings( - num_msgs=len(self.SETTINGS.sizes), - output_fn=self.SETTINGS.output_fn, - ) - ) - - def network(self) -> ez.NetworkDefinition: - return ((self.PUB.OUTPUT, self.SUB.INPUT),) - - def process_components(self): - # Force PUB and SUB into separate processes so the boundary uses the - # cross-process SHM transport (and thus the grow/reattach path). - return (self.PUB, self.SUB) +from ezmsg.core.shm_grow_test_support import GrowSystem, GrowSystemSettings @pytest.mark.parametrize( "sizes", [ - # small, small, GROW (exceeds buf_size), small, small (8, 8, 16384, 8, 8), - # two successive grows of increasing size -> two reattaches (8, 16384, 8, 65536, 8), - # grow on the very first message (empty cache to repopulate) (16384, 8, 8), ], ) @@ -155,6 +46,5 @@ def test_cross_process_grow_delivers_all(sizes): for line in file: results.append(json.loads(line)) - # Every message survived the mid-stream grow: none dropped, in order. assert [r["seq"] for r in results] == list(range(len(sizes))) assert [r["len"] for r in results] == list(sizes) From 37e2faf56f754eecce12bdcfb60196b6099a1ca1 Mon Sep 17 00:00:00 2001 From: Griffin Milsap Date: Mon, 10 Aug 2026 14:39:10 -0400 Subject: [PATCH 7/7] fix for inflight message handling on shm resize --- src/ezmsg/core/messagechannel.py | 48 +++++++++++++++--------- tests/test_channel.py | 63 ++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 17 deletions(-) diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index 93251f4f..e6479486 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -266,11 +266,16 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: channel_kind = ProfileChannelType.SHM shm_name = await read_str(reader) - if self.shm is not None and self.shm.name != shm_name: - shm_entries = self.cache.keys() + if self.shm is None or self.shm.name != shm_name: + preserved_cache = self._snapshot_cached_messages() self.cache.clear() - self.shm.close() - await self.shm.wait_closed() + for preserved_msg in preserved_cache: + self.cache.put_from_mem(preserved_msg) + + if self.shm is not None: + old_shm = self.shm + old_shm.close() + await old_shm.wait_closed() try: self.shm = await GraphService( @@ -285,19 +290,6 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: ) self.shm = None - if self.shm is not None: - for id in shm_entries: - shm_buf = self.shm[id % self.num_buffers] - # A carried-over slot may not have survived the - # grow/copy; msg_id() raises UninitializedMemory - # on an uninitialized slot. Skip it rather than - # let that escape and kill the channel task. - try: - if MessageMarshal.msg_id(shm_buf) == id: - self.cache.put_from_mem(shm_buf) - except UninitializedMemory: - pass - if self.shm is None: logger.warning( "Channel %s dropping message %s from publisher %s because its SHM generation is stale", @@ -348,6 +340,13 @@ async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError): logger.debug(f"connection fail: channel:{self.id} - pub:{self.pub_id}") + except Exception: + logger.exception( + "Channel %s publisher connection crashed for pub %s", + self.id, + self.pub_id, + ) + raise finally: self.cache.clear() @@ -435,6 +434,21 @@ def release_without_get(self, msg_id: int, client_id: UUID) -> None: """ self._release_backpressure(msg_id, client_id) + def _snapshot_cached_messages(self) -> list[memoryview]: + if self.shm is None: + return [] + + preserved: list[memoryview] = [] + for msg_id in self.cache.keys(): + shm_buf = self.shm[msg_id % self.num_buffers] + try: + if MessageMarshal.msg_id(shm_buf) == msg_id: + preserved.append(memoryview(bytes(shm_buf)).toreadonly()) + except UninitializedMemory: + pass + + return preserved + def _release_backpressure(self, msg_id: int, client_id: UUID) -> None: """ Internal method to release backpressure for a message. diff --git a/tests/test_channel.py b/tests/test_channel.py index e45f7192..9db30ddb 100644 --- a/tests/test_channel.py +++ b/tests/test_channel.py @@ -5,6 +5,7 @@ from ezmsg.core.messagechannel import Channel from ezmsg.core.messagecache import CacheMiss +from ezmsg.core.messagemarshal import MessageMarshal from ezmsg.core.netprotocol import Command, uint64_to_bytes from ezmsg.core.backpressure import Backpressure @@ -16,6 +17,12 @@ def __init__(self): def write(self, data: bytes) -> None: self.buffer.append(data) + def close(self) -> None: + return None + + async def wait_closed(self) -> None: + return None + def _resolved_task(): loop = asyncio.get_running_loop() @@ -24,6 +31,28 @@ def _resolved_task(): return fut +class FakeSHM: + def __init__(self, name: str, slots: list[memoryview], on_wait_closed=None): + self.name = name + self._slots = slots + self._on_wait_closed = on_wait_closed + + def __getitem__(self, idx: int) -> memoryview: + return self._slots[idx] + + def close(self) -> None: + return None + + async def wait_closed(self) -> None: + if self._on_wait_closed is not None: + self._on_wait_closed() + + +def _raw_message(msg_id: int, payload) -> memoryview: + with MessageMarshal.serialize(msg_id, payload) as (_, header, buffers): + return memoryview(header + b"".join(buffers)).toreadonly() + + @pytest.mark.asyncio async def test_channel_acknowledges_remote_messages(): channel = Channel(uuid4(), uuid4(), 2, None, None, Channel._SENTINEL) @@ -91,3 +120,37 @@ def test_channel_put_local_requires_local_backpressure(): channel = Channel(uuid4(), uuid4(), 1, None, None, Channel._SENTINEL) with pytest.raises(ValueError): channel.put_local(1, "no pub") + + +@pytest.mark.asyncio +async def test_channel_preserves_cached_message_during_shm_reattach(monkeypatch): + old_slots = [_raw_message(0, {"value": 0}), _raw_message(1, {"value": 1})] + new_slots = [_raw_message(4, {"value": 4}), _raw_message(1, {"value": 1}), _raw_message(2, {"value": 2})] + + channel = Channel(uuid4(), uuid4(), 3, None, None, Channel._SENTINEL) + channel._pub_writer = DummyWriter() + channel._pub_task = _resolved_task() + channel._graph_task = _resolved_task() + preserved_during_reattach = False + + def assert_cached_before_close(): + nonlocal preserved_during_reattach + assert channel.cache[1] == {"value": 1} + preserved_during_reattach = True + + channel.shm = FakeSHM("old", old_slots, on_wait_closed=assert_cached_before_close) + channel.cache.put_from_mem(old_slots[1]) + + async def fake_attach_shm(self, shm_name): + assert shm_name == "new" + return FakeSHM("new", new_slots) + + monkeypatch.setattr("ezmsg.core.messagechannel.GraphService.attach_shm", fake_attach_shm) + + reader = asyncio.StreamReader() + reader.feed_data(Command.TX_SHM.value + uint64_to_bytes(2) + uint64_to_bytes(3) + b"new") + reader.feed_eof() + + await channel._publisher_connection(reader) + + assert preserved_during_reattach