Cross-process SHM grow (buf_size exceeded) can silently and permanently kill a subscriber channel via uncaught assert in MessageChannel.monitor

[cboulay]

Summary

A subscriber's MessageChannel can be permanently and silently killed when a publisher grows its shared-memory segment mid-stream (i.e. when a single message exceeds the current buf_size). The reattach path in MessageChannel.monitor uses bare asserts that are not covered by the loop's caught-exception set, so any transient inconsistency during the resize raises out of the receive loop, ends the channel task, and the subscriber stops receiving from that publisher forever — with no error surfaced to the application. Downstream units that depend on that stream simply stall.

Observed in ezmsg 3.9.0 on a real multi-process pipeline: a bursty upstream (data acquisition catch-up flood) occasionally emits messages far larger than the default 64 KB buf_size, forcing the grow, after which the consuming process silently stops receiving.

Where

Subscriber side — src/ezmsg/core/messagechannel.py, MessageChannel.monitor:

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])   # (1)

assert self.shm is not None                                        # (2)
assert MessageMarshal.msg_id(self.shm[buf_idx]) == msg_id          # (3)
self.cache.put_from_mem(self.shm[buf_idx])

The enclosing try only catches transport errors:

except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError):
    logger.debug(...)

So:

• (3) — if the buffer at buf_idx in the new segment does not carry the expected msg_id (a race window once the publisher has advanced/reused slots during/after the grow), the AssertionError is not caught → propagates out of monitor → the channel task dies. The subscriber silently stops receiving from this publisher for the remainder of the run.
• (1) — the repopulation loop reads carried-over cache entries out of the new segment; a slot that didn't survive the grow/copy raises UninitializedMemory, which is likewise uncaught and fatal.

Publisher side — src/ezmsg/core/pubclient.py, PublisherClient.broadcast (the grow):

if self._shm.buf_size < total_size:
    new_shm = await GraphService(self._graph_address).create_shm(self._num_buffers, total_size * 2)
    for i in range(self._num_buffers):
        try:
            with self._shm.buffer(i, readonly=True) as from_buf:
                with new_shm.buffer(i) as to_buf:
                    MessageMarshal.copy_obj(from_buf, to_buf)
        except UninitializedMemory:
            pass
    self._shm.close()            # can raise BufferError: cannot close exported pointers exist
    await self._shm.wait_closed()
    self._shm = new_shm

self._shm.close() here can raise BufferError: cannot close exported pointers exist if a memoryview into the old segment is still leased when the grow happens. That also propagates out of broadcast (uncaught), killing the publisher's send path.

Why it wedges silently

The failure is an AssertionError/UninitializedMemory/BufferError escaping the channel (or publisher) task rather than a transport exception. monitor's except list does not include these, so the coroutine terminates in its finally (closing the SHM) with no application-visible error. From the app's perspective the stream just goes quiet; nothing raises, nothing logs at warning/error level.

Impact

Any pipeline where cross-process messages can occasionally exceed buf_size (bursty sources, variable-size frames, catch-up floods after a stall) is exposed. Because it only triggers on the grow, it can run for a long time and then wedge under load — and the symptom (a downstream unit stops producing) points nowhere near the transport.

Reproduction

The grow path itself is easy to drive across processes; the wedge additionally needs the racy interleaving (publisher advancing/reusing slots while the subscriber reattaches), which requires real concurrency and is hard to make deterministic. A scaffold that forces the cross-process grow (small buf_size, one oversized message mid-stream, publisher and subscriber in separate processes) reliably exercises (1)/(2)/(3):

from dataclasses import dataclass
from collections.abc import AsyncGenerator
import ezmsg.core as ez

@dataclass
class Blob:
    seq: int
    payload: bytes

class Gen(ez.Unit):
    OUTPUT = ez.OutputStream(Blob, num_buffers=4, buf_size=4096, allow_local=False)
    @ez.publisher(OUTPUT)
    async def spawn(self) -> AsyncGenerator:
        for seq, size in enumerate([8, 8, 16384, 8, 8]):  # 16384 > buf_size -> grow
            yield self.OUTPUT, Blob(seq, b"x" * size)
        raise ez.Complete

class Recv(ez.Unit):
    INPUT = ez.InputStream(Blob)
    @ez.subscriber(INPUT)
    async def on_msg(self, msg: Blob) -> None:
        ...  # count / assert all 5 arrive in order

class Sys(ez.Collection):
    PUB = Gen()
    SUB = Recv()
    def network(self):
        return ((self.PUB.OUTPUT, self.SUB.INPUT),)
    def process_components(self):
        return (self.PUB, self.SUB)   # separate processes -> cross-process SHM

ez.run(SYSTEM=Sys())

Suggested fix

Make the reattach path resilient rather than fatal:

• Guard the repopulation loop against (UninitializedMemory, AssertionError) per slot (skip a slot that didn't survive the grow).
• Replace the bare assert on the published buf_idx with an explicit check: on msg_id mismatch (or lost SHM), log a warning, _acknowledge(msg_id) to release the publisher's backpressure for that buffer, and continue — dropping at most one message instead of killing the channel.
• On the publisher side, handle the BufferError from close() when exported pointers are still held during the grow (e.g. defer the close / retry after leases drain) so broadcast cannot die on a resize.

Happy to open a PR — I have a working patch for the subscriber-side hardening plus a cross-process grow test, and it passes the existing core suite (404 passed).

Candidate patch (subscriber-side hardening)

--- a/src/ezmsg/core/messagechannel.py
+++ b/src/ezmsg/core/messagechannel.py
@@ -7,7 +7,7 @@ from uuid import UUID
 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 .graphserver import GraphService
@@ -267,6 +267,14 @@ class Channel:
                     shm_name = await read_str(reader)
 
                     if self.shm is not None and self.shm.name != shm_name:
+                        # Publisher grew its shared-memory segment (a message
+                        # exceeded buf_size, see PublisherClient.broadcast) and
+                        # handed us a new name; detach and reattach. This path is
+                        # racy under load, and a transient inconsistency here must
+                        # NOT escape the receive loop: an uncaught error would
+                        # kill this channel task and silently, permanently stall
+                        # the subscriber. Recover instead (log + ack to keep the
+                        # publisher's backpressure moving).
                         shm_entries = self.cache.keys()
                         self.cache.clear()
                         self.shm.close()
@@ -282,11 +290,41 @@ class Channel:
                             )
                             raise
 
+                        # Repopulate carried-over entries from the new segment. A
+                        # slot that didn't survive the grow/copy (uninitialized or
+                        # already reused) is skipped rather than fatal.
                         for id in shm_entries:
-                            self.cache.put_from_mem(self.shm[id % self.num_buffers])
+                            try:
+                                self.cache.put_from_mem(self.shm[id % self.num_buffers])
+                            except (UninitializedMemory, AssertionError):
+                                pass
+
+                    if self.shm is None:
+                        logger.warning(
+                            f"channel:{self.id} -> pub:{self.pub_id}: no SHM for "
+                            f"msg {msg_id}; acking to release publisher"
+                        )
+                        self._acknowledge(msg_id)
+                        continue
+
+                    # The published buffer must carry the id we were told to read.
+                    # After a mid-stream resize a stale/reused slot can mismatch;
+                    # drop this message rather than deliver wrong bytes, but still
+                    # ack so the publisher's backpressure for this buffer frees
+                    # (otherwise it wedges once it wraps num_buffers sends).
+                    try:
+                        slot_msg_id = MessageMarshal.msg_id(self.shm[buf_idx])
+                    except UninitializedMemory:
+                        slot_msg_id = None
+                    if slot_msg_id != msg_id:
+                        logger.warning(
+                            f"channel:{self.id} -> pub:{self.pub_id}: SHM buffer "
+                            f"{buf_idx} holds msg {slot_msg_id} != expected "
+                            f"{msg_id} (post-resize skew); dropping + ack"
+                        )
+                        self._acknowledge(msg_id)
+                        continue
 
-                    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])

(This addresses the subscriber-side silent kill. The publisher-side BufferError-on-close during the grow is left for discussion — deferring/retrying the close until leases drain seems the natural fix.)

[griffinmilsap]

I suspect the issue is that a bursty publisher resizes memory twice before rx message channel has an opportunity to even get one message out of SHM

1. Publisher starts with SHM A, resizes to SHM B, for message ID 1 then resizes to SHM C. for message ID 2
2. After resize SHM C/message ID 2, publisher relinquishes lease for SHM B, but rx message channel has not yet acquired lease for SHM B -- SHM B no longer has any leases causing it to be released
3. RX Message channel receives message indicating Message ID 1 is located in SHM B, but then tries to attach to SHM B, and it doesn't exist, raising an error that is unhandled causing a silent failure of the message channel task

I have a minimal repro here

import asyncio
import os
import time

from dataclasses import dataclass

import ezmsg.core as ez

INITIAL_SHM_SIZE = 64 # bytes
NUM_MSGS = 4

@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
        for itr in range(NUM_MSGS):
            
            # SHM internally resizes to new message size * 2 
            # we force resizes on each new message by multiplying cur size * 3
            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:
        ez.logger.info(
            "subscriber received seq=%s payload=%s pid=%s",
            msg.seq,
            len(msg.payload),
            os.getpid(),
        )

        # Curiously, this sleep seems to be pretty required to reliably trigger the hang...
        # This is non-intuitive because MessageChannel should be receiving/caching messages
        # in its own task on the backend unrelated to how quickly we can clear the queue
        await asyncio.sleep(0.001)
        
        self.STATE.cur_msg += 1

        if self.STATE.cur_msg == NUM_MSGS:
            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__":
    print(
        "Running SHM resize race reproducer. "
        "Use EZMSG_LOGLEVEL=DEBUG to watch channel / SHM behavior."
    )
    print(
        "This intentionally sends one rapid burst of ever-larger messages "
        "to force repeated SHM growth under a slow cross-process subscriber."
    )

    ez.run(SYSTEM=ReproSystem())

[griffinmilsap]

There's a fix for this in #254, assuming its the same failure mechanism as the repro example I posted above.