topic reader: add flush method

Author: Valeriya Popova

ydb/_topic_reader/datatypes.py

@@ -96,7 +96,7 @@ def add_waiter(self, end_offset: int) -> "PartitionSession.CommitAckWaiter":
             return waiter
 
         # fast way
-        if len(self._ack_waiters) > 0 and self._ack_waiters[-1].end_offset < end_offset:
+        if self._ack_waiters and self._ack_waiters[-1].end_offset < end_offset:
             self._ack_waiters.append(waiter)
         else:
             bisect.insort(self._ack_waiters, waiter)
@@ -106,25 +106,23 @@ def add_waiter(self, end_offset: int) -> "PartitionSession.CommitAckWaiter":
     def _create_future(self) -> asyncio.Future:
         if self._loop:
             return self._loop.create_future()
-        else:
-            return asyncio.Future()
+        return asyncio.Future()
 
     def ack_notify(self, offset: int):
         self._ensure_not_closed()
 
         self.committed_offset = offset
 
-        if len(self._ack_waiters) == 0:
+        if not self._ack_waiters:
             # todo log warning
             # must be never receive ack for not sended request
             return
 
-        while len(self._ack_waiters) > 0:
-            if self._ack_waiters[0].end_offset <= offset:
-                waiter = self._ack_waiters.popleft()
-                waiter._finish_ok()
-            else:
+        while self._ack_waiters:
+            if self._ack_waiters[0].end_offset > offset:
                 break
+            waiter = self._ack_waiters.popleft()
+            waiter._finish_ok()
 
     def close(self):
         if self.closed:

ydb/_topic_reader/topic_reader_asyncio.py

@@ -120,12 +120,20 @@ async def commit_with_ack(
         waiter = self._reconnector.commit(batch)
         await waiter.future
 
-    async def close(self):
+    async def flush(self):
+        """
+        force send all commit messages from internal buffers to server and wait acks for all of them.
+
+        use asyncio.wait_for for wait with timeout.
+        """
+        await self._reconnector.flush()
+
+    async def close(self, flush: bool = True):
         if self._closed:
             raise TopicReaderClosedError()
 
         self._closed = True
-        await self._reconnector.close()
+        await self._reconnector.close(flush)
 
 
 class ReaderReconnector:
@@ -199,14 +207,20 @@ def commit(
     ) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)
 
-    async def close(self):
+    async def close(self, flush: bool):
         if self._stream_reader:
+            if flush:
+                await self.flush()
             await self._stream_reader.close()
         for task in self._background_tasks:
             task.cancel()
 
         await asyncio.wait(self._background_tasks)
 
+    async def flush(self):
+        if self._stream_reader:
+            await self._stream_reader.flush()
+
     def _set_first_error(self, err: issues.Error):
         try:
             self._first_error.set_result(err)
@@ -641,6 +655,16 @@ def _get_first_error(self) -> Optional[YdbError]:
         if self._first_error.done():
             return self._first_error.result()
 
+    async def flush(self):
+        if self._closed:
+            raise RuntimeError("Flush on closed Stream")
+
+        futures = []
+        for session in self._partition_sessions.values():
+            futures.extend(w.future for w in session._ack_waiters)
+
+        await asyncio.gather(*futures)
+
     async def close(self):
         if self._closed:
             return

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -389,6 +389,35 @@ async def test_close_ack_waiters_when_close_stream_reader(
         with pytest.raises(topic_reader_asyncio.TopicReaderCommitToExpiredPartition):
             waiter.future.result()
 
+    async def test_flush(
+        self, stream, stream_reader_started: ReaderStream, partition_session
+    ):
+        offset = self.partition_session_committed_offset + 1
+        waiter = partition_session.add_waiter(offset)
+
+        with pytest.raises(WaitConditionError):
+            await wait_for_fast(stream_reader_started.flush(), timeout=0.1)
+
+        stream.from_server.put_nowait(
+            StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.CommitOffsetResponse(
+                    partitions_committed_offsets=[
+                        StreamReadMessage.CommitOffsetResponse.PartitionCommittedOffset(
+                            partition_session_id=partition_session.id,
+                            committed_offset=offset,
+                        ),
+                    ]
+                ),
+            )
+        )
+
+        await stream_reader_started.flush()
+        # don't raises
+        assert waiter.future.result() is None
+
+        await wait_for_fast(stream_reader_started.close())
+
     async def test_commit_ranges_for_received_messages(
         self, stream, stream_reader_started: ReaderStream, partition_session
     ):