diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index 130895f1..e6479486 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -7,9 +7,9 @@ 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 +from .messagecache import MessageCache, CacheMiss from .graphserver import GraphService from .netprotocol import ( Command, @@ -266,28 +266,60 @@ 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( 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 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] + # 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, + 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 @@ -308,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() @@ -395,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. @@ -407,7 +461,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: 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/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_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 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) diff --git a/tests/test_shm_grow.py b/tests/test_shm_grow.py new file mode 100644 index 00000000..f4d76e8a --- /dev/null +++ b/tests/test_shm_grow.py @@ -0,0 +1,50 @@ +"""Cross-process SHM grow / reattach coverage. + +When a published message exceeds the publisher's shared-memory ``buf_size``, +``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. +""" + +import json + +import pytest + +import ezmsg.core as ez + +from ez_test_utils import get_test_fn +from ezmsg.core.shm_grow_test_support import GrowSystem, GrowSystemSettings + + +@pytest.mark.parametrize( + "sizes", + [ + (8, 8, 16384, 8, 8), + (8, 16384, 8, 65536, 8), + (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)) + + assert [r["seq"] for r in results] == list(range(len(sizes))) + assert [r["len"] for r in results] == list(sizes)