Merge pull request #491 from ydb-platform/topic_batch_messages

Ability to batch messages in topic reader

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -40,15 +40,14 @@ async def test_read_and_commit_with_close_reader(self, driver, topic_with_messag
             assert message != message2
 
     async def test_read_and_commit_with_ack(self, driver, topic_with_messages, topic_consumer):
-        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
-        batch = await reader.receive_batch()
-        await reader.commit_with_ack(batch)
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            message = await reader.receive_message()
+            await reader.commit_with_ack(message)
 
-        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
-        batch2 = await reader.receive_batch()
-        assert batch.messages[0] != batch2.messages[0]
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            batch = await reader.receive_batch()
 
-        await reader.close()
+        assert message != batch.messages[0]
 
     async def test_read_compressed_messages(self, driver, topic_path, topic_consumer):
         async with driver.topic_client.writer(topic_path, codec=ydb.TopicCodec.GZIP) as writer:
@@ -147,12 +146,12 @@ def test_read_and_commit_with_close_reader(self, driver_sync, topic_with_message
 
     def test_read_and_commit_with_ack(self, driver_sync, topic_with_messages, topic_consumer):
         reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
-        batch = reader.receive_batch()
-        reader.commit_with_ack(batch)
+        message = reader.receive_message()
+        reader.commit_with_ack(message)
 
         reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
-        batch2 = reader.receive_batch()
-        assert batch.messages[0] != batch2.messages[0]
+        batch = reader.receive_batch()
+        assert message != batch.messages[0]
 
     def test_read_compressed_messages(self, driver_sync, topic_path, topic_consumer):
         with driver_sync.topic_client.writer(topic_path, codec=ydb.TopicCodec.GZIP) as writer:

tests/topics/test_topic_writer.py

@@ -41,10 +41,10 @@ async def test_random_producer_id(self, driver: ydb.aio.Driver, topic_path, topi
         async with driver.topic_client.writer(topic_path) as writer:
             await writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
-        batch1 = await topic_reader.receive_batch()
-        batch2 = await topic_reader.receive_batch()
+        msg1 = await topic_reader.receive_message()
+        msg2 = await topic_reader.receive_message()
 
-        assert batch1.messages[0].producer_id != batch2.messages[0].producer_id
+        assert msg1.producer_id != msg2.producer_id
 
     async def test_auto_flush_on_close(self, driver: ydb.aio.Driver, topic_path):
         async with driver.topic_client.writer(
@@ -77,18 +77,16 @@ async def test_write_multi_message_with_ack(
             assert res1.offset == 0
             assert res2.offset == 1
 
-        batch = await topic_reader.receive_batch()
+        msg1 = await topic_reader.receive_message()
+        msg2 = await topic_reader.receive_message()
 
-        assert batch.messages[0].offset == 0
-        assert batch.messages[0].seqno == 1
-        assert batch.messages[0].data == "123".encode()
+        assert msg1.offset == 0
+        assert msg1.seqno == 1
+        assert msg1.data == "123".encode()
 
-        # remove second recieve batch when implement batching
-        # https://github.com/ydb-platform/ydb-python-sdk/issues/142
-        batch = await topic_reader.receive_batch()
-        assert batch.messages[0].offset == 1
-        assert batch.messages[0].seqno == 2
-        assert batch.messages[0].data == "456".encode()
+        assert msg2.offset == 1
+        assert msg2.seqno == 2
+        assert msg2.data == "456".encode()
 
     @pytest.mark.parametrize(
         "codec",
@@ -186,10 +184,10 @@ def test_random_producer_id(
         with driver_sync.topic_client.writer(topic_path) as writer:
             writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
-        batch1 = topic_reader_sync.receive_batch()
-        batch2 = topic_reader_sync.receive_batch()
+        msg1 = topic_reader_sync.receive_message()
+        msg2 = topic_reader_sync.receive_message()
 
-        assert batch1.messages[0].producer_id != batch2.messages[0].producer_id
+        assert msg1.producer_id != msg2.producer_id
 
     def test_write_multi_message_with_ack(
         self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReader
@@ -202,18 +200,16 @@ def test_write_multi_message_with_ack(
                 ]
             )
 
-        batch = topic_reader_sync.receive_batch()
+        msg1 = topic_reader_sync.receive_message()
+        msg2 = topic_reader_sync.receive_message()
 
-        assert batch.messages[0].offset == 0
-        assert batch.messages[0].seqno == 1
-        assert batch.messages[0].data == "123".encode()
+        assert msg1.offset == 0
+        assert msg1.seqno == 1
+        assert msg1.data == "123".encode()
 
-        # remove second recieve batch when implement batching
-        # https://github.com/ydb-platform/ydb-python-sdk/issues/142
-        batch = topic_reader_sync.receive_batch()
-        assert batch.messages[0].offset == 1
-        assert batch.messages[0].seqno == 2
-        assert batch.messages[0].data == "456".encode()
+        assert msg2.offset == 1
+        assert msg2.seqno == 2
+        assert msg2.data == "456".encode()
 
     @pytest.mark.parametrize(
         "codec",

ydb/_topic_reader/datatypes.py

@@ -7,7 +7,7 @@
 from collections import deque
 from dataclasses import dataclass, field
 import datetime
-from typing import Union, Any, List, Dict, Deque, Optional
+from typing import Union, Any, List, Dict, Deque, Optional, Tuple
 
 from ydb._grpc.grpcwrapper.ydb_topic import OffsetsRange, Codec
 from ydb._topic_reader import topic_reader_asyncio
@@ -171,3 +171,11 @@ def alive(self) -> bool:
 
     def pop_message(self) -> PublicMessage:
         return self.messages.pop(0)
+
+    def _extend(self, batch: PublicBatch) -> None:
+        self.messages.extend(batch.messages)
+        self._bytes_size += batch._bytes_size
+
+    def _pop(self) -> Tuple[List[PublicMessage], bool]:
+        msgs_left = True if len(self.messages) > 1 else False
+        return self.messages.pop(0), msgs_left

ydb/_topic_reader/topic_reader_asyncio.py

@@ -5,7 +5,7 @@
 import gzip
 import typing
 from asyncio import Task
-from collections import deque
+from collections import OrderedDict
 from typing import Optional, Set, Dict, Union, Callable
 
 import ydb
@@ -266,7 +266,7 @@ class ReaderStream:
 
     _state_changed: asyncio.Event
     _closed: bool
-    _message_batches: typing.Deque[datatypes.PublicBatch]
+    _message_batches: typing.Dict[int, datatypes.PublicBatch]  # keys are partition session ID
     _first_error: asyncio.Future[YdbError]
 
     _update_token_interval: Union[int, float]
@@ -298,7 +298,7 @@ def __init__(
         self._closed = False
         self._first_error = asyncio.get_running_loop().create_future()
         self._batches_to_decode = asyncio.Queue()
-        self._message_batches = deque()
+        self._message_batches = OrderedDict()
 
         self._update_token_interval = settings.update_token_interval
         self._get_token_function = get_token_function
@@ -379,29 +379,38 @@ async def wait_messages(self):
             await self._state_changed.wait()
             self._state_changed.clear()
 
+    def _get_first_batch(self) -> typing.Tuple[int, datatypes.PublicBatch]:
+        partition_session_id, batch = self._message_batches.popitem(last=False)
+        return partition_session_id, batch
+
     def receive_batch_nowait(self):
         if self._get_first_error():
             raise self._get_first_error()
 
         if not self._message_batches:
             return None
 
-        batch = self._message_batches.popleft()
+        _, batch = self._get_first_batch()
         self._buffer_release_bytes(batch._bytes_size)
+
         return batch
 
     def receive_message_nowait(self):
         if self._get_first_error():
             raise self._get_first_error()
 
-        try:
-            batch = self._message_batches[0]
-            message = batch.pop_message()
-        except IndexError:
+        if not self._message_batches:
             return None
 
-        if batch.empty():
-            self.receive_batch_nowait()
+        part_sess_id, batch = self._get_first_batch()
+
+        message, msgs_left = batch._pop()
+
+        if not msgs_left:
+            self._buffer_release_bytes(batch._bytes_size)
+        else:
+            # TODO: we should somehow release bytes from single message as well
+            self._message_batches[part_sess_id] = batch
 
         return message
 
@@ -625,9 +634,17 @@ async def _decode_batches_loop(self):
         while True:
             batch = await self._batches_to_decode.get()
             await self._decode_batch_inplace(batch)
-            self._message_batches.append(batch)
+            self._add_batch_to_queue(batch)
             self._state_changed.set()
 
+    def _add_batch_to_queue(self, batch: datatypes.PublicBatch):
+        part_sess_id = batch._partition_session.id
+        if part_sess_id in self._message_batches:
+            self._message_batches[part_sess_id]._extend(batch)
+            return
+
+        self._message_batches[part_sess_id] = batch
+
     async def _decode_batch_inplace(self, batch):
         if batch._codec == Codec.CODEC_RAW:
             return