add stop marker

fix close path

Author: rekby

tests/topics/test_topic_reader.py

@@ -9,3 +9,4 @@ async def test_read_message(
         reader = driver.topic_client.topic_reader(topic_consumer, topic_path)
 
         assert await reader.receive_batch() is not None
+        await reader.close()

ydb/_grpc/grpcwrapper/common_utils.py

@@ -69,6 +69,9 @@ class UnknownGrpcMessageError(issues.Error):
     pass
 
 
+_stop_grpc_connection_marker = object()
+
+
 class QueueToIteratorAsyncIO:
     __slots__ = ("_queue",)
 
@@ -80,7 +83,7 @@ def __aiter__(self):
 
     async def __anext__(self):
         item = await self._queue.get()
-        if item is None:
+        if item is _stop_grpc_connection_marker:
             raise StopAsyncIteration()
         return item
 
@@ -101,7 +104,7 @@ def __iter__(self):
 
     def __next__(self):
         item = asyncio.run_coroutine_threadsafe(self._queue.get(), self._loop).result()
-        if item is None:
+        if item is _stop_grpc_connection_marker:
             raise StopIteration()
         return item
 
@@ -161,8 +164,9 @@ async def start(self, driver: SupportedDriverType, stub, method):
         self._connection_state = "started"
 
     def close(self):
-        self.from_client_grpc.put_nowait(None)
-        self._stream_call.cancel()
+        self.from_client_grpc.put_nowait(_stop_grpc_connection_marker)
+        if self._stream_call:
+            self._stream_call.cancel()
 
     async def _start_asyncio_driver(self, driver: ydb.aio.Driver, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)

ydb/_topic_writer/topic_writer_asyncio.py

@@ -296,6 +296,7 @@ async def _connection_loop(self):
             pending = []
 
             # noinspection PyBroadException
+            stream_writer = None
             try:
                 stream_writer = await WriterAsyncIOStream.create(
                     self._driver, self._init_message, self._get_token
@@ -339,6 +340,8 @@ async def _connection_loop(self):
                 self._stop(err)
                 return
             finally:
+                if stream_writer:
+                    stream_writer.close()
                 if len(pending) > 0:
                     for task in pending:
                         task.cancel()