Fix hang while write messages to closed driver.

Close #296

Author: rekby

tests/topics/test_topic_writer.py

@@ -96,12 +96,38 @@ async def test_write_multi_message_with_ack(
             None,
         ],
     )
-    async def test_write_encoded(self, driver: ydb.Driver, topic_path: str, codec):
+    async def test_write_encoded(self, driver: ydb.aio.Driver, topic_path: str, codec):
         async with driver.topic_client.writer(topic_path, codec=codec) as writer:
             await writer.write("a" * 1000)
             await writer.write("b" * 1000)
             await writer.write("c" * 1000)
 
+    async def test_create_writer_after_stop(self, driver: ydb.aio.Driver, topic_path: str):
+        await driver.stop()
+        with pytest.raises(ydb.Error):
+            async with driver.topic_client.writer(topic_path) as writer:
+                await writer.write_with_ack("123")
+
+    async def test_send_message_after_stop(self, driver: ydb.aio.Driver, topic_path: str):
+        writer = driver.topic_client.writer(topic_path)
+        await driver.stop()
+        with pytest.raises(ydb.Error):
+            await writer.write_with_ack("123")
+
+    async def test_preserve_exception_on_cm_close(self, driver: ydb.aio.Driver, topic_path: str):
+        class TestException(Exception):
+            pass
+
+        with pytest.raises(TestException):
+            async with driver.topic_client.writer(topic_path) as writer:
+                driver.stop()  # will raise exception on topic writer __exit__
+                try:
+                    writer.write("123")
+                except ydb.Error:
+                    pass
+
+                raise TestException()
+
 
 class TestTopicWriterSync:
     def test_send_message(self, driver_sync: ydb.Driver, topic_path):
@@ -212,3 +238,29 @@ def test_start_many_sync_writers_in_parallel(self, driver_sync: ydb.Driver, topi
 
         for writer in writers:
             writer.close()
+
+    def test_create_writer_after_stop(self, driver_sync: ydb.Driver, topic_path: str):
+        driver_sync.stop()
+        with pytest.raises(ydb.Error):
+            with driver_sync.topic_client.writer(topic_path) as writer:
+                writer.write_with_ack("123")
+
+    def test_send_message_after_stop(self, driver_sync: ydb.Driver, topic_path: str):
+        writer = driver_sync.topic_client.writer(topic_path)
+        driver_sync.stop()
+        with pytest.raises(ydb.Error):
+            writer.write_with_ack("123")
+
+    def test_preserve_exception_on_cm_close(self, driver_sync: ydb.Driver, topic_path: str):
+        class TestException(Exception):
+            pass
+
+        with pytest.raises(TestException):
+            with driver_sync.topic_client.writer(topic_path) as writer:
+                driver_sync.stop()  # will raise exception on topic writer __exit__
+                try:
+                    writer.write("123")
+                except ydb.Error:
+                    pass
+
+                raise TestException()

ydb/_topic_writer/topic_writer_asyncio.py

@@ -330,7 +330,7 @@ def _prepare_internal_messages(self, messages: List[PublicMessage]) -> List[Inte
 
     def _check_stop(self):
         if self._stop_reason.done():
-            raise self._stop_reason.result()
+            raise self._stop_reason.exception()
 
     async def _connection_loop(self):
         retry_settings = RetrySettings()  # todo
@@ -543,7 +543,7 @@ def _stop(self, reason: BaseException):
         if self._stop_reason.done():
             return
 
-        self._stop_reason.set_result(reason)
+        self._stop_reason.set_exception(reason)
 
         for f in self._messages_future:
             f.set_exception(reason)

ydb/_topic_writer/topic_writer_sync.py

@@ -5,6 +5,7 @@
 from concurrent.futures import Future
 from typing import Union, List, Optional
 
+import ydb
 from .._grpc.grpcwrapper.common_utils import SupportedDriverType
 from .topic_writer import (
     PublicWriterSettings,
@@ -56,7 +57,12 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        try:
+            self.close()
+        except ydb.Error:
+            if exc_val:
+                raise exc_val
+            raise
 
     def __del__(self):
         self.close(flush=False)

ydb/aio/pool.py

@@ -241,6 +241,8 @@ async def __call__(
         preferred_endpoint=None,
         fast_fail=False,
     ):
+        if self._stopped:
+            raise issues.Error("Driver was stopped")
         wait_timeout = settings.timeout if settings else 10
         try:
             connection = await self._store.get(preferred_endpoint, fast_fail=fast_fail, wait_timeout=wait_timeout)

ydb/pool.py

@@ -429,6 +429,9 @@ def __call__(
 
         :return: A result of computation
         """
+        if self._stopped:
+            raise issues.Error("Driver was stopped")
+
         tracing.trace(self.tracer, {"request": request, "stub": stub, "rpc_name": rpc_name})
         try:
             connection = self._store.get(preferred_endpoint)

ydb/topic.py

@@ -337,7 +337,7 @@ def writer(
         encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None,
         encoder_executor: Optional[concurrent.futures.Executor] = None,  # default shared client executor pool
     ) -> TopicWriter:
-        args = locals()
+        args = locals().copy()
         del args["self"]
         self._check_closed()