Merge pull request #248 sync-from-v3

Author: rekby

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

@@ -75,7 +75,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     def __del__(self):
         if not self._closed:
-            self._loop.create_task(self.close(), name="close reader")
+            self._loop.create_task(self.close(flush=False), name="close reader")
 
     async def receive_batch(
         self,
@@ -120,12 +120,12 @@ async def commit_with_ack(
         waiter = self._reconnector.commit(batch)
         await waiter.future
 
-    async def close(self):
+    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 +199,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 +647,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
     ):

ydb/_topic_reader/topic_reader_sync.py

@@ -50,7 +50,7 @@ async def create_reader():
         ).result()
 
     def __del__(self):
-        self.close()
+        self.close(flush=False)
 
     def __enter__(self):
         return self
@@ -152,13 +152,13 @@ def async_commit_with_ack(
             self._async_reader.commit_with_ack(mess)
         )
 
-    def close(self, *, timeout: TimeoutType = None):
+    def close(self, *, flush: bool = True, timeout: TimeoutType = None):
         if self._closed:
             return
 
         self._closed = True
 
-        self._caller.safe_call_with_result(self._async_reader.close(), timeout)
+        self._caller.safe_call_with_result(self._async_reader.close(flush), timeout)
 
     def _check_closed(self):
         if self._closed:

ydb/_topic_writer/topic_writer_asyncio.py

@@ -6,6 +6,8 @@
 from collections import deque
 from typing import Deque, AsyncIterator, Union, List, Optional, Dict, Callable
 
+import logging
+
 import ydb
 from .topic_writer import (
     PublicWriterSettings,
@@ -38,6 +40,8 @@
     GrpcWrapperAsyncIO,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class WriterAsyncIO:
     _loop: asyncio.AbstractEventLoop
@@ -61,7 +65,7 @@ def __del__(self):
         if self._closed or self._loop.is_closed():
             return
 
-        self._loop.call_soon(self.close)
+        self._loop.call_soon(self.close, False)
 
     async def close(self, *, flush: bool = True):
         if self._closed:
@@ -154,7 +158,6 @@ class WriterAsyncIOReconnector:
     _credentials: Union[ydb.credentials.Credentials, None]
     _driver: ydb.aio.Driver
     _init_message: StreamWriteMessage.InitRequest
-    _init_info: asyncio.Future
     _stream_connected: asyncio.Event
     _settings: WriterSettings
     _codec: PublicCodec
@@ -164,25 +167,30 @@ class WriterAsyncIOReconnector:
     _codec_selector_last_codec: Optional[PublicCodec]
     _codec_selector_check_batches_interval: int
 
-    _last_known_seq_no: int
     if typing.TYPE_CHECKING:
         _messages_for_encode: asyncio.Queue[List[InternalMessage]]
     else:
         _messages_for_encode: asyncio.Queue
     _messages: Deque[InternalMessage]
     _messages_future: Deque[asyncio.Future]
     _new_messages: asyncio.Queue
-    _stop_reason: asyncio.Future
     _background_tasks: List[asyncio.Task]
 
+    _state_changed: asyncio.Event
+    if typing.TYPE_CHECKING:
+        _stop_reason: asyncio.Future[BaseException]
+    else:
+        _stop_reason: asyncio.Future
+    _init_info: Optional[PublicWriterInitInfo]
+
     def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
         self._closed = False
         self._loop = asyncio.get_running_loop()
         self._driver = driver
         self._credentials = driver._credentials
         self._init_message = settings.create_init_request()
         self._new_messages = asyncio.Queue()
-        self._init_info = self._loop.create_future()
+        self._init_info = None
         self._stream_connected = asyncio.Event()
         self._settings = settings
 
@@ -219,14 +227,17 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
             asyncio.create_task(self._encode_loop(), name="encode_loop"),
         ]
 
+        self._state_changed = asyncio.Event()
+
     async def close(self, flush: bool):
         if self._closed:
             return
+        self._closed = True
+        logger.debug("Close writer reconnector")
 
         if flush:
             await self.flush()
 
-        self._closed = True
         self._stop(TopicWriterStopped())
 
         for task in self._background_tasks:
@@ -240,19 +251,20 @@ async def close(self, flush: bool):
             pass
 
     async def wait_init(self) -> PublicWriterInitInfo:
-        done, _ = await asyncio.wait(
-            [self._init_info, self._stop_reason], return_when=asyncio.FIRST_COMPLETED
-        )
-        res = done.pop()  # type: asyncio.Future
-        res_val = res.result()
+        while True:
+            if self._stop_reason.done():
+                raise self._stop_reason.exception()
 
-        if isinstance(res_val, BaseException):
-            raise res_val
+            if self._init_info:
+                return self._init_info
 
-        return res_val
+            await self._state_changed.wait()
 
-    async def wait_stop(self) -> Exception:
-        return await self._stop_reason
+    async def wait_stop(self) -> BaseException:
+        try:
+            await self._stop_reason
+        except BaseException as stop_reason:
+            return stop_reason
 
     async def write_with_ack_future(
         self, messages: List[PublicMessage]
@@ -343,13 +355,14 @@ async def _connection_loop(self):
                     self._settings.update_token_interval,
                 )
                 try:
-                    self._last_known_seq_no = stream_writer.last_seqno
-                    self._init_info.set_result(
-                        PublicWriterInitInfo(
+                    if self._init_info is None:
+                        self._last_known_seq_no = stream_writer.last_seqno
+                        self._init_info = PublicWriterInitInfo(
                             last_seqno=stream_writer.last_seqno,
                             supported_codecs=stream_writer.supported_codecs,
                         )
-                    )
+                        self._state_changed.set()
+
                 except asyncio.InvalidStateError:
                     pass
 
@@ -369,9 +382,6 @@ async def _connection_loop(self):
                 await stream_writer.close()
                 done.pop().result()
             except issues.Error as err:
-                # todo log error
-                print(err)
-
                 err_info = check_retriable_error(err, retry_settings, attempt)
                 if not err_info.is_retriable:
                     self._stop(err)
@@ -550,8 +560,13 @@ def _stop(self, reason: Exception):
 
         self._stop_reason.set_result(reason)
 
+        for f in self._messages_future:
+            f.set_exception(reason)
+
+        self._state_changed.set()
+        logger.info("Stop topic writer: %s" % reason)
+
     async def flush(self):
-        self._check_stop()
         if not self._messages_future:
             return
 

ydb/_topic_writer/topic_writer_sync.py

@@ -56,6 +56,9 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
+    def __del__(self):
+        self.close(flush=False)
+
     def close(self, *, flush: bool = True, timeout: TimeoutType = None):
         if self._closed:
             return

ydb/topic.py

@@ -208,7 +208,7 @@ def close(self):
             return
 
         self._closed = True
-        self._executor.shutdown(wait=False, cancel_futures=True)
+        self._executor.shutdown(wait=False)
 
     def _check_closed(self):
         if not self._closed:
@@ -366,7 +366,7 @@ def close(self):
             return
 
         self._closed = True
-        self._executor.shutdown(wait=False, cancel_futures=True)
+        self._executor.shutdown(wait=False)
 
     def _check_closed(self):
         if not self._closed: