Merge pull request #164 from ydb-platform/topic-writer-flush

topic writer add flush

Author: Valeria1235

examples/topic/writer_async_example.py

@@ -65,7 +65,7 @@ async def send_message_without_block_if_internal_buffer_is_full(
         return False
 
 
-def send_messages_with_manual_seqno(writer: ydb.TopicWriter):
+async def send_messages_with_manual_seqno(writer: ydb.TopicWriter):
     await writer.write(ydb.TopicWriterMessage("mess"))  # send text
 
 

tests/topics/test_topic_writer.py

@@ -9,7 +9,8 @@ async def test_send_message(self, driver: ydb.aio.Driver, topic_path):
         writer = driver.topic_client.topic_writer(
             topic_path, producer_and_message_group_id="test"
         )
-        writer.write(ydb.TopicWriterMessage(data="123".encode()))
+        await writer.write(ydb.TopicWriterMessage(data="123".encode()))
+        await writer.close()
 
     async def test_wait_last_seqno(self, driver: ydb.aio.Driver, topic_path):
         async with driver.topic_client.topic_writer(
@@ -28,3 +29,24 @@ async def test_wait_last_seqno(self, driver: ydb.aio.Driver, topic_path):
         ) as writer2:
             init_info = await writer2.wait_init()
             assert init_info.last_seqno == 5
+
+    async def test_auto_flush_on_close(self, driver: ydb.aio.Driver, topic_path):
+        async with driver.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            auto_seqno=False,
+        ) as writer:
+            last_seqno = 0
+            for i in range(10):
+                last_seqno = i + 1
+                await writer.write(
+                    ydb.TopicWriterMessage(data=f"msg-{i}", seqno=last_seqno)
+                )
+
+        async with driver.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            get_last_seqno=True,
+        ) as writer:
+            init_info = await writer.wait_init()
+            assert init_info.last_seqno == last_seqno

ydb/_topic_writer/topic_writer.py

@@ -104,7 +104,7 @@ def async_flush(self):
         """
         raise NotImplementedError()
 
-    def flush(self, timeout: Union[float, None] = None) -> concurrent.futures.Future:
+    def flush(self, timeout: Optional[float] = None) -> concurrent.futures.Future:
         """
         Force send all messages from internal buffer and wait acks from server for all
         messages.
@@ -122,7 +122,7 @@ def async_wait_init(self) -> concurrent.futures.Future:
         """
         raise NotImplementedError()
 
-    def wait_init(self, timeout: Union[float, None] = None):
+    def wait_init(self, timeout: Optional[float] = None):
         """
         Wait until underling connection established
 
@@ -141,15 +141,15 @@ class PublicWriterSettings:
     session_metadata: Optional[Dict[str, str]] = None
     encoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
     serializer: Union[Callable[[Any], bytes], None] = None
-    send_buffer_count: Union[int, None] = 10000
-    send_buffer_bytes: Union[int, None] = 100 * 1024 * 1024
+    send_buffer_count: Optional[int] = 10000
+    send_buffer_bytes: Optional[int] = 100 * 1024 * 1024
     partition_id: Optional[int] = None
-    codec: Union[int, None] = None
+    codec: Optional[int] = None
     codec_autoselect: bool = True
     auto_seqno: bool = True
     auto_created_at: bool = True
     get_last_seqno: bool = False
-    retry_policy: Union["RetryPolicy", None] = None
+    retry_policy: Optional["RetryPolicy"] = None
     update_token_interval: Union[int, float] = 3600
 
 
@@ -251,7 +251,7 @@ def to_message_data(self) -> StreamWriteMessage.WriteRequest.MessageData:
 
 
 class MessageSendResult:
-    offset: Union[None, int]
+    offset: Optional[int]
     write_status: "MessageWriteStatus"
 
 

ydb/_topic_writer/topic_writer_asyncio.py

@@ -86,21 +86,14 @@ async def write_with_ack(
 
         For wait with timeout use asyncio.wait_for.
         """
-        if isinstance(messages, PublicMessage):
-            futures = await self._reconnector.write_with_ack([messages])
-            return await futures[0]
-        if isinstance(messages, list):
-            for m in messages:
-                if not isinstance(m, PublicMessage):
-                    raise NotImplementedError()
+        futures = await self.write_with_ack_future(messages, *args)
+        if not isinstance(futures, list):
+            futures = [futures]
 
-            futures = await self._reconnector.write_with_ack(messages)
-            await asyncio.wait(futures)
+        await asyncio.wait(futures)
+        results = [f.result() for f in futures]
 
-            results = [f.result() for f in futures]
-            return results
-
-        raise NotImplementedError()
+        return results if isinstance(messages, list) else results[0]
 
     async def write_with_ack_future(
         self,
@@ -145,7 +138,7 @@ async def flush(self):
 
         For wait with timeout use asyncio.wait_for.
         """
-        raise NotImplementedError()
+        return await self._reconnector.flush()
 
     async def wait_init(self) -> PublicWriterInitInfo:
         """
@@ -192,10 +185,13 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
             asyncio.create_task(self._connection_loop(), name="connection_loop")
         ]
 
-    async def close(self):
+    async def close(self, flush: bool = True):
         if self._closed:
             return
 
+        if flush:
+            await self.flush()
+
         self._closed = True
         self._stop(TopicWriterStopped())
 
@@ -396,6 +392,14 @@ def _stop(self, reason: Exception):
     def _get_token(self) -> str:
         raise NotImplementedError()
 
+    async def flush(self):
+        self._check_stop()
+        if not self._messages_future:
+            return
+
+        # wait last message
+        await asyncio.wait((self._messages_future[-1],))
+
 
 class WriterAsyncIOStream:
     # todo slots

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -239,6 +239,20 @@ def default_write_statistic(
             topic_quota_wait_time=datetime.timedelta(milliseconds=5),
         )
 
+    def make_default_ack_message(self, seq_no=1) -> StreamWriteMessage.WriteResponse:
+        return StreamWriteMessage.WriteResponse(
+            partition_id=1,
+            acks=[
+                StreamWriteMessage.WriteResponse.WriteAck(
+                    seq_no=seq_no,
+                    message_write_status=StreamWriteMessage.WriteResponse.WriteAck.StatusWritten(
+                        offset=1
+                    ),
+                )
+            ],
+            write_statistics=self.default_write_statistic,
+        )
+
     @pytest.fixture
     async def reconnector(
         self, default_driver, default_settings
@@ -275,20 +289,7 @@ async def test_reconnect_and_resent_non_acked_messages_on_retriable_error(
         assert [InternalMessage(message2)] == messages
 
         # ack first message
-        stream_writer.from_server.put_nowait(
-            StreamWriteMessage.WriteResponse(
-                partition_id=1,
-                acks=[
-                    StreamWriteMessage.WriteResponse.WriteAck(
-                        seq_no=1,
-                        message_write_status=StreamWriteMessage.WriteResponse.WriteAck.StatusWritten(
-                            offset=1
-                        ),
-                    )
-                ],
-                write_statistics=default_write_statistic,
-            )
-        )
+        stream_writer.from_server.put_nowait(self.make_default_ack_message(seq_no=1))
 
         stream_writer.from_server.put_nowait(issues.Overloaded("test"))
 
@@ -297,6 +298,8 @@ async def test_reconnect_and_resent_non_acked_messages_on_retriable_error(
 
         expected_messages = [InternalMessage(message2)]
         assert second_sent_msg == expected_messages
+
+        second_writer.from_server.put_nowait(self.make_default_ack_message(seq_no=2))
         await reconnector.close()
 
     async def test_stop_on_unexpected_exception(
@@ -323,7 +326,7 @@ async def wait_stop():
             await asyncio.wait_for(wait_stop(), 1)
 
         with pytest.raises(TestException):
-            await reconnector.close()
+            await reconnector.close(flush=False)
 
     async def test_wait_init(self, default_driver, default_settings, get_stream_writer):
         init_seqno = 100
@@ -350,7 +353,7 @@ async def test_wait_init(self, default_driver, default_settings, get_stream_writ
             info = await reconnector.wait_init()
             assert info == expected_init_info
 
-        await reconnector.close()
+        await reconnector.close(flush=False)
 
     async def test_write_message(
         self, reconnector: WriterAsyncIOReconnector, get_stream_writer
@@ -365,7 +368,7 @@ async def test_write_message(
         sent_messages = await asyncio.wait_for(stream_writer.from_client.get(), 1)
         assert sent_messages == [InternalMessage(message)]
 
-        await reconnector.close()
+        await reconnector.close(flush=False)
 
     async def test_auto_seq_no(
         self, default_driver, default_settings, get_stream_writer
@@ -399,7 +402,7 @@ async def test_auto_seq_no(
                 [PublicMessage(seqno=last_seq_no + 3, data="123")]
             )
 
-        await reconnector.close()
+        await reconnector.close(flush=False)
 
     async def test_deny_double_seqno(self, reconnector: WriterAsyncIOReconnector):
         await reconnector.write_with_ack([PublicMessage(seqno=10, data="123")])
@@ -412,7 +415,7 @@ async def test_deny_double_seqno(self, reconnector: WriterAsyncIOReconnector):
 
         await reconnector.write_with_ack([PublicMessage(seqno=11, data="123")])
 
-        await reconnector.close()
+        await reconnector.close(flush=False)
 
     @freezegun.freeze_time("2022-01-13 20:50:00", tz_offset=0)
     async def test_auto_created_at(
@@ -431,7 +434,7 @@ async def test_auto_created_at(
         assert [
             InternalMessage(PublicMessage(seqno=4, data="123", created_at=now))
         ] == sent
-        await reconnector.close()
+        await reconnector.close(flush=False)
 
 
 @pytest.mark.asyncio