Fix asyncio.Cancelled error propagated to writer/reader (#800)

Author: vgvoleg

ydb/_topic_reader/topic_reader_asyncio.py

@@ -727,6 +727,11 @@ async def _read_messages_loop(self):
                     logger.exception("unexpected message in stream reader: %s" % e)
 
                 self._state_changed.set()
+        except asyncio.CancelledError as e:
+            logger.debug("reader stream %s error: %s", self._id, e)
+            if not self._closed:
+                self._set_first_error(issues.ConnectionLost("gRPC stream cancelled"))
+            raise
         except Exception as e:
             logger.debug("reader stream %s error: %s", self._id, e)
             self._set_first_error(e)

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -1566,3 +1566,40 @@ async def stream_create(
 
         reader_stream_mock_with_error.wait_error.assert_any_await()
         reader_stream_mock_with_error.wait_messages.assert_any_await()
+
+    async def test_wait_error_returns_on_cancelled_error_from_receive(self, default_reader_settings):
+        receive_call = 0
+
+        async def receive(timeout=None):
+            nonlocal receive_call
+            receive_call += 1
+            if receive_call == 1:
+                # first call: return init response (from _start)
+                return StreamReadMessage.FromServer(
+                    server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                    server_message=StreamReadMessage.InitResponse(session_id="test-session"),
+                )
+            # subsequent calls (from _read_messages_loop): simulate gRPC-level CancelledError
+            raise asyncio.CancelledError()
+
+        stream = mock.MagicMock(spec=StreamMock)
+        stream.receive = receive
+        stream.write = mock.Mock()
+        stream.close = mock.Mock()
+
+        reader = ReaderStream(4, default_reader_settings)
+        await reader._start(stream, default_reader_settings._init_message())
+
+        # Bug: wait_error() hangs forever because _first_error is never set.
+        # After the fix: wait_error() returns with a retriable ConnectionLost error.
+        try:
+            await asyncio.wait_for(reader.wait_error(), timeout=1.0)
+        except asyncio.TimeoutError:
+            pytest.fail(
+                "wait_error() hung forever after CancelledError from receive() — "
+                "bug: CancelledError bypasses 'except Exception' in _read_messages_loop"
+            )
+        except Exception:
+            pass  # any error is fine, we just need wait_error() to not hang
+
+        await reader.close(False)

ydb/_topic_writer/topic_writer_asyncio.py

@@ -478,7 +478,12 @@ async def _connection_loop(self):
                 tasks = [send_loop, receive_loop]
                 done, _ = await asyncio.wait([send_loop, receive_loop], return_when=asyncio.FIRST_COMPLETED)
                 done.pop().result()  # need for raise exception - reason of stop task
-            except issues.Error as err:
+            except (asyncio.CancelledError, issues.Error) as err:
+                if isinstance(err, asyncio.CancelledError):
+                    if self._closed:
+                        return
+                    err = issues.ConnectionLost("gRPC stream cancelled")
+
                 err_info = check_retriable_error(err, retry_settings, attempt)
                 if not err_info.is_retriable or self._tx is not None:  # no retries in tx writer
                     logger.debug("writer reconnector %s stop connection loop due to %s", self._id, err)
@@ -492,7 +497,7 @@ async def _connection_loop(self):
                 )
                 await asyncio.sleep(err_info.sleep_timeout_seconds)
 
-            except (asyncio.CancelledError, Exception) as err:
+            except Exception as err:
                 self._stop(err)
                 return
             finally:

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -455,6 +455,44 @@ async def test_reconnect_and_resent_non_acked_messages_on_retriable_error(
         second_writer.from_server.put_nowait(self.make_default_ack_message(seq_no=2))
         await reconnector.close(flush=True)
 
+    async def test_reconnect_on_cancelled_error_from_receive(self, default_driver, default_settings, monkeypatch):
+        stream_creates = 0
+        stream_2_created = asyncio.Event()
+
+        class StreamWriterCancelOnFirstReceive(TestWriterAsyncIOReconnector.StreamWriterMock):
+            def __init__(self):
+                super().__init__()
+                self._first_receive = True
+
+            async def receive(self):
+                if self._first_receive:
+                    self._first_receive = False
+                    raise asyncio.CancelledError()
+                await asyncio.Future()  # stream 2 stays alive
+
+        async def create_mock(*args, **kwargs):
+            nonlocal stream_creates
+            stream_creates += 1
+            writer = StreamWriterCancelOnFirstReceive()
+            writer.last_seqno = TestWriterAsyncIOReconnector.init_last_seqno
+            if stream_creates >= 2:
+                stream_2_created.set()
+            return writer
+
+        with mock.patch.object(WriterAsyncIOStream, "create", create_mock):
+            reconnector = WriterAsyncIOReconnector(default_driver, default_settings)
+            try:
+                # Bug: stream 2 is never created — _stop(CancelledError) kills the writer permanently.
+                # After the fix: writer reconnects and stream 2 is created.
+                await asyncio.wait_for(stream_2_created.wait(), timeout=2.0)
+            except asyncio.TimeoutError:
+                pytest.fail(
+                    "Writer did not reconnect after CancelledError from receive() — "
+                    "bug: _stop(CancelledError) permanently kills writer"
+                )
+            finally:
+                await reconnector.close(False)
+
     async def test_stop_on_unexpected_exception(self, reconnector: WriterAsyncIOReconnector, get_stream_writer):
         class TestException(Exception):
             pass