Improve callback handlers

rstream/client.py

@@ -11,6 +11,7 @@
 import time
 from collections import defaultdict
 from contextlib import suppress
+from dataclasses import dataclass
 from typing import (
     Annotated,
     Any,
@@ -57,6 +58,12 @@ class Addr(NamedTuple):
     port: int
 
 
+@dataclass
+class FrameEntry:
+    handler: HT[Any]
+    frame: schema.Frame
+
+
 class BaseClient:
     def __init__(
         self,
@@ -95,6 +102,7 @@ def __init__(
 
         self._last_heartbeat: float = 0
         self._connection_closed_handler = connection_closed_handler
+        self._frames: asyncio.Queue = asyncio.Queue()
 
     def start_task(self, name: str, coro: Awaitable[None]) -> None:
         assert name not in self._tasks
@@ -172,9 +180,22 @@ async def start(self) -> None:
         self._conn = Connection(self.host, self.port, self._ssl_context)
         await self._conn.open()
         self.start_task("listener", self._listener())
+        self.start_task("run_delivery_handlers", self._run_delivery_handlers())
         self.add_handler(schema.Heartbeat, self._on_heartbeat)
         self.add_handler(schema.Close, self._on_close)
 
+    async def _run_delivery_handlers(self):
+
+        while True:
+            try:
+                frame_entry = await self._frames.get()
+                maybe_coro = await frame_entry.handler(frame_entry.frame)
+                if maybe_coro is not None:
+                    await maybe_coro
+            except Exception as e:
+                logger.debug("Error reading item from _run_delivery_handlers: " + str(e))
+                break
+
     async def _listener(self) -> None:
         assert self._conn
         while True:
@@ -206,9 +227,12 @@ async def _listener(self) -> None:
 
             for _, handler in self._handlers.get(frame.__class__, {}).items():
                 try:
-                    maybe_coro = handler(frame)
-                    if maybe_coro is not None:
-                        await maybe_coro
+                    if frame.__class__ == schema.Deliver:
+                        await self._frames.put(FrameEntry(handler, frame))
+                    else:
+                        maybe_coro = handler(frame)
+                        if maybe_coro is not None:
+                            await maybe_coro
                 except Exception:
                     logger.exception("Error while running handler %s of frame %s", handler, frame)
 
@@ -243,6 +267,7 @@ def is_started(self) -> bool:
 
     async def close(self) -> None:
         logger.info("Stopping client %s:%s", self.host, self.port)
+
         if self._conn is None:
             return
 
@@ -256,6 +281,7 @@ async def close(self) -> None:
                 resp_schema=schema.CloseResponse,
             )
 
+        await self.stop_task("run_delivery_handlers")
         await self.stop_task("listener")
 
         await self._conn.close()

tests/conftest.py

@@ -107,6 +107,24 @@ async def consumer(pytestconfig, ssl_context):
         await consumer.close()
 
 
+@pytest.fixture()
+async def test_consumer_close(pytestconfig, ssl_context):
+    consumer = Consumer(
+        host=pytestconfig.getoption("rmq_host"),
+        port=pytestconfig.getoption("rmq_port"),
+        ssl_context=ssl_context,
+        username=pytestconfig.getoption("rmq_username"),
+        password=pytestconfig.getoption("rmq_password"),
+        frame_max=1024 * 1024,
+        heartbeat=60,
+    )
+    await consumer.start()
+    try:
+        yield consumer
+    finally:
+        await consumer.close()
+
+
 @pytest.fixture()
 async def producer(pytestconfig, ssl_context):
     producer = Producer(

tests/test_consumer.py

@@ -502,3 +502,17 @@ async def test_consume_superstream_with_callback_offset(
     consumers_set = consumers_set.union(consumer_stream_list3)
 
     assert len(consumers_set) == 3
+
+
+async def test_callback_sync_request(stream: str, test_consumer_close: Consumer, producer: Producer) -> None:
+    captured: list[bytes] = []
+
+    async def on_message_first(msg: AMQPMessage, message_context: MessageContext):
+        captured.append(bytes(msg))
+        await test_consumer_close.close()
+
+    await test_consumer_close.subscribe(stream, callback=on_message_first)
+    messages = [str(i).encode() for i in range(0, 1)]
+    await producer.send_batch(stream, messages)
+
+    await wait_for(lambda: len(captured) >= 1)