another approach to safe error handling

Author: vgvoleg

tests/slo/playground/configs/chaos.sh

@@ -31,7 +31,7 @@ do
   sh -c "docker stop ${nodeForChaos} -t 10"
   sh -c "docker start ${nodeForChaos}"
 
-  sleep 60
+  sleep 30
 done
 
 # for i in $(seq 1 3)

tests/slo/slo_runner.sh

@@ -3,7 +3,7 @@ docker compose -f playground/configs/compose.yaml up -d --wait
 
 ../../.venv/bin/python ./src topic-create grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic
 
-# ../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --write-rps 1 --time 600 --async
+../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --write-rps 1 --time 120
 
-../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --time 10
-../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --write-threads 0 --read-rps 1 --time 600 
+# ../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --read-threads 0 --time 5
+# ../../.venv/bin/python ./src topic-run grpc://localhost:2135 /Root/testdb --path /Root/testdb/slo_topic --prom-pgw "" --write-threads 0 --read-rps 1 --time 200

ydb/_grpc/grpcwrapper/common_utils.py

@@ -165,7 +165,6 @@ def __init__(self, convert_server_grpc_to_wrapper):
         self._connection_state = "new"
         self._stream_call = None
         self._wait_executor = None
-        self._on_disconnected_callback = None
 
         self._stream_settings: BaseRequestSettings = (
             BaseRequestSettings()
@@ -197,32 +196,28 @@ def _clean_executor(self, wait: bool):
 
     async def _start_asyncio_driver(self, driver: DriverIO, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)
-        stream_call, on_disconnected_callback = await driver(
+        stream_call = await driver(
             requests_iterator,
             stub,
             method,
             settings=self._stream_settings,
-            include_disconnected_callback_to_result=True,
         )
         self._stream_call = stream_call
-        self._on_disconnected_callback = on_disconnected_callback
         self.from_server_grpc = stream_call.__aiter__()
 
     async def _start_sync_driver(self, driver: Driver, stub, method):
         requests_iterator = AsyncQueueToSyncIteratorAsyncIO(self.from_client_grpc)
         self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
-        stream_call, on_disconnected_callback = await to_thread(
+        stream_call = await to_thread(
             driver,
             requests_iterator,
             stub,
             method,
             executor=self._wait_executor,
             settings=self._stream_settings,
-            include_disconnected_callback_to_result=True,
         )
         self._stream_call = stream_call
-        self._on_disconnected_callback = on_disconnected_callback
         self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
 
     async def receive(self, timeout: Optional[int] = None) -> Any:
@@ -238,11 +233,6 @@ async def get_response():
                 grpc_message = await asyncio.wait_for(get_response(), timeout)
 
         except (grpc.RpcError, grpc.aio.AioRpcError) as e:
-            if self._on_disconnected_callback:
-                coro = self._on_disconnected_callback()
-                if asyncio.iscoroutine(coro):
-                    await coro
-
             raise connection._rpc_error_handler(self._connection_state, e)
 
         issues._process_response(grpc_message)

ydb/_topic_reader/topic_reader_asyncio.py

@@ -501,9 +501,6 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
             )  # type: StreamReadMessage.FromServer
         except asyncio.TimeoutError:
             raise TopicReaderError("Timeout waiting for init response")
-        except Exception as e:
-            logger.debug("reader stream %s init request error %s", self._id, e)
-            raise e
 
         if isinstance(init_response.server_message, StreamReadMessage.InitResponse):
             self._session_id = init_response.server_message.session_id

ydb/_topic_writer/topic_writer_asyncio.py

@@ -488,12 +488,12 @@ async def _connection_loop(self):
                     self._stop(err)
                     return
 
-                await asyncio.sleep(err_info.sleep_timeout_seconds)
                 logger.debug(
                     "writer reconnector %s retry in %s seconds",
                     self._id,
                     err_info.sleep_timeout_seconds,
                 )
+                await asyncio.sleep(err_info.sleep_timeout_seconds)
 
             except (asyncio.CancelledError, Exception) as err:
                 self._stop(err)

ydb/aio/connection.py

@@ -103,6 +103,28 @@ def future(self, *args, **kwargs):
         raise NotImplementedError
 
 
+class _SafeAsyncIterator:
+    def __init__(self, resp, rpc_state, on_disconnected_callback):
+        self.resp = resp
+        self.it = resp.__aiter__()
+        self.rpc_state = rpc_state
+        self.on_disconnected_callback = on_disconnected_callback
+
+    def cancel(self):
+        self.resp.cancel()
+        return self
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.it.__anext__()
+        except grpc.RpcError as rpc_error:
+            await self.on_disconnected_callback()
+            raise _rpc_error_handler(self.rpc_state, rpc_error)
+
+
 class Connection:
     __slots__ = (
         "endpoint",
@@ -192,6 +214,10 @@ async def __call__(
 
             response = await feature
             _log_response(rpc_state, response)
+
+            if hasattr(response, "__aiter__"):
+                response = _SafeAsyncIterator(response, rpc_state, on_disconnected)
+
             return response if wrap_result is None else wrap_result(rpc_state, response, *wrap_args)
         except grpc.RpcError as rpc_error:
             if on_disconnected:

ydb/aio/pool.py

@@ -259,7 +259,6 @@ async def __call__(
         wrap_args=(),
         preferred_endpoint=None,
         fast_fail=False,
-        include_disconnected_callback_to_result=False,
     ):
         if self._stopped:
             raise issues.Error("Driver was stopped")
@@ -271,18 +270,14 @@ async def __call__(
                 self._discovery.notify_disconnected()
             raise
 
-        on_disconnected = self._on_disconnected(connection)
-
         res = await connection(
             request,
             stub,
             rpc_name,
             wrap_result,
             settings,
             wrap_args,
-            on_disconnected,
+            self._on_disconnected(connection),
         )
 
-        if include_disconnected_callback_to_result:
-            return res, on_disconnected
         return res

ydb/connection.py

@@ -74,7 +74,7 @@ def _rpc_error_handler(
     :param rpc_error: an underlying rpc error to handle
     :param on_disconnected: a handler to call on disconnected connection
     """
-    logger.info("%s: received error, %s", rpc_state, rpc_error)
+    logger.debug("%s: received error, %s", rpc_state, rpc_error)
     if isinstance(rpc_error, (grpc.RpcError, grpc.aio.AioRpcError, grpc.Call, grpc.aio.Call)):
         if rpc_error.code() == grpc.StatusCode.UNAUTHENTICATED:
             return issues.Unauthenticated(rpc_error.details())
@@ -325,6 +325,28 @@ def __init__(self, endpoint, node_id):
         self.node_id = node_id
 
 
+class _SafeSyncIterator:
+    def __init__(self, resp, rpc_state, on_disconnected_callback):
+        self.resp = resp
+        self.it = resp.__iter__()
+        self.rpc_state = rpc_state
+        self.on_disconnected_callback = on_disconnected_callback
+
+    def cancel(self):
+        self.resp.cancel()
+        return self
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return self.it.__next__()
+        except grpc.RpcError as rpc_error:
+            self.on_disconnected_callback()
+            raise _rpc_error_handler(self.rpc_state, rpc_error)
+
+
 class Connection(object):
     __slots__ = (
         "endpoint",
@@ -466,6 +488,10 @@ def __call__(
                 compression=getattr(settings, "compression", None),
             )
             _log_response(rpc_state, response)
+
+            if hasattr(response, "__iter__"):
+                response = _SafeSyncIterator(response, rpc_state, on_disconnected)
+
             return response if wrap_result is None else wrap_result(rpc_state, response, *wrap_args)
         except grpc.RpcError as rpc_error:
             raise _rpc_error_handler(rpc_state, rpc_error, on_disconnected)

ydb/pool.py

@@ -437,7 +437,6 @@ def __call__(
         settings=None,
         wrap_args=(),
         preferred_endpoint=None,
-        include_disconnected_callback_to_result=False,
     ):
         """
         Synchronously sends request constructed by client library
@@ -462,22 +461,18 @@ def __call__(
             self._discovery_thread.notify_disconnected()
             raise
 
-        on_disconnected = lambda: self._on_disconnected(connection)  # noqa: E731
-
         res = connection(
             request,
             stub,
             rpc_name,
             wrap_result,
             settings,
             wrap_args,
-            on_disconnected,
+            lambda: self._on_disconnected(connection),
         )
 
         tracing.trace(self.tracer, {"response": res}, trace_level=tracing.TraceLevel.DEBUG)
 
-        if include_disconnected_callback_to_result:
-            return res, on_disconnected
         return res
 
     @_utilities.wrap_async_call_exceptions