addressed copilot review

Author: griffinmilsap

examples/simple_async_subscriber.py

@@ -2,7 +2,6 @@
 
 import ezmsg.core as ez
 
-PORT = 12345
 TOPIC = "/TEST"
 
 

examples/simple_subscriber.py

@@ -1,4 +1,3 @@
-import time
 import ezmsg.core as ez
 
 TOPIC = "/TEST"
@@ -10,6 +9,7 @@ def main(host: str = "127.0.0.1", port: int = 12345) -> None:
 
         def on_message(msg: str) -> None:
             # Uncomment if you want to witness backpressure!
+            # import time
             # time.sleep(1.0)
             print(msg)
 

src/ezmsg/core/graphserver.py

@@ -56,6 +56,7 @@ class GraphServer(threading.Thread):
     _shutdown: threading.Event
 
     _sock: socket.socket
+    _address: Address | None
     _loop: asyncio.AbstractEventLoop
 
     graph: DAG
@@ -78,11 +79,13 @@ def __init__(self, **kwargs) -> None:
         self.clients = {}
         self._client_tasks = {}
         self.shms = {}
+        self._address = None
 
     @property
     def address(self) -> Address:
-        return Address(*self._sock.getsockname())
-
+        assert self._address is not None, "GraphServer not up yet"
+        return self._address
+        
     def start(self, address: AddressType | None = None) -> None:  # type: ignore[override]
         if address is not None:
             self._sock = create_socket(*address)
@@ -92,10 +95,13 @@ def start(self, address: AddressType | None = None) -> None:  # type: ignore[ove
             )
             self._sock = create_socket(start_port=start_port)
 
+        # Cache address immediately to avoid touching a possibly-closed socket later.
+        self._address = Address(*self._sock.getsockname())
+
         self._loop = asyncio.new_event_loop()
         super().start()
         self._server_up.wait()
-        logger.info(f'Started GraphServer at {address}')
+        logger.info(f'Started GraphServer at {self.address}')
 
     def stop(self) -> None:
         self._shutdown.set()

src/ezmsg/core/sync.py

@@ -44,13 +44,13 @@ def sync(self, timeout: float | None = None) -> None:
         _future_result(fut, timeout)
 
     def pause(self) -> None:
-        self._pub.pause()
+        self._loop.call_soon_threadsafe(self._pub.pause)
 
     def resume(self) -> None:
-        self._pub.resume()
+        self._loop.call_soon_threadsafe(self._pub.resume)
 
     def close(self) -> None:
-        self._pub.close()
+        self._loop.call_soon_threadsafe(self._pub.close)
 
     def wait_closed(self, timeout: float | None = None) -> None:
         fut = asyncio.run_coroutine_threadsafe(self._pub.wait_closed(), self._loop)
@@ -101,7 +101,7 @@ def recv_zero_copy(self, timeout: float | None = None) -> _SyncZeroCopy:
         return _SyncZeroCopy(self._sub, self._loop, timeout)
 
     def close(self) -> None:
-        self._sub.close()
+        self._loop.call_soon_threadsafe(self._sub.close)
 
     def wait_closed(self, timeout: float | None = None) -> None:
         fut = asyncio.run_coroutine_threadsafe(self._sub.wait_closed(), self._loop)
@@ -134,6 +134,14 @@ def graph_address(self) -> AddressType | None:
         return self._graph_context.graph_address
 
     def __enter__(self) -> "SyncContext":
+
+        # SyncContext instances are single-use: they cannot be re-entered after shutdown.
+        if self._closed:
+            raise RuntimeError(
+                "SyncContext instances cannot be reused after shutdown; "
+                "create a new SyncContext instead."
+            )
+        
         if self._loop_cm is not None:
             return self
 
@@ -368,6 +376,22 @@ async def _recv_any(
     entries: Iterable[tuple[SyncSubscriber, Callable[[Any], None], bool]],
     timeout: float | None,
 ) -> tuple[tuple[SyncSubscriber, Callable[[Any], None], bool], Any, Any] | None:
+    async def _cleanup_result(result: Any) -> None:
+        if isinstance(result, BaseException):
+            return
+        try:
+            _, cm, _ = result
+        except Exception:
+            return
+        try:
+            await cm.__aexit__(None, None, None)
+        except CacheMiss:
+            logger.warning(
+                "Cache miss while releasing message; publisher likely exited."
+            )
+        except Exception:
+            logger.exception("Failed while releasing message backpressure")
+
     async def _recv_entry(
         entry: tuple[SyncSubscriber, Callable[[Any], None], bool]
     ) -> tuple[tuple[SyncSubscriber, Callable[[Any], None], bool], Any, Any]:
@@ -392,19 +416,17 @@ async def _recv_entry(
                         task.cancel()
                     except RuntimeError:
                         pass
-                await asyncio.gather(*pending, return_exceptions=True)
+                pending_results = await asyncio.gather(
+                    *pending, return_exceptions=True
+                )
+                for result in pending_results:
+                    await _cleanup_result(result)
                 return None
 
-            for task in pending:
-                try:
-                    task.cancel()
-                except RuntimeError:
-                    pass
-            await asyncio.gather(*pending, return_exceptions=True)
-
+            winner_result = None
             for task in done:
                 try:
-                    return task.result()
+                    result = task.result()
                 except CacheMiss:
                     # Likely stale notification after publisher exit; keep waiting.
                     continue
@@ -413,6 +435,24 @@ async def _recv_entry(
                 except Exception:
                     logger.exception("Sync subscription receive failed")
                     continue
+                if winner_result is None:
+                    winner_result = result
+                else:
+                    await _cleanup_result(result)
+
+            for task in pending:
+                try:
+                    task.cancel()
+                except RuntimeError:
+                    pass
+            pending_results = await asyncio.gather(
+                *pending, return_exceptions=True
+            )
+            for result in pending_results:
+                await _cleanup_result(result)
+
+            if winner_result is not None:
+                return winner_result
 
             # Only CacheMiss/cancelled/error occurred; continue within timeout window.
             if deadline is not None and loop.time() >= deadline:

tests/test_sync_api.py

@@ -1,3 +1,4 @@
+import asyncio
 import socket
 import threading
 import time
@@ -56,7 +57,8 @@ def on_msg(msg: str) -> None:
         t2 = time.monotonic()
 
         assert t2 - t1 >= 0.15
-        done.wait(2.0)
+        assert done.wait(2.0), "Timed out waiting for messages to be received"
+        assert received == ["one", "two"]
 
         ctx.shutdown()
         spin_thread.join(timeout=1.0)
@@ -77,3 +79,64 @@ def __init__(self) -> None:
     entry = (DummySyncSub(), lambda _: None, True)
     result = await sync_mod._recv_any([entry], timeout=0.01)
     assert result is None
+
+
+@pytest.mark.asyncio
+async def test_recv_any_cleans_up_non_winner_contexts():
+    class DummyCM:
+        def __init__(self, name: str, gate: asyncio.Event | None) -> None:
+            self._name = name
+            self._gate = gate
+            self.exited = asyncio.Event()
+
+        async def __aenter__(self):
+            if self._gate is not None:
+                try:
+                    await self._gate.wait()
+                except asyncio.CancelledError:
+                    await self._gate.wait()
+            return f"msg-{self._name}"
+
+        async def __aexit__(self, exc_type, exc, tb):
+            self.exited.set()
+
+    class DummySub:
+        def __init__(self, cm: DummyCM) -> None:
+            self._cm = cm
+
+        def recv_zero_copy(self):
+            return self._cm
+
+    class DummySyncSub:
+        def __init__(self, cm: DummyCM) -> None:
+            self._sub = DummySub(cm)
+
+    release = asyncio.Event()
+    cms: list[DummyCM] = []
+    entries = []
+
+    def add_entry(name: str, gate: asyncio.Event | None) -> None:
+        cm = DummyCM(name, gate)
+        cms.append(cm)
+        entries.append((DummySyncSub(cm), lambda _: None, True))
+
+    add_entry("fast", None)
+    add_entry("slow1", release)
+    add_entry("slow2", release)
+
+    async def _release_pending() -> None:
+        await asyncio.sleep(0.01)
+        release.set()
+
+    asyncio.create_task(_release_pending())
+
+    result = await sync_mod._recv_any(entries, timeout=0.2)
+    assert result is not None
+
+    _, winner_cm, _ = result
+    await winner_cm.__aexit__(None, None, None)
+
+    for cm in cms:
+        if cm is winner_cm:
+            continue
+        assert cm.exited.is_set()