fixed bug where messages could be dropped

Author: griffinmilsap

src/ezmsg/core/sync.py

@@ -257,10 +257,10 @@ def spin(self, poll_interval: float = 0.1) -> None:
 
     def spin_once(self, timeout: float | None = 0.0) -> bool:
         """
-        Process at most one subscription callback.
+        Process any subscription callbacks ready within the timeout window.
 
-        :param timeout: Seconds to wait for a callback. Use None to block forever.
-        :return: True if a callback was processed, False otherwise.
+        :param timeout: Seconds to wait for callbacks. Use None to block forever.
+        :return: True if any callback was processed, False otherwise.
         """
         self._ensure_started()
         if self._shutdown_requested.is_set():
@@ -289,31 +289,32 @@ def spin_once(self, timeout: float | None = 0.0) -> bool:
         finally:
             if not keep_future:
                 self._spin_future = None
-        if result is None:
+        if not result:
             return False
+        processed = False
+        for entry, cm, msg in result:
+            _, callback, zero_copy = entry
 
-        entry, cm, msg = result
-        _, callback, zero_copy = entry
-
-        try:
-            if not zero_copy:
-                msg = deepcopy(msg)
-            callback(msg)
-        except Exception:
-            logger.exception("Unhandled exception in subscription callback")
-        finally:
-            exit_fut = asyncio.run_coroutine_threadsafe(
-                cm.__aexit__(None, None, None), self._loop
-            )
             try:
-                _future_result(exit_fut, None)
-            except CacheMiss:
-                logger.warning(
-                    "Cache miss while releasing message; publisher likely exited."
-                )
+                if not zero_copy:
+                    msg = deepcopy(msg)
+                callback(msg)
             except Exception:
-                logger.exception("Failed while releasing message backpressure")
-        return True
+                logger.exception("Unhandled exception in subscription callback")
+            finally:
+                exit_fut = asyncio.run_coroutine_threadsafe(
+                    cm.__aexit__(None, None, None), self._loop
+                )
+                try:
+                    _future_result(exit_fut, None)
+                except CacheMiss:
+                    logger.warning(
+                        "Cache miss while releasing message; publisher likely exited."
+                    )
+                except Exception:
+                    logger.exception("Failed while releasing message backpressure")
+            processed = True
+        return processed
 
     def shutdown(self) -> None:
         self._shutdown_requested.set()
@@ -368,30 +369,14 @@ def spin(context: SyncContext, poll_interval: float = 0.1) -> None:
 
 
 def spin_once(context: SyncContext, timeout: float | None = 0.0) -> bool:
-    """Process at most one subscription callback."""
+    """Process any subscription callbacks ready within the timeout window."""
     return context.spin_once(timeout=timeout)
 
 
 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")
-
+) -> list[tuple[tuple[SyncSubscriber, Callable[[Any], None], bool], Any, Any]]:
     async def _recv_entry(
         entry: tuple[SyncSubscriber, Callable[[Any], None], bool]
     ) -> tuple[tuple[SyncSubscriber, Callable[[Any], None], bool], Any, Any]:
@@ -410,20 +395,9 @@ async def _recv_entry(
             done, pending = await asyncio.wait(
                 tasks, timeout=remaining, return_when=asyncio.FIRST_COMPLETED
             )
-            if not done:
-                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)
-                return None
-
-            winner_result = None
+            results: list[
+                tuple[tuple[SyncSubscriber, Callable[[Any], None], bool], Any, Any]
+            ] = []
             for task in done:
                 try:
                     result = task.result()
@@ -435,28 +409,32 @@ 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)
+                results.append(result)
 
             for task in pending:
                 try:
                     task.cancel()
                 except RuntimeError:
                     pass
-            pending_results = 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)
+                if isinstance(result, CacheMiss):
+                    continue
+                if isinstance(result, asyncio.CancelledError):
+                    continue
+                if isinstance(result, BaseException):
+                    logger.exception(
+                        "Sync subscription receive failed", exc_info=result
+                    )
+                    continue
+                results.append(result)
 
-            if winner_result is not None:
-                return winner_result
+            if results:
+                return results
 
             # Only CacheMiss/cancelled/error occurred; continue within timeout window.
             if deadline is not None and loop.time() >= deadline:
-                return None
+                return []
         finally:
             for task in tasks:
                 if not task.done():

tests/test_sync_api.py

@@ -64,6 +64,34 @@ def on_msg(msg: str) -> None:
         spin_thread.join(timeout=1.0)
 
 
+def test_spin_once_processes_all_ready_callbacks():
+    host = "127.0.0.1"
+    port = _free_port()
+
+    with ez.sync.init((host, port), auto_start=True) as ctx:
+        received: list[tuple[str, str]] = []
+
+        def on_a(msg: str) -> None:
+            received.append(("A", msg))
+
+        def on_b(msg: str) -> None:
+            received.append(("B", msg))
+
+        ctx.create_subscription("/A", callback=on_a, zero_copy=True)
+        ctx.create_subscription("/B", callback=on_b, zero_copy=True)
+
+        pub_a = ctx.create_publisher("/A", num_buffers=1, force_tcp=True)
+        pub_b = ctx.create_publisher("/B", num_buffers=1, force_tcp=True)
+
+        time.sleep(0.05)
+        pub_a.publish("one")
+        pub_b.publish("two")
+        time.sleep(0.05)
+
+        assert ctx.spin_once(timeout=0.2) is True
+        assert set(received) == {("A", "one"), ("B", "two")}
+
+
 @pytest.mark.asyncio
 async def test_recv_any_cachemiss_does_not_raise():
     class DummySub:
@@ -78,11 +106,11 @@ def __init__(self) -> None:
 
     entry = (DummySyncSub(), lambda _: None, True)
     result = await sync_mod._recv_any([entry], timeout=0.01)
-    assert result is None
+    assert result == []
 
 
 @pytest.mark.asyncio
-async def test_recv_any_cleans_up_non_winner_contexts():
+async def test_recv_any_returns_all_ready_contexts():
     class DummyCM:
         def __init__(self, name: str, gate: asyncio.Event | None) -> None:
             self._name = name
@@ -131,12 +159,13 @@ async def _release_pending() -> None:
     asyncio.create_task(_release_pending())
 
     result = await sync_mod._recv_any(entries, timeout=0.2)
-    assert result is not None
+    assert len(result) == 3
+
+    returned_cms = {cm for _, cm, _ in result}
+    assert returned_cms == set(cms)
 
-    _, winner_cm, _ = result
-    await winner_cm.__aexit__(None, None, None)
+    for _, cm, _ in result:
+        await cm.__aexit__(None, None, None)
 
     for cm in cms:
-        if cm is winner_cm:
-            continue
         assert cm.exited.is_set()