Remove 10min timeout from topic stream

Author: vgvoleg

ydb/_grpc/grpcwrapper/common_utils.py

@@ -34,6 +34,10 @@
     from ..common.protos import ydb_topic_pb2, ydb_issue_message_pb2
 
 from ... import issues, connection
+from ...settings import BaseRequestSettings
+
+
+DEFAULT_LONG_TIMEOUT = 31536000  # year
 
 
 class IFromProto(abc.ABC):
@@ -131,7 +135,7 @@ async def __anext__(self):
 
 class IGrpcWrapperAsyncIO(abc.ABC):
     @abc.abstractmethod
-    async def receive(self) -> Any:
+    async def receive(self, timeout: Optional[int] = None) -> Any:
         ...
 
     @abc.abstractmethod
@@ -161,6 +165,13 @@ def __init__(self, convert_server_grpc_to_wrapper):
         self._stream_call = None
         self._wait_executor = None
 
+        self._stream_settings: BaseRequestSettings = (
+            BaseRequestSettings()
+            .with_operation_timeout(DEFAULT_LONG_TIMEOUT)
+            .with_cancel_after(DEFAULT_LONG_TIMEOUT)
+            .with_timeout(DEFAULT_LONG_TIMEOUT)
+        )
+
     def __del__(self):
         self._clean_executor(wait=False)
 
@@ -188,6 +199,7 @@ async def _start_asyncio_driver(self, driver: DriverIO, stub, method):
             requests_iterator,
             stub,
             method,
+            settings=self._stream_settings,
         )
         self._stream_call = stream_call
         self.from_server_grpc = stream_call.__aiter__()
@@ -196,14 +208,29 @@ async def _start_sync_driver(self, driver: Driver, stub, method):
         requests_iterator = AsyncQueueToSyncIteratorAsyncIO(self.from_client_grpc)
         self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
-        stream_call = await to_thread(driver, requests_iterator, stub, method, executor=self._wait_executor)
+        stream_call = await to_thread(
+            driver,
+            requests_iterator,
+            stub,
+            method,
+            executor=self._wait_executor,
+            settings=self._stream_settings,
+        )
         self._stream_call = stream_call
         self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
 
-    async def receive(self) -> Any:
+    async def receive(self, timeout: Optional[int] = None) -> Any:
         # todo handle grpc exceptions and convert it to internal exceptions
         try:
-            grpc_message = await self.from_server_grpc.__anext__()
+            if timeout is None:
+                grpc_message = await self.from_server_grpc.__anext__()
+            else:
+
+                async def get_response():
+                    return await self.from_server_grpc.__anext__()
+
+                grpc_message = await asyncio.wait_for(get_response(), timeout)
+
         except (grpc.RpcError, grpc.aio.AioRpcError) as e:
             raise connection._rpc_error_handler(self._connection_state, e)
 

ydb/_topic_common/test_helpers.py

@@ -15,7 +15,7 @@ def __init__(self):
         self.from_client = asyncio.Queue()
         self._closed = False
 
-    async def receive(self) -> typing.Any:
+    async def receive(self, timeout: typing.Optional[int] = None) -> typing.Any:
         if self._closed:
             raise Exception("read from closed StreamMock")
 

ydb/_topic_reader/topic_reader_asyncio.py

@@ -490,7 +490,11 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
         logger.debug("reader stream %s send init request", self._id)
 
         stream.write(StreamReadMessage.FromClient(client_message=init_message))
-        init_response = await stream.receive()  # type: StreamReadMessage.FromServer
+        try:
+            init_response = await stream.receive(timeout=10)  # type: StreamReadMessage.FromServer
+        except asyncio.TimeoutError:
+            raise TopicReaderError("Timeout waiting for init response")
+
         if isinstance(init_response.server_message, StreamReadMessage.InitResponse):
             self._session_id = init_response.server_message.session_id
             logger.debug("reader stream %s initialized session=%s", self._id, self._session_id)

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -14,7 +14,7 @@
 from . import datatypes, topic_reader_asyncio
 from .datatypes import PublicBatch, PublicMessage
 from .topic_reader import PublicReaderSettings
-from .topic_reader_asyncio import ReaderStream, ReaderReconnector
+from .topic_reader_asyncio import ReaderStream, ReaderReconnector, TopicReaderError
 from .._grpc.grpcwrapper.common_utils import SupportedDriverType, ServerStatus
 from .._grpc.grpcwrapper.ydb_topic import (
     StreamReadMessage,
@@ -1475,6 +1475,46 @@ def logged():
 
         await wait_condition(logged)
 
+    async def test_init_timeout_parameter(self, stream, default_reader_settings):
+        """Test that ReaderStream._start calls stream.receive with timeout=10"""
+        reader = ReaderStream(self.default_reader_reconnector_id, default_reader_settings)
+        init_message = default_reader_settings._init_message()
+
+        # Mock stream.receive to check if timeout=10 is passed
+        with mock.patch.object(stream, 'receive') as mock_receive:
+            mock_receive.return_value = StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.InitResponse(session_id="test_session")
+            )
+
+            await reader._start(stream, init_message)
+
+            # Verify that receive was called with timeout=10
+            mock_receive.assert_called_with(timeout=10)
+
+        await reader.close(False)
+
+    async def test_init_timeout_behavior(self, stream, default_reader_settings):
+        """Test that ReaderStream._start raises TopicReaderError when receive times out"""
+        reader = ReaderStream(self.default_reader_reconnector_id, default_reader_settings)
+        init_message = default_reader_settings._init_message()
+
+        # Mock stream.receive to directly raise TimeoutError when called with timeout=10
+        async def timeout_receive(timeout=None):
+            if timeout == 10:
+                raise asyncio.TimeoutError("Simulated timeout")
+            return StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.InitResponse(session_id="test_session")
+            )
+
+        with mock.patch.object(stream, 'receive', side_effect=timeout_receive):
+            # Should raise TopicReaderError with timeout message
+            with pytest.raises(TopicReaderError, match="Timeout waiting for init response"):
+                await reader._start(stream, init_message)
+
+        await reader.close(False)
+
 
 @pytest.mark.asyncio
 class TestReaderReconnector:

ydb/_topic_writer/topic_writer_asyncio.py

@@ -799,7 +799,11 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamWriteMes
         logger.debug("writer stream %s send init request", self._id)
         stream.write(StreamWriteMessage.FromClient(init_message))
 
-        resp = await stream.receive()
+        try:
+            resp = await stream.receive(timeout=10)
+        except asyncio.TimeoutError:
+            raise TopicWriterError("Timeout waiting for init response")
+
         self._ensure_ok(resp)
         if not isinstance(resp, StreamWriteMessage.InitResponse):
             raise TopicWriterError("Unexpected answer for init request: %s" % resp)

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -231,6 +231,55 @@ async def test_update_token(self, stream: StreamMock):
 
         await writer.close()
 
+    async def test_init_timeout_parameter(self, stream):
+        """Test that WriterAsyncIOStream._start calls stream.receive with timeout=10"""
+        writer_id = 1
+        settings = WriterSettings(PublicWriterSettings("test-topic", "test-producer"))
+
+        # Mock stream.receive to check if timeout=10 is passed
+        with mock.patch.object(stream, 'receive') as mock_receive:
+            mock_receive.return_value = StreamWriteMessage.InitResponse(
+                last_seq_no=0,
+                session_id="test_session",
+                partition_id=1,
+                supported_codecs=[Codec.CODEC_RAW],
+                status=ServerStatus(StatusCode.SUCCESS, [])
+            )
+
+            writer = WriterAsyncIOStream(writer_id, settings)
+            await writer._start(stream, settings.create_init_request())
+
+            # Verify that receive was called with timeout=10
+            mock_receive.assert_called_with(timeout=10)
+
+        await writer.close()
+
+    async def test_init_timeout_behavior(self, stream):
+        """Test that WriterAsyncIOStream._start raises TopicWriterError when receive times out"""
+        writer_id = 1
+        settings = WriterSettings(PublicWriterSettings("test-topic", "test-producer"))
+
+        # Mock stream.receive to directly raise TimeoutError when called with timeout=10
+        async def timeout_receive(timeout=None):
+            if timeout == 10:
+                raise asyncio.TimeoutError("Simulated timeout")
+            return StreamWriteMessage.InitResponse(
+                last_seq_no=0,
+                session_id="test_session",
+                partition_id=1,
+                supported_codecs=[Codec.CODEC_RAW],
+                status=ServerStatus(StatusCode.SUCCESS, [])
+            )
+
+        with mock.patch.object(stream, 'receive', side_effect=timeout_receive):
+            writer = WriterAsyncIOStream(writer_id, settings)
+
+            # Should raise TopicWriterError with timeout message
+            with pytest.raises(TopicWriterError, match="Timeout waiting for init response"):
+                await writer._start(stream, settings.create_init_request())
+
+        # Don't close writer since _start failed and _stream was never set
+
 
 @pytest.mark.asyncio
 class TestWriterAsyncIOReconnector: