Repository: rabbitmq-community/rstream (language: python)

Issue:
specify a timeout option to send_wait
At the moment send_wait waits until publish_confirmation is received from the server.

Better to put a timeout that will allow the function to continue if the timeout expires even if the message is not yet confirmed


Discussion / hints:
This task also to revisit send_wait (we did most of the tests for send and send_batch at the moment)

Fix patch:
diff --git a/docs/examples/basic_producers/README.md b/docs/examples/basic_producers/README.md
index d04d294..6518cb7 100644
--- a/docs/examples/basic_producers/README.md
+++ b/docs/examples/basic_producers/README.md
@@ -8,7 +8,8 @@ The library is using a dedicated background thread to send the messages to the R
 See producer_send.py
 * Using send_batch: Send batch is synchronous. You send a batch of messages and wait till they are sent to the server
 See producer_send_batch.py
-* Using send_wait: send_wait is synchronous and it will also wait until confirmation from the server have been received.
+* Using send_wait: send_wait is synchronous and it will block until confirmation from the server have been received.
+A timeout value can be specified (default 5s) in the method, after that an exception is raised.
 
 send and send_batch by default don't wait for confirmation. You can enable asynchronous confirmation through callbacks.
 You can find examples in producers_with_confirmations folder.
diff --git a/rstream/producer.py b/rstream/producer.py
index defdfff..396c4da 100644
--- a/rstream/producer.py
+++ b/rstream/producer.py
@@ -256,23 +256,78 @@ async def send_batch(
         on_publish_confirm: Optional[CB[ConfirmationStatus]] = None,
     ) -> list[int]:
 
+        await self.check_connection()
+
+        if len(batch) == 0:
+            raise ValueError("Empty batch")
+
         if self._close_called:
             return []
 
-        wrapped_batch = []
+        return await self._send_batch(
+            stream=stream, batch=batch, publisher_name=publisher_name, callback=on_publish_confirm, sync=False
+        )
+
+    # used by send_batch and send_wait
+    async def _send_batch(
+        self,
+        stream: str,
+        batch: list[MessageT],
+        publisher_name: Optional[str] = None,
+        callback: Optional[CB[ConfirmationStatus]] = None,
+        sync: bool = True,
+        timeout: Optional[int] = None,
+    ) -> list[int]:
+
+        messages = []
+        publishing_ids = set()
+
+        async with self._lock:
+            publisher = await self._get_or_create_publisher(stream, publisher_name=publisher_name)
+
         for item in batch:
-            wrapped_item = _MessageNotification(
-                entry=item, callback=on_publish_confirm, publisher_name=publisher_name
+
+            msg = RawMessage(item) if isinstance(item, bytes) else item
+
+            if msg.publishing_id is None:
+                msg.publishing_id = publisher.sequence.next()
+
+            messages.append(
+                schema.Message(
+                    publishing_id=msg.publishing_id,
+                    data=bytes(msg),
+                )
             )
-            wrapped_batch.append(wrapped_item)
 
-        return await self._send_batch(stream, wrapped_batch, sync=False)
+        if len(messages) > 0:
+            await publisher.client.send_frame(
+                schema.Publish(
+                    publisher_id=publisher.id,
+                    messages=messages,
+                ),
+            )
+            publishing_ids.update([m.publishing_id for m in messages])
 
-    async def _send_batch(
+        if not sync:
+            if callback is not None:
+                if callback not in self._waiting_for_confirm[publisher.reference]:
+                    self._waiting_for_confirm[publisher.reference][callback] = set()
+
+                self._waiting_for_confirm[publisher.reference][callback].update(publishing_ids)
+
+        # this is just called in case of send_wait
+        else:
+            future: asyncio.Future[None] = asyncio.Future()
+            self._waiting_for_confirm[publisher.reference][future] = publishing_ids.copy()
+            await asyncio.wait_for(future, timeout)
+
+        return list(publishing_ids)
+
+    # used by send, send_sub_batch (with compression)
+    async def _send_batch_async(
         self,
         stream: str,
         batch: list[_MessageNotification],
-        sync: bool = True,
     ) -> list[int]:
         if len(batch) == 0:
             raise ValueError("Empty batch")
@@ -353,11 +408,6 @@ async def _send_batch(
                 self._waiting_for_confirm[publisher.reference][callback] = set()
 
             self._waiting_for_confirm[publisher.reference][callback].update(publishing_ids_callback[callback])
-        # this is just called in case of send_wait
-        if sync:
-            future: asyncio.Future[None] = asyncio.Future()
-            self._waiting_for_confirm[publisher.reference][future] = publishing_ids.copy()
-            await future
 
         return list(publishing_ids)
 
@@ -366,16 +416,17 @@ async def send_wait(
         stream: str,
         message: MessageT,
         publisher_name: Optional[str] = None,
+        timeout: Optional[int] = 5,
     ) -> int:
 
-        wrapped_message: _MessageNotification = _MessageNotification(
-            entry=message, callback=None, publisher_name=publisher_name
-        )
+        await self.check_connection()
 
         publishing_ids = await self._send_batch(
             stream,
-            [wrapped_message],
+            [message],
+            publisher_name=publisher_name,
             sync=True,
+            timeout=timeout,
         )
         return publishing_ids[0]
 
@@ -460,7 +511,7 @@ async def _publish_buffered_messages(self, stream: str) -> None:
 
         async with self._buffered_messages_lock:
             if len(self._buffered_messages[stream]):
-                await self._send_batch(stream, self._buffered_messages[stream], sync=False)
+                await self._send_batch_async(stream, self._buffered_messages[stream])
                 self._buffered_messages[stream].clear()
 
     async def _on_publish_confirm(self, frame: schema.PublishConfirm, publisher: _Publisher) -> None:


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index d48fb3c..5081bfa 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -89,6 +89,15 @@ async def stream(client: Client):
         await client.delete_stream("test-stream")
 
 
+@pytest.fixture()
+async def stream2(client: Client):
+    await client.create_stream("test-stream-2")
+    try:
+        yield "test-stream-2"
+    finally:
+        await client.delete_stream("test-stream-2")
+
+
 @pytest.fixture()
 async def consumer(pytestconfig, ssl_context):
     consumer = Consumer(
diff --git a/tests/test_producer.py b/tests/test_producer.py
index e7e3a56..a3676e7 100644
--- a/tests/test_producer.py
+++ b/tests/test_producer.py
@@ -77,6 +77,27 @@ async def test_publishing_several_messages(stream: str, producer: Producer, cons
     await wait_for(lambda: len(captured) == 100000)
 
 
+async def test_publishing_several_messages_different_streams(
+    stream: str, stream2: str, producer: Producer, consumer: Consumer
+) -> None:
+    captured_stream_1: list[bytes] = []
+    captured_stream_2: list[bytes] = []
+    await consumer.subscribe(
+        stream, callback=lambda message, message_context: captured_stream_1.append(bytes(message))
+    )
+    await consumer.subscribe(
+        stream2, callback=lambda message, message_context: captured_stream_2.append(bytes(message))
+    )
+
+    for i in range(0, 100000):
+        await producer.send(stream, b"one")
+    for i in range(0, 100000):
+        await producer.send(stream2, b"one")
+
+    await wait_for(lambda: len(captured_stream_1) == 100000)
+    await wait_for(lambda: len(captured_stream_2) == 100000)
+
+
 async def test_publishing_sequence_subbatching_nocompression(
     stream: str, producer: Producer, consumer: Consumer
 ) -> None: