Fix: Channel failure on shm resize due to CacheMiss

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:

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

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)