close grpc streams when stream reader/writer closed

Author: rekby

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Close grpc streams while closing readers/writers
 * Add control plane operations for topic api: create, drop
 
 ## 3.0.1b4 ##

ydb/_grpc/grpcwrapper/common_utils.py

@@ -79,10 +79,10 @@ def __aiter__(self):
         return self
 
     async def __anext__(self):
-        try:
-            return await self._queue.get()
-        except asyncio.QueueEmpty:
+        item = await self._queue.get()
+        if item is None:
             raise StopAsyncIteration()
+        return item
 
 
 class AsyncQueueToSyncIteratorAsyncIO:
@@ -100,13 +100,10 @@ def __iter__(self):
         return self
 
     def __next__(self):
-        try:
-            res = asyncio.run_coroutine_threadsafe(
-                self._queue.get(), self._loop
-            ).result()
-            return res
-        except asyncio.QueueEmpty:
+        item = asyncio.run_coroutine_threadsafe(self._queue.get(), self._loop).result()
+        if item is None:
             raise StopIteration()
+        return item
 
 
 class SyncIteratorToAsyncIterator:
@@ -133,6 +130,10 @@ async def receive(self) -> Any:
     def write(self, wrap_message: IToProto):
         ...
 
+    @abc.abstractmethod
+    def close(self):
+        ...
+
 
 SupportedDriverType = Union[ydb.Driver, ydb.aio.Driver]
 
@@ -142,11 +143,15 @@ class GrpcWrapperAsyncIO(IGrpcWrapperAsyncIO):
     from_server_grpc: AsyncIterator
     convert_server_grpc_to_wrapper: Callable[[Any], Any]
     _connection_state: str
+    _stream_call: Optional[
+        Union[grpc.aio.StreamStreamCall, "grpc._channel._MultiThreadedRendezvous"]
+    ]
 
     def __init__(self, convert_server_grpc_to_wrapper):
         self.from_client_grpc = asyncio.Queue()
         self.convert_server_grpc_to_wrapper = convert_server_grpc_to_wrapper
         self._connection_state = "new"
+        self._stream_call = None
 
     async def start(self, driver: SupportedDriverType, stub, method):
         if asyncio.iscoroutinefunction(driver.__call__):
@@ -155,13 +160,18 @@ async def start(self, driver: SupportedDriverType, stub, method):
             await self._start_sync_driver(driver, stub, method)
         self._connection_state = "started"
 
+    def close(self):
+        self.from_client_grpc.put_nowait(None)
+        self._stream_call.cancel()
+
     async def _start_asyncio_driver(self, driver: ydb.aio.Driver, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)
         stream_call = await driver(
             requests_iterator,
             stub,
             method,
         )
+        self._stream_call = stream_call
         self.from_server_grpc = stream_call.__aiter__()
 
     async def _start_sync_driver(self, driver: ydb.Driver, stub, method):
@@ -172,6 +182,7 @@ async def _start_sync_driver(self, driver: ydb.Driver, stub, method):
             stub,
             method,
         )
+        self._stream_call = stream_call
         self.from_server_grpc = SyncIteratorToAsyncIterator(stream_call.__iter__())
 
     async def receive(self) -> Any:

ydb/_topic_common/test_helpers.py

@@ -8,20 +8,36 @@
 class StreamMock(IGrpcWrapperAsyncIO):
     from_server: asyncio.Queue
     from_client: asyncio.Queue
+    _closed: bool
 
     def __init__(self):
         self.from_server = asyncio.Queue()
         self.from_client = asyncio.Queue()
+        self._closed = False
 
     async def receive(self) -> typing.Any:
+        if self._closed:
+            raise Exception("read from closed StreamMock")
+
         item = await self.from_server.get()
+        if item is None:
+            raise StopAsyncIteration()
         if isinstance(item, Exception):
             raise item
         return item
 
     def write(self, wrap_message: IToProto):
+        if self._closed:
+            raise Exception("write to closed StreamMock")
         self.from_client.put_nowait(wrap_message)
 
+    def close(self):
+        if self._closed:
+            return
+
+        self._closed = True
+        self.from_server.put_nowait(None)
+
 
 async def wait_condition(f: typing.Callable[[], bool], timeout=1):
     start = time.monotonic()

ydb/_topic_reader/topic_reader_asyncio.py

@@ -496,6 +496,7 @@ async def close(self):
         self._closed = True
         self._set_first_error(TopicReaderStreamClosedError())
         self._state_changed.set()
+        self._stream.close()
 
         for task in self._background_tasks:
             task.cancel()

ydb/_topic_writer/topic_writer_asyncio.py

@@ -322,6 +322,7 @@ async def _connection_loop(self):
                 done, pending = await asyncio.wait(
                     [send_loop, receive_loop], return_when=asyncio.FIRST_COMPLETED
                 )
+                stream_writer.close()
                 done.pop().result()
             except issues.Error as err:
                 # todo log error
@@ -417,6 +418,9 @@ def __init__(
     ):
         self._token_getter = token_getter
 
+    def close(self):
+        self._stream.close()
+
     @staticmethod
     async def create(
         driver: SupportedDriverType,

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -158,20 +158,37 @@ class StreamWriterMock:
         from_client: asyncio.Queue
         from_server: asyncio.Queue
 
+        _closed: bool
+
         def __init__(self):
             self.last_seqno = 0
             self.from_server = asyncio.Queue()
             self.from_client = asyncio.Queue()
+            self._closed = False
 
         def write(self, messages: typing.List[InternalMessage]):
+            if self._closed:
+                raise Exception("write to closed StreamWriterMock")
+
             self.from_client.put_nowait(messages)
 
         async def receive(self) -> StreamWriteMessage.WriteResponse:
+            if self._closed:
+                raise Exception("read from closed StreamWriterMock")
+
             item = await self.from_server.get()
             if isinstance(item, Exception):
                 raise item
             return item
 
+        def close(self):
+            if self._closed:
+                return
+
+            self.from_server.put_nowait(
+                Exception("waited message while StreamWriterMock closed")
+            )
+
     @pytest.fixture(autouse=True)
     async def stream_writer_double_queue(self, monkeypatch):
         class DoubleQueueWriters: