allow save link to parent client - for prevent stop it (and underlay executors) by GC.

Author: rekby

tests/topics/test_topic_reader.py

@@ -16,6 +16,10 @@ async def test_read_batch(
 
         await reader.close()
 
+    async def test_link_to_client(self, driver, topic_path, topic_consumer):
+        reader = driver.topic_client.reader(topic_path, topic_consumer)
+        assert reader._parent is driver.topic_client
+
     async def test_read_message(
         self, driver, topic_path, topic_with_messages, topic_consumer
     ):
@@ -84,6 +88,10 @@ def test_read_batch(
 
         reader.close()
 
+    def test_link_to_client(self, driver_sync, topic_path, topic_consumer):
+        reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
+        assert reader._parent is driver_sync.topic_client
+
     def test_read_message(
         self, driver_sync, topic_path, topic_with_messages, topic_consumer
     ):

tests/topics/test_topic_writer.py

@@ -27,6 +27,10 @@ async def test_wait_last_seqno(self, driver: ydb.aio.Driver, topic_path):
             init_info = await writer2.wait_init()
             assert init_info.last_seqno == 5
 
+    async def test_link_to_client(self, driver, topic_path, topic_consumer):
+        writer = driver.topic_client.writer(topic_path)
+        assert writer._parent is driver.topic_client
+
     async def test_random_producer_id(
         self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
     ):
@@ -138,6 +142,10 @@ def test_auto_flush_on_close(self, driver_sync: ydb.Driver, topic_path):
             init_info = writer.wait_init()
             assert init_info.last_seqno == last_seqno
 
+    def test_link_to_client(self, driver_sync, topic_path, topic_consumer):
+        writer = driver_sync.topic_client.writer(topic_path)
+        assert writer._parent is driver_sync.topic_client
+
     def test_random_producer_id(
         self,
         driver_sync: ydb.aio.Driver,

ydb/_topic_reader/topic_reader_asyncio.py

@@ -26,6 +26,7 @@
     Codec,
 )
 from .._errors import check_retriable_error
+from .. import topic
 
 
 class TopicReaderError(YdbError):
@@ -61,11 +62,19 @@ class PublicAsyncIOReader:
     _loop: asyncio.AbstractEventLoop
     _closed: bool
     _reconnector: ReaderReconnector
+    _parent: typing.Any  # need for prevent close parent client by GC
 
-    def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
+    def __init__(
+            self,
+            driver: Driver,
+            settings: topic_reader.PublicReaderSettings,
+            *,
+            _parent=None,
+    ):
         self._loop = asyncio.get_running_loop()
         self._closed = False
         self._reconnector = ReaderReconnector(driver, settings)
+        self._parent = _parent
 
     async def __aenter__(self):
         return self
@@ -78,7 +87,7 @@ def __del__(self):
             self._loop.create_task(self.close(flush=False), name="close reader")
 
     async def receive_batch(
-        self,
+            self,
     ) -> typing.Union[datatypes.PublicBatch, None]:
         """
         Get one messages batch from reader.
@@ -99,7 +108,7 @@ async def receive_message(self) -> typing.Optional[datatypes.PublicMessage]:
         return self._reconnector.receive_message_nowait()
 
     def commit(
-        self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
+            self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
     ):
         """
         Write commit message to a buffer.
@@ -110,7 +119,7 @@ def commit(
         self._reconnector.commit(batch)
 
     async def commit_with_ack(
-        self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
+            self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
     ):
         """
         write commit message to a buffer and wait ack from the server.
@@ -195,7 +204,7 @@ def receive_message_nowait(self):
         return self._stream_reader.receive_message_nowait()
 
     def commit(
-        self, batch: datatypes.ICommittable
+            self, batch: datatypes.ICommittable
     ) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)
 
@@ -254,10 +263,10 @@ class ReaderStream:
     _get_token_function: Callable[[], str]
 
     def __init__(
-        self,
-        reader_reconnector_id: int,
-        settings: topic_reader.PublicReaderSettings,
-        get_token_function: Optional[Callable[[], str]] = None,
+            self,
+            reader_reconnector_id: int,
+            settings: topic_reader.PublicReaderSettings,
+            get_token_function: Optional[Callable[[], str]] = None,
     ):
         self._loop = asyncio.get_running_loop()
         self._id = ReaderStream._static_id_counter.inc_and_get()
@@ -286,9 +295,9 @@ def __init__(
 
     @staticmethod
     async def create(
-        reader_reconnector_id: int,
-        driver: SupportedDriverType,
-        settings: topic_reader.PublicReaderSettings,
+            reader_reconnector_id: int,
+            driver: SupportedDriverType,
+            settings: topic_reader.PublicReaderSettings,
     ) -> "ReaderStream":
         stream = GrpcWrapperAsyncIO(StreamReadMessage.FromServer.from_proto)
 
@@ -306,7 +315,7 @@ async def create(
         return reader
 
     async def _start(
-        self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest
+            self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest
     ):
         if self._started:
             raise TopicReaderError("Double start ReaderStream")
@@ -372,7 +381,7 @@ def receive_message_nowait(self):
         return message
 
     def commit(
-        self, batch: datatypes.ICommittable
+            self, batch: datatypes.ICommittable
     ) -> datatypes.PartitionSession.CommitAckWaiter:
         partition_session = batch._commit_get_partition_session()
 
@@ -426,19 +435,19 @@ async def _read_messages_loop(self):
                     self._on_read_response(message.server_message)
 
                 elif isinstance(
-                    message.server_message, StreamReadMessage.CommitOffsetResponse
+                        message.server_message, StreamReadMessage.CommitOffsetResponse
                 ):
                     self._on_commit_response(message.server_message)
 
                 elif isinstance(
-                    message.server_message,
-                    StreamReadMessage.StartPartitionSessionRequest,
+                        message.server_message,
+                        StreamReadMessage.StartPartitionSessionRequest,
                 ):
                     self._on_start_partition_session(message.server_message)
 
                 elif isinstance(
-                    message.server_message,
-                    StreamReadMessage.StopPartitionSessionRequest,
+                        message.server_message,
+                        StreamReadMessage.StopPartitionSessionRequest,
                 ):
                     self._on_partition_session_stop(message.server_message)
 
@@ -470,12 +479,12 @@ async def _update_token(self, token: str):
             self._update_token_event.clear()
 
     def _on_start_partition_session(
-        self, message: StreamReadMessage.StartPartitionSessionRequest
+            self, message: StreamReadMessage.StartPartitionSessionRequest
     ):
         try:
             if (
-                message.partition_session.partition_session_id
-                in self._partition_sessions
+                    message.partition_session.partition_session_id
+                    in self._partition_sessions
             ):
                 raise TopicReaderError(
                     "Double start partition session: %s"
@@ -506,7 +515,7 @@ def _on_start_partition_session(
             self._set_first_error(err)
 
     def _on_partition_session_stop(
-        self, message: StreamReadMessage.StopPartitionSessionRequest
+            self, message: StreamReadMessage.StopPartitionSessionRequest
     ):
         if message.partition_session_id not in self._partition_sessions:
             # may if receive stop partition with graceful=false after response on stop partition
@@ -554,7 +563,7 @@ def _buffer_release_bytes(self, bytes_size):
         )
 
     def _read_response_to_batches(
-        self, message: StreamReadMessage.ReadResponse
+            self, message: StreamReadMessage.ReadResponse
     ) -> typing.List[datatypes.PublicBatch]:
         batches = []
 
@@ -564,7 +573,7 @@ def _read_response_to_batches(
 
         bytes_per_batch = message.bytes_size // batch_count
         additional_bytes_to_last_batch = (
-            message.bytes_size - bytes_per_batch * batch_count
+                message.bytes_size - bytes_per_batch * batch_count
         )
 
         for partition_data in message.partition_data:

ydb/_topic_reader/topic_reader_sync.py

@@ -25,13 +25,15 @@ class TopicReaderSync:
     _caller: CallFromSyncToAsync
     _async_reader: PublicAsyncIOReader
     _closed: bool
+    _parent: typing.Any  # need for prevent stop the client by GC
 
     def __init__(
         self,
         driver: SupportedDriverType,
         settings: PublicReaderSettings,
         *,
         eventloop: Optional[asyncio.AbstractEventLoop] = None,
+        _parent=None,  # need for prevent stop the client by GC
     ):
         self._closed = False
 
@@ -49,6 +51,8 @@ async def create_reader():
             create_reader(), loop
         ).result()
 
+        self._parent = _parent
+
     def __del__(self):
         self.close(flush=False)
 

ydb/_topic_writer/topic_writer_asyncio.py

@@ -40,20 +40,29 @@
     GrpcWrapperAsyncIO,
 )
 
+from .. import topic
+
 logger = logging.getLogger(__name__)
 
 
 class WriterAsyncIO:
     _loop: asyncio.AbstractEventLoop
     _reconnector: "WriterAsyncIOReconnector"
     _closed: bool
+    _parent: typing.Any  # need for prevent close parent client by GC
 
-    def __init__(self, driver: SupportedDriverType, settings: PublicWriterSettings):
+    def __init__(
+            self,
+            driver: SupportedDriverType,
+            settings: PublicWriterSettings,
+            _client=None,
+    ):
         self._loop = asyncio.get_running_loop()
         self._closed = False
         self._reconnector = WriterAsyncIOReconnector(
             driver=driver, settings=WriterSettings(settings)
         )
+        self._parent = _client
 
     async def __aenter__(self) -> "WriterAsyncIO":
         return self

ydb/_topic_writer/topic_writer_sync.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import asyncio
+import typing
 from concurrent.futures import Future
 from typing import Union, List, Optional
 
+from .. import topic
 from .._grpc.grpcwrapper.common_utils import SupportedDriverType
 from .topic_writer import (
     PublicWriterSettings,
@@ -25,13 +27,15 @@ class WriterSync:
     _caller: CallFromSyncToAsync
     _async_writer: WriterAsyncIO
     _closed: bool
+    _parent: typing.Any  # need for prevent close parent client by GC
 
     def __init__(
-        self,
-        driver: SupportedDriverType,
-        settings: PublicWriterSettings,
-        *,
-        eventloop: Optional[asyncio.AbstractEventLoop] = None,
+            self,
+            driver: SupportedDriverType,
+            settings: PublicWriterSettings,
+            *,
+            eventloop: Optional[asyncio.AbstractEventLoop] = None,
+            _parent=None
     ):
 
         self._closed = False
@@ -49,6 +53,7 @@ async def create_async_writer():
         self._async_writer = self._caller.safe_call_with_result(
             create_async_writer(), None
         )
+        self._parent = _parent
 
     def __enter__(self):
         return self
@@ -96,17 +101,17 @@ def wait_init(self, *, timeout: TimeoutType = None) -> PublicWriterInitInfo:
         )
 
     def write(
-        self,
-        messages: Union[Message, List[Message]],
-        timeout: TimeoutType = None,
+            self,
+            messages: Union[Message, List[Message]],
+            timeout: TimeoutType = None,
     ):
         self._check_closed()
 
         self._caller.safe_call_with_result(self._async_writer.write(messages), timeout)
 
     def async_write_with_ack(
-        self,
-        messages: Union[Message, List[Message]],
+            self,
+            messages: Union[Message, List[Message]],
     ) -> Future[Union[PublicWriteResult, List[PublicWriteResult]]]:
         self._check_closed()
 
@@ -115,9 +120,9 @@ def async_write_with_ack(
         )
 
     def write_with_ack(
-        self,
-        messages: Union[Message, List[Message]],
-        timeout: Union[float, None] = None,
+            self,
+            messages: Union[Message, List[Message]],
+            timeout: Union[float, None] = None,
     ) -> Union[PublicWriteResult, List[PublicWriteResult]]:
         self._check_closed()