sync

Author: rekby

ydb/_topic_reader/topic_reader.py

@@ -16,6 +16,7 @@
     Any, Dict,
 )
 
+from ydb import RetrySettings
 from ydb._topic_wrapper.common import OffsetsRange, TokenGetterFuncType
 from ydb._topic_wrapper.reader import StreamReadMessage
 
@@ -281,6 +282,9 @@ def _init_message(self) -> StreamReadMessage.InitRequest:
             consumer=self.consumer,
         )
 
+    def _retry_settings(self)->RetrySettings:
+        return RetrySettings(idempotent=True)
+
 
 class Events:
     class OnCommit:

ydb/_topic_reader/topic_reader_asyncio.py

@@ -8,13 +8,18 @@
 
 import grpc
 
-from .. import _apis, issues
+import ydb
+from .. import _apis, issues, RetrySettings
 from ..aio import Driver
-from ..issues import Error as YdbError
+from ..issues import (
+    Error as YdbError,
+    _process_response
+)
 from .datatypes import PartitionSession, PublicMessage, PublicBatch
 from .topic_reader import PublicReaderSettings
 from .._topic_wrapper.common import TokenGetterFuncType, IGrpcWrapperAsyncIO, SupportedDriverType, GrpcWrapperAsyncIO
 from .._topic_wrapper.reader import StreamReadMessage
+from .._errors import check_retriable_error
 
 
 class TopicReaderError(YdbError):
@@ -48,22 +53,42 @@ class ReaderReconnector:
 
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
+    _first_error: asyncio.Future[ydb.Error]
 
     def __init__(self, driver: Driver, settings: PublicReaderSettings):
         self._settings = settings
         self._driver = driver
         self._background_tasks = set()
+        self._retry_settins = RetrySettings(idempotent=True)  # get from settings
 
         self._state_changed = asyncio.Event()
         self._stream_reader = None
-        self._background_tasks.add(asyncio.create_task(self.start()))
+        self._background_tasks.add(asyncio.create_task(self._connection_loop()))
+        self._first_error = asyncio.get_running_loop().create_future()
 
-    async def start(self):
-        self._stream_reader = await ReaderStream.create(self._driver, self._settings)
-        self._state_changed.set()
+    async def _connection_loop(self):
+        attempt = 0
+        while True:
+            try:
+                self._stream_reader = await ReaderStream.create(self._driver, self._settings)
+                self._state_changed.set()
+                self._stream_reader._state_changed.wait()
+            except Exception as err:
+                # todo reset attempts when connection established
+
+                retry_info = check_retriable_error(err, self._settings._retry_settings(), attempt)
+                if not retry_info.is_retriable:
+                    self._set_first_error(err)
+                    return
+                await asyncio.sleep(retry_info.sleep_timeout_seconds)
+
+                attempt += 1
 
     async def wait_message(self):
         while True:
+            if self._first_error.done():
+                raise self._first_error.result()
+
             if self._stream_reader is not None:
                 await self._stream_reader.wait_messages()
                 return
@@ -81,6 +106,13 @@ async def close(self):
 
         await asyncio.wait(self._background_tasks)
 
+    def _set_first_error(self, err: issues.Error):
+        try:
+            self._first_error.set_result(err)
+            self._state_changed.set()
+        except asyncio.InvalidStateError:
+            # skip if already has result
+            pass
 
 class ReaderStream:
     _token_getter: Optional[TokenGetterFuncType]
@@ -93,8 +125,8 @@ class ReaderStream:
 
     _state_changed: asyncio.Event
     _closed: bool
-    _first_error: Optional[YdbError]
     _message_batches: typing.Deque[PublicBatch]
+    first_error: asyncio.Future[YdbError]
 
     def __init__(self, settings: PublicReaderSettings):
         self._token_getter = settings._token_getter
@@ -107,7 +139,7 @@ def __init__(self, settings: PublicReaderSettings):
 
         self._state_changed = asyncio.Event()
         self._closed = False
-        self._first_error = None
+        self.first_error = asyncio.get_running_loop().create_future()
         self._message_batches = deque()
 
     @staticmethod
@@ -144,8 +176,8 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
 
     async def wait_messages(self):
         while True:
-            if self._first_error is not None:
-                raise self._first_error
+            if self._get_first_error() is not None:
+                raise self._get_first_error()
 
             if len(self._message_batches) > 0:
                 return
@@ -154,8 +186,8 @@ async def wait_messages(self):
             self._state_changed.clear()
 
     def receive_batch_nowait(self):
-        if self._first_error is not None:
-            raise self._first_error
+        if self._get_first_error() is not None:
+            raise self._get_first_error()
 
         try:
             batch = self._message_batches.popleft()
@@ -173,6 +205,7 @@ async def _read_messages_loop(self, stream: IGrpcWrapperAsyncIO):
             ))
             while True:
                 message = await stream.receive()  # type: StreamReadMessage.FromServer
+                _process_response(message.server_status)
                 if isinstance(message.server_message, StreamReadMessage.ReadResponse):
                     self._on_read_response(message.server_message)
                 elif isinstance(message.server_message, StreamReadMessage.StartPartitionSessionRequest):
@@ -285,9 +318,18 @@ def _read_response_to_batches(self, message: StreamReadMessage.ReadResponse) ->
         return batches
 
     def _set_first_error(self, err):
-        if self._first_error is None:
-            self._first_error = err
+        try:
+            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()
+        else:
+            return None
 
     async def close(self):
         if self._closed:

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -6,11 +6,11 @@
 import pytest
 
 import ydb
-from ydb import aio
+from ydb import aio, issues
 from .datatypes import PublicBatch, PublicMessage
 from .topic_reader import PublicReaderSettings
-from .topic_reader_asyncio import ReaderStream, PartitionSession
-from .._topic_wrapper.common import OffsetsRange, Codec, ServerStatus
+from .topic_reader_asyncio import ReaderStream, PartitionSession, ReaderReconnector
+from .._topic_wrapper.common import OffsetsRange, Codec, ServerStatus, UpdateTokenResponse
 from .._topic_wrapper.reader import StreamReadMessage
 from .._topic_wrapper.test_helpers import StreamMock, wait_condition, wait_for_fast
 from ..issues import Unavailable
@@ -135,14 +135,14 @@ async def stream_reader_started(self, stream, default_reader_settings, partition
     async def stream_reader(self, stream_reader_started: ReaderStream):
         yield stream_reader_started
 
-        assert stream_reader_started._first_error is None
+        assert stream_reader_started._get_first_error() is None
         await stream_reader_started.close()
 
     @pytest.fixture()
     async def stream_reader_finish_with_error(self, stream_reader_started: ReaderStream):
         yield stream_reader_started
 
-        assert stream_reader_started._first_error is not None
+        assert stream_reader_started._get_first_error() is not None
         await stream_reader_started.close()
 
 
@@ -195,16 +195,9 @@ def batch_count():
         )))
         await wait_condition(lambda: batch_count() > initial_batches)
 
-    async def test_first_error(self, stream, stream_reader_finish_with_error):
-        class TestError(grpc.RpcError, grpc.Call):
-            def __init__(self):
-                pass
-
-            def code(self):
-                return grpc.StatusCode.UNAUTHENTICATED
-
-            def details(self):
-                return "test error"
+    async def test_unknown_error(self, stream, stream_reader_finish_with_error):
+        class TestError(Exception):
+            pass
 
         test_err = TestError()
         stream.from_server.put_nowait(test_err)
@@ -215,6 +208,24 @@ def details(self):
         with pytest.raises(TestError):
             stream_reader_finish_with_error.receive_batch_nowait()
 
+    async def test_error_from_status_code(self, stream, stream_reader_finish_with_error):
+        # noinspection PyTypeChecker
+        stream.from_server.put_nowait(
+            StreamReadMessage.FromServer(
+                server_status=ServerStatus(
+                    status=issues.StatusCode.OVERLOADED,
+                    issues=[],
+                ),
+                server_message=None,
+            )
+        )
+
+        with pytest.raises(issues.Overloaded):
+            await wait_for_fast(stream_reader_finish_with_error.wait_messages())
+
+        with pytest.raises(issues.Overloaded):
+            stream_reader_finish_with_error.receive_batch_nowait()
+
     async def test_init_reader(self, stream, default_reader_settings):
         reader = ReaderStream(default_reader_settings)
         init_message = StreamReadMessage.InitRequest(
@@ -619,8 +630,15 @@ async def test_receive_batch_nowait(self, stream, stream_reader, partition_sessi
         with pytest.raises(asyncio.QueueEmpty):
             stream.from_client.get_nowait()
 
+
 @pytest.mark.asyncio
 class TestReaderReconnector:
-    async def test_start(self):
-        pass
+    async def test_reconnect_on_repeatable_error(self, monkeypatch):
+        def stream_create():
+            pass
+
+        with mock.patch.object(ReaderStream, "create", stream_create):
+            reconnector = ReaderReconnector(None, PublicReaderSettings("", ""))
+            await reconnector.wait_message()
 
+        raise NotImplementedError()

ydb/_topic_wrapper/common.py

@@ -211,11 +211,11 @@ class ServerStatus(IFromProto):
 
     def __init__(
             self,
-            status_code: ydb_status_codes_pb2.StatusIds.StatusCode,
-            grpc_issues: typing.Iterable[ydb_issue_message_pb2.IssueMessage],
+            status: issues.StatusCode,
+            issues: typing.Iterable[typing.Any],
     ):
-        self._grpc_status_code = status_code
-        self._issues = grpc_issues
+        self.status = status
+        self.issues = issues
 
     def __str__(self):
         return self.__repr__()
@@ -228,7 +228,7 @@ def from_proto(msg: typing.Union[
         return ServerStatus(msg.status, msg.issues)
 
     def is_success(self) -> bool:
-        return self._grpc_status_code == ydb_status_codes_pb2.StatusIds.SUCCESS
+        return self.status == issues.StatusCode.SUCCESS
 
     @classmethod
     def issue_to_str(cls, issue: ydb_issue_message_pb2.IssueMessage):

ydb/_topic_wrapper/common_test.py

@@ -3,7 +3,7 @@
 import grpc
 import pytest
 
-from .common import callback_from_asyncio, GrpcWrapperAsyncIO
+from .common import callback_from_asyncio, GrpcWrapperAsyncIO, ServerStatus
 from .. import issues
 
 # Workaround for good autocomplete in IDE and universal import at runtime
@@ -90,3 +90,22 @@ async def __anext__(self):
         with pytest.raises(issues.Overloaded):
             await wrapper.receive()
 
+
+class TestServerStatus:
+    def test_success(self):
+        status = ServerStatus(
+            status=ydb_status_codes_pb2.StatusIds.SUCCESS,
+            issues=[],
+        )
+        assert status.is_success()
+        assert issues._process_response(status) is None
+
+    def test_failed(self):
+        status = ServerStatus(
+            status=ydb_status_codes_pb2.StatusIds.OVERLOADED,
+            issues=[],
+        )
+        assert not status.is_success()
+        with pytest.raises(issues.Overloaded):
+            issues._process_response(status)
+

ydb/_topic_wrapper/reader.py

@@ -247,17 +247,20 @@ class FromServer(IFromProto):
         @staticmethod
         def from_proto(msg: ydb_topic_pb2.StreamReadMessage.FromServer) -> "StreamReadMessage.FromServer":
             mess_type = msg.WhichOneof("server_message")
-            server_status = ServerStatus.from_proto(ms)
+            server_status = ServerStatus.from_proto(msg)
             if mess_type == "read_response":
                 return StreamReadMessage.FromServer(
-                    server_message=StreamReadMessage.ReadResponse.from_proto(msg.read_response)
+                    server_status=server_status,
+                    server_message=StreamReadMessage.ReadResponse.from_proto(msg.read_response),
                 )
             elif mess_type == "init_response":
                 return StreamReadMessage.FromServer(
+                    server_status=server_status,
                     server_message=StreamReadMessage.InitResponse.from_proto(msg.init_response),
                 )
             elif mess_type == "start_partition_session_request":
                 return StreamReadMessage.FromServer(
+                    server_status=server_status,
                     server_message=StreamReadMessage.StartPartitionSessionRequest.from_proto(msg.start_partition_session_request)
                 )
 

ydb/issues.py

@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+import abc
+
 from google.protobuf import text_format
 import enum
 from six.moves import queue