sync

Author: rekby

ydb/_topic_reader/topic_reader_asyncio.py

@@ -6,9 +6,7 @@
 from collections import deque
 from typing import Optional, Set, Dict
 
-import grpc
 
-import ydb
 from .. import _apis, issues, RetrySettings
 from ..aio import Driver
 from ..issues import (
@@ -53,7 +51,7 @@ class ReaderReconnector:
 
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
-    _first_error: asyncio.Future[ydb.Error]
+    _first_error: asyncio.Future[YdbError]
 
     def __init__(self, driver: Driver, settings: PublicReaderSettings):
         self._settings = settings
@@ -71,11 +69,10 @@ async def _connection_loop(self):
         while True:
             try:
                 self._stream_reader = await ReaderStream.create(self._driver, self._settings)
+                attempt = 0
                 self._state_changed.set()
-                self._stream_reader._state_changed.wait()
-            except Exception as err:
-                # todo reset attempts when connection established
-
+                await self._stream_reader.wait_error()
+            except issues.Error as err:
                 retry_info = check_retriable_error(err, self._settings._retry_settings(), attempt)
                 if not retry_info.is_retriable:
                     self._set_first_error(err)
@@ -90,8 +87,11 @@ async def wait_message(self):
                 raise self._first_error.result()
 
             if self._stream_reader is not None:
-                await self._stream_reader.wait_messages()
-                return
+                try:
+                    await self._stream_reader.wait_messages()
+                    return
+                except YdbError:
+                    pass  # handle errors in reconnection loop
 
             await self._state_changed.wait()
             self._state_changed.clear()
@@ -114,6 +114,7 @@ def _set_first_error(self, err: issues.Error):
             # skip if already has result
             pass
 
+
 class ReaderStream:
     _token_getter: Optional[TokenGetterFuncType]
     _session_id: str
@@ -126,7 +127,7 @@ class ReaderStream:
     _state_changed: asyncio.Event
     _closed: bool
     _message_batches: typing.Deque[PublicBatch]
-    first_error: asyncio.Future[YdbError]
+    _first_error: asyncio.Future[YdbError]
 
     def __init__(self, settings: PublicReaderSettings):
         self._token_getter = settings._token_getter
@@ -139,7 +140,7 @@ def __init__(self, settings: PublicReaderSettings):
 
         self._state_changed = asyncio.Event()
         self._closed = False
-        self.first_error = asyncio.get_running_loop().create_future()
+        self._first_error = asyncio.get_running_loop().create_future()
         self._message_batches = deque()
 
     @staticmethod
@@ -174,6 +175,9 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
         read_messages_task = asyncio.create_task(self._read_messages_loop(stream))
         self._background_tasks.add(read_messages_task)
 
+    async def wait_error(self):
+        raise await self._first_error
+
     async def wait_messages(self):
         while True:
             if self._get_first_error() is not None:
@@ -317,17 +321,17 @@ def _read_response_to_batches(self, message: StreamReadMessage.ReadResponse) ->
         batches[-1]._bytes_size += additional_bytes_to_last_batch
         return batches
 
-    def _set_first_error(self, err):
+    def _set_first_error(self, err: ydb.Error):
         try:
-            self.first_error.set_result(err)
+            self._first_error.set_result(err)
             self._state_changed.set()
         except asyncio.InvalidStateError:
             # skip later set errors
             pass
 
-    def _get_first_error(self):
-        if self.first_error.done():
-            return self.first_error.result()
+    def _get_first_error(self) -> Optional[ydb.Error]:
+        if self._first_error.done():
+            return self._first_error.result()
         else:
             return None
 

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -10,7 +10,7 @@
 from .datatypes import PublicBatch, PublicMessage
 from .topic_reader import PublicReaderSettings
 from .topic_reader_asyncio import ReaderStream, PartitionSession, ReaderReconnector
-from .._topic_wrapper.common import OffsetsRange, Codec, ServerStatus, UpdateTokenResponse
+from .._topic_wrapper.common import OffsetsRange, Codec, ServerStatus, UpdateTokenResponse, SupportedDriverType
 from .._topic_wrapper.reader import StreamReadMessage
 from .._topic_wrapper.test_helpers import StreamMock, wait_condition, wait_for_fast
 from ..issues import Unavailable
@@ -23,6 +23,13 @@
     from .._grpc.common.protos import ydb_status_codes_pb2
 
 
+@pytest.fixture(autouse=True)
+def handle_exceptions(event_loop):
+    def handler(loop, context):
+        print(context)
+    event_loop.set_exception_handler(handler)
+
+
 @pytest.fixture()
 def default_reader_settings():
     return PublicReaderSettings(
@@ -634,11 +641,38 @@ async def test_receive_batch_nowait(self, stream, stream_reader, partition_sessi
 @pytest.mark.asyncio
 class TestReaderReconnector:
     async def test_reconnect_on_repeatable_error(self, monkeypatch):
-        def stream_create():
-            pass
+        test_error = issues.Overloaded("test error")
+
+        async def wait_error():
+            raise test_error
+
+        reader_stream_mock_with_error = mock.Mock(ReaderStream)
+        reader_stream_mock_with_error.wait_error = mock.AsyncMock(side_effect=wait_error)
+
+        async def wait_messages():
+            raise test_error
+
+        reader_stream_mock_with_error.wait_messages = mock.AsyncMock(side_effect=wait_messages)
+
+        reader_stream_with_messages = mock.Mock(ReaderStream)
+        reader_stream_with_messages.wait_error.return_value = asyncio.Future()
+        reader_stream_with_messages.wait_messages.return_value = None
+
+        stream_index = 0
+
+        async def stream_create(driver: SupportedDriverType, settings: PublicReaderSettings,):
+            nonlocal stream_index
+            stream_index += 1
+            if stream_index == 1:
+                return reader_stream_mock_with_error
+            elif stream_index == 2:
+                return reader_stream_with_messages
+            else:
+                raise Exception("unexpected create stream")
 
         with mock.patch.object(ReaderStream, "create", stream_create):
-            reconnector = ReaderReconnector(None, PublicReaderSettings("", ""))
+            reconnector = ReaderReconnector(mock.Mock(), PublicReaderSettings("", ""))
             await reconnector.wait_message()
 
-        raise NotImplementedError()
+        reader_stream_mock_with_error.wait_error.assert_any_await()
+        reader_stream_mock_with_error.wait_messages.assert_any_await()