Merge pull request #218 from ezmsg-org/217-subscriber-never-updates-its-set-of-current-publishers-after-processing-update-preventing-disconnection-of-old-publishers

hot fix 217 - recalculate Subscriber's pub keys on demand

Author: cboulay

src/ezmsg/core/subclient.py

@@ -39,7 +39,6 @@ class Subscriber:
 
     _graph_address: AddressType | None
     _graph_task: asyncio.Task[None]
-    _cur_pubs: set[UUID]
     _incoming: NotificationQueue
 
     # FIXME: This event allows Subscriber.create to block until
@@ -128,7 +127,6 @@ def __init__(
         self.leaky = leaky
         self._graph_address = graph_address
 
-        self._cur_pubs = set()
         self._channels = dict()
         if self.leaky:
             self._incoming = LeakyQueue(
@@ -223,8 +221,10 @@ async def _graph_connection(
                     pub_ids = (
                         set([UUID(id) for id in update.split(",")]) if update else set()
                     )
+                    cur_pubs = set(self._channels.keys())
 
-                    for pub_id in set(pub_ids - self._cur_pubs):
+                    # Register new channels
+                    for pub_id in set(pub_ids - cur_pubs):
                         channel = await CHANNELS.register(
                             pub_id, self.id, self._incoming, self._graph_address
                         )
@@ -239,7 +239,8 @@ async def _graph_connection(
 
                         self._channels[pub_id] = channel
 
-                    for pub_id in set(self._cur_pubs - pub_ids):
+                    # Unregister expired channels
+                    for pub_id in set(cur_pubs - pub_ids):
                         await CHANNELS.unregister(pub_id, self.id, self._graph_address)
                         del self._channels[pub_id]
 
@@ -288,7 +289,11 @@ async def recv_zero_copy(self) -> typing.AsyncGenerator[typing.Any, None]:
         :return: Context manager yielding the received message.
         :rtype: collections.abc.AsyncGenerator[typing.Any, None]
         """
-        pub_id, msg_id = await self._incoming.get()
+        while True:
+            pub_id, msg_id = await self._incoming.get()
+            if pub_id in self._channels:
+                break
+            # Stale notification from an unregistered publisher — skip.
 
         with self._channels[pub_id].get(msg_id, self.id) as msg:
             yield msg

tests/test_subclient.py

@@ -0,0 +1,148 @@
+import asyncio
+from contextlib import contextmanager
+from uuid import uuid4
+
+import pytest
+
+from ezmsg.core.subclient import Subscriber
+from ezmsg.core.netprotocol import Command, encode_str
+from ezmsg.core import channelmanager as channelmanager_module
+from ezmsg.core import subclient as subclient_module
+
+
+class DummyChannel:
+    """Minimal Channel stand-in for Subscriber tests."""
+
+    def __init__(self):
+        self.clients = {}
+        self.closed = False
+        self.waited = False
+        self.topic = "test"
+        self.num_buffers = 8
+
+    def register_client(self, client_id, queue, local_backpressure=None):
+        self.clients[client_id] = queue
+
+    def unregister_client(self, client_id):
+        del self.clients[client_id]
+
+    def close(self):
+        self.closed = True
+
+    async def wait_closed(self):
+        self.waited = True
+
+    @contextmanager
+    def get(self, msg_id, client_id):
+        yield f"msg-{msg_id}"
+
+
+class DummyWriter:
+    """Minimal asyncio.StreamWriter stand-in."""
+
+    def __init__(self):
+        self.buffer = []
+        self._closed = False
+
+    def write(self, data):
+        self.buffer.append(data)
+
+    async def drain(self):
+        pass
+
+    def close(self):
+        self._closed = True
+
+    async def wait_closed(self):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_subscriber_unregisters_removed_publisher(monkeypatch):
+    """A graph UPDATE that drops a publisher must remove it from _channels.
+
+    Before the fix (PR #218), Subscriber._cur_pubs was initialised to an
+    empty set and never updated, so ``cur_pubs - pub_ids`` was always
+    empty and stale publishers were never unregistered.
+    """
+    async def fake_create(pub_id, address):
+        return DummyChannel()
+
+    monkeypatch.setattr(channelmanager_module.Channel, "create", fake_create)
+    monkeypatch.setattr(
+        subclient_module, "CHANNELS", channelmanager_module.ChannelManager()
+    )
+
+    pub_a = uuid4()
+
+    reader = asyncio.StreamReader()
+    writer = DummyWriter()
+
+    sub = Subscriber(
+        id=uuid4(),
+        topic="test/topic",
+        graph_address=None,
+        _guard=Subscriber._SENTINEL,
+    )
+
+    # Protocol sequence:
+    #   UPDATE [pub_a]  ->  subscriber registers pub_a
+    #   COMPLETE        ->  _initialized is set
+    #   UPDATE []       ->  subscriber should unregister pub_a
+    reader.feed_data(
+        Command.UPDATE.value
+        + encode_str(str(pub_a))
+        + Command.COMPLETE.value
+        + Command.UPDATE.value
+        + encode_str("")
+    )
+    # No EOF yet — connection stays open after the second UPDATE.
+
+    task = asyncio.create_task(sub._graph_connection(reader, writer))
+
+    # Wait until the subscriber has written two COMPLETE responses
+    # (one per UPDATE), meaning both UPDATEs have been processed.
+    for _ in range(200):
+        if len(writer.buffer) >= 2:
+            break
+        await asyncio.sleep(0)
+    else:
+        pytest.fail("Subscriber did not process both UPDATEs")
+
+    # Connection is still open — this is mid-session state, not post-cleanup.
+    # Before the fix, pub_a would still be in _channels here.
+    assert pub_a not in sub._channels, (
+        "Publisher should be removed from _channels "
+        "when a graph UPDATE no longer includes it"
+    )
+
+    reader.feed_eof()
+    await task
+
+
+@pytest.mark.asyncio
+async def test_recv_zero_copy_skips_stale_notification():
+    """recv_zero_copy must skip notifications from unregistered publishers.
+
+    Before PR #218 there was no guard; a stale notification would cause
+    a KeyError on ``self._channels[pub_id]``.
+    """
+    sub = Subscriber(
+        id=uuid4(),
+        topic="test/topic",
+        graph_address=None,
+        _guard=Subscriber._SENTINEL,
+    )
+
+    stale_pub = uuid4()
+    valid_pub = uuid4()
+
+    sub._channels[valid_pub] = DummyChannel()
+
+    # Stale notification first, then a valid one.
+    await sub._incoming.put((stale_pub, 0))
+    await sub._incoming.put((valid_pub, 1))
+
+    # Before the fix this would raise KeyError for stale_pub.
+    async with sub.recv_zero_copy() as msg:
+        assert msg == "msg-1"