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/_errors.py

@@ -1,13 +1,12 @@
 from dataclasses import dataclass
-from typing import Optional, Union
-
-import grpc
+from typing import Optional
 
 from . import issues
 
 _errors_retriable_fast_backoff_types = [
     issues.Unavailable,
     issues.ClientInternalError,
+    issues.ConnectionLost,
 ]
 _errors_retriable_slow_backoff_types = [
     issues.Aborted,
@@ -54,26 +53,3 @@ def check_retriable_error(err, retry_settings, attempt):
 class ErrorRetryInfo:
     is_retriable: bool
     sleep_timeout_seconds: Optional[float]
-
-
-def stream_error_converter(exc: BaseException) -> Union[issues.Error, BaseException]:
-    """Converts gRPC stream errors to appropriate YDB exception types.
-
-    This function takes a base exception and converts specific gRPC aio stream errors
-    to their corresponding YDB exception types for better error handling and semantic
-    clarity.
-
-    Args:
-        exc (BaseException): The original exception to potentially convert.
-
-    Returns:
-        BaseException: Either a converted YDB exception or the original exception
-                      if no specific conversion rule applies.
-    """
-    if isinstance(exc, (grpc.RpcError, grpc.aio.AioRpcError)):
-        if exc.code() == grpc.StatusCode.UNAVAILABLE:
-            return issues.Unavailable(exc.details() or "")
-        if exc.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-            return issues.DeadlineExceed("Deadline exceeded on request")
-        return issues.Error("Stream has been terminated. Original exception: {}".format(str(exc.details())))
-    return exc

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/_utilities.py

@@ -2,10 +2,9 @@
 
 
 class AsyncResponseIterator(object):
-    def __init__(self, it, wrapper, error_converter=None):
+    def __init__(self, it, wrapper):
         self.it = it.__aiter__()
         self.wrapper = wrapper
-        self.error_converter = error_converter
 
     def cancel(self):
         self.it.cancel()
@@ -18,12 +17,7 @@ def __aiter__(self):
         return self
 
     async def _next(self):
-        try:
-            res = self.wrapper(await self.it.__anext__())
-        except BaseException as e:
-            if self.error_converter:
-                raise self.error_converter(e) from e
-            raise e
+        res = self.wrapper(await self.it.__anext__())
 
         if res is not None:
             return res

ydb/aio/connection.py

@@ -11,6 +11,7 @@
     _log_request,
     _log_response,
     _rpc_error_handler,
+    _is_disconnect_needed,
     _get_request_timeout,
     _set_server_timeouts,
     _RpcState as RpcState,
@@ -103,6 +104,30 @@ 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:
+            ydb_error = _rpc_error_handler(self.rpc_state, rpc_error, use_unavailable=True)
+            if _is_disconnect_needed(ydb_error):
+                await self.on_disconnected_callback()
+            raise ydb_error
+
+
 class Connection:
     __slots__ = (
         "endpoint",
@@ -192,6 +217,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/aio/query/session.py

@@ -23,7 +23,6 @@
 )
 
 from ..._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT
-from ..._errors import stream_error_converter
 
 
 class QuerySession(BaseQuerySession):
@@ -164,7 +163,6 @@ async def execute(
                 session=self,
                 settings=self._settings,
             ),
-            error_converter=stream_error_converter,
         )
 
     async def explain(