test autosplit message order

ydb/_topic_reader/topic_reader_asyncio.py

@@ -394,7 +394,6 @@ def _return_batch_to_queue(self, part_sess_id: int, batch: datatypes.PublicBatch
         # In case of auto-split we should return all parent messages ASAP
         # without queue rotation to prevent child's messages before parent's.
         if part_sess_id in self._partition_sessions and self._partition_sessions[part_sess_id].ended:
-            print(f"part_sess_id: {part_sess_id} is ended, return to beginning of queue")
             self._message_batches.move_to_end(part_sess_id, last=False)
 
     def receive_batch_nowait(self, max_messages: Optional[int] = None):

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -52,8 +52,8 @@ def default_executor():
     executor.shutdown()
 
 
-def stub_partition_session(id: int = 0):
-    return datatypes.PartitionSession(
+def stub_partition_session(id: int = 0, ended: bool = False):
+    partition_session = datatypes.PartitionSession(
         id=id,
         state=datatypes.PartitionSession.State.Active,
         topic_path="asd",
@@ -63,6 +63,11 @@ def stub_partition_session(id: int = 0):
         reader_stream_id=513,
     )
 
+    if ended:
+        partition_session.end()
+
+    return partition_session
+
 
 def stub_message(id: int):
     return PublicMessage(
@@ -73,7 +78,7 @@ def stub_message(id: int):
         offset=0,
         written_at=datetime.datetime(2023, 3, 18, 14, 15),
         producer_id="",
-        data=bytes(),
+        data=id,
         metadata_items={},
         _partition_session=stub_partition_session(),
         _commit_start_offset=0,
@@ -746,6 +751,31 @@ def session_count():
         with pytest.raises(asyncio.QueueEmpty):
             stream.from_client.get_nowait()
 
+    async def test_end_partition_session(self, stream, stream_reader, partition_session):
+        def session_count():
+            return len(stream_reader._partition_sessions)
+
+        initial_session_count = session_count()
+
+        stream.from_server.put_nowait(
+            StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.EndPartitionSession(
+                    partition_session_id=partition_session.id,
+                    adjacent_partition_ids=[],
+                    child_partition_ids=[20, 30],
+                ),
+            )
+        )
+
+        await asyncio.sleep(0)  # wait next loop
+        with pytest.raises(asyncio.QueueEmpty):
+            stream.from_client.get_nowait()
+
+        assert session_count() == initial_session_count
+        assert partition_session.id in stream_reader._partition_sessions
+        assert partition_session.ended
+
     @pytest.mark.parametrize(
         "graceful",
         (
@@ -1168,6 +1198,82 @@ async def test_read_message(
         assert mess == expected_message
         assert dict(stream_reader._message_batches) == batches_after
 
+    @pytest.mark.parametrize(
+        "batches,expected_order",
+        [
+            (
+                {
+                    0: PublicBatch(
+                        messages=[stub_message(1)],
+                        _partition_session=stub_partition_session(0, ended=True),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    )
+                },
+                [1],
+            ),
+            (
+                {
+                    0: PublicBatch(
+                        messages=[stub_message(1), stub_message(2)],
+                        _partition_session=stub_partition_session(0, ended=True),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    ),
+                    1: PublicBatch(
+                        messages=[stub_message(3), stub_message(4)],
+                        _partition_session=stub_partition_session(1),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    ),
+                },
+                [1, 2, 3, 4],
+            ),
+            (
+                {
+                    0: PublicBatch(
+                        messages=[stub_message(1), stub_message(2)],
+                        _partition_session=stub_partition_session(0),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    ),
+                    1: PublicBatch(
+                        messages=[stub_message(3), stub_message(4)],
+                        _partition_session=stub_partition_session(1, ended=True),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    ),
+                    2: PublicBatch(
+                        messages=[stub_message(5)],
+                        _partition_session=stub_partition_session(2),
+                        _bytes_size=0,
+                        _codec=Codec.CODEC_RAW,
+                    ),
+                },
+                [1, 3, 4, 5, 2],
+            ),
+        ],
+    )
+    async def test_read_message_autosplit_order(
+        self,
+        stream_reader,
+        batches: typing.Dict[int, datatypes.PublicBatch],
+        expected_order: typing.List[int],
+    ):
+        stream_reader._message_batches = OrderedDict(batches)
+
+        for id, batch in batches.items():
+            ps = batch._partition_session
+            stream_reader._partition_sessions[id] = ps
+
+        result = []
+        for _ in range(len(expected_order)):
+            mess = stream_reader.receive_message_nowait()
+            result.append(mess.data)
+
+        assert result == expected_order
+        assert stream_reader.receive_message_nowait() is None
+
     @pytest.mark.parametrize(
         "batches_before,max_messages,actual_messages,batches_after",
         [