Fix topic stream error handling

Author: vgvoleg

ydb/_grpc/grpcwrapper/common_utils.py

@@ -6,6 +6,7 @@
 import contextvars
 import datetime
 import functools
+import logging
 import typing
 from typing import (
     Optional,
@@ -37,6 +38,8 @@
 from ...settings import BaseRequestSettings
 from ..._constants import DEFAULT_LONG_STREAM_TIMEOUT
 
+logger = logging.getLogger(__name__)
+
 
 class IFromProto(abc.ABC):
     @staticmethod
@@ -193,28 +196,32 @@ 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 = await driver(
+        stream_call, on_disconnected_lambda = await driver(
             requests_iterator,
             stub,
             method,
             settings=self._stream_settings,
+            include_disconnected_lambda_to_result=True,
         )
         self._stream_call = stream_call
+        self._on_disconnected_lambda = on_disconnected_lambda
         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 = await to_thread(
+        stream_call, on_disconnected_lambda = await to_thread(
             driver,
             requests_iterator,
             stub,
             method,
             executor=self._wait_executor,
             settings=self._stream_settings,
+            include_disconnected_lambda_to_result=True,
         )
         self._stream_call = stream_call
+        self._on_disconnected_lambda = on_disconnected_lambda
         self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
 
     async def receive(self, timeout: Optional[int] = None) -> Any:
@@ -230,6 +237,11 @@ 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_lambda:
+                coro = self._on_disconnected_lambda()
+                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

@@ -248,12 +248,15 @@ async def _connection_loop(self):
                 self._state_changed.set()
                 await self._stream_reader.wait_error()
             except BaseException as err:
+                logger.debug("reader %s, attempt %s connection loop error %s", self._id, attempt, err)
                 retry_info = check_retriable_error(err, self._settings._retry_settings(), attempt)
                 if not retry_info.is_retriable:
                     logger.debug("reader %s stop connection loop due to %s", self._id, err)
                     self._set_first_error(err)
                     return
 
+                logger.debug("sleep before retry for %s seconds", retry_info.sleep_timeout_seconds)
+
                 await asyncio.sleep(retry_info.sleep_timeout_seconds)
 
                 attempt += 1
@@ -498,6 +501,9 @@ 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/aio/pool.py

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

ydb/pool.py

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