Fix: Channel failure on shm resize due to CacheMiss

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:

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)

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())