impl writer

Author: rekby

ydb/_topic_common/common.py

@@ -7,7 +7,7 @@
 from .. import operation, issues
 from .._grpc.grpcwrapper.common_utils import IFromProtoWithProtoType
 
-TimeoutType = typing.Union[int, float]
+TimeoutType = typing.Union[int, float, None]
 
 
 def wrap_operation(rpc_state, response_pb, driver=None):
@@ -86,7 +86,7 @@ def unsafe_call_with_result(self, coro: typing.Coroutine, timeout: typing.Union[
         finally:
             f.cancel()
 
-    def safe_call_with_result(self, coro: typing.Coroutine, timeout: typing.Union[int, float]):
+    def safe_call_with_result(self, coro: typing.Coroutine, timeout: TimeoutType):
         """
         no lost returned value from coro, but may be slower especially timeout latency - it wait coroutine cancelation.
         """

ydb/_topic_writer/topic_writer.py

@@ -176,6 +176,11 @@ def __init__(self, message: str):
         super(TopicWriterError, self).__init__(message)
 
 
+class TopicWriterClosedError(ydb.Error):
+    def __init__(self):
+        super(TopicWriterClosedError, self).__init__("Topic writer already closed")
+
+
 class TopicWriterRepeatableError(TopicWriterError):
     pass
 

ydb/_topic_writer/topic_writer_sync.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import typing
 from concurrent.futures import Future
 from typing import Union, List, Optional, Coroutine
 
@@ -10,15 +11,15 @@
     TopicWriterError,
     PublicWriterInitInfo,
     PublicWriteResult,
-    Message,
+    Message, TopicWriterClosedError,
 )
 
 from .topic_writer_asyncio import WriterAsyncIO
-from .._topic_common.common import _get_shared_event_loop, TimeoutType
+from .._topic_common.common import _get_shared_event_loop, TimeoutType, CallFromSyncToAsync
 
 
 class WriterSync:
-    _loop: asyncio.AbstractEventLoop
+    _caller: CallFromSyncToAsync
     _async_writer: WriterAsyncIO
     _closed: bool
 
@@ -33,80 +34,79 @@ def __init__(
         self._closed = False
 
         if eventloop:
-            self._loop = eventloop
+            loop = eventloop
         else:
-            self._loop = _get_shared_event_loop()
+            loop = _get_shared_event_loop()
+
+        self._caller = CallFromSyncToAsync(loop)
 
         async def create_async_writer():
             return WriterAsyncIO(driver, settings)
 
-        self._async_writer = asyncio.run_coroutine_threadsafe(
-            create_async_writer(), self._loop
-        ).result()
+        self._async_writer = self._caller.safe_call_with_result(create_async_writer(), None)
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def _call(self, coro):
+    def close(self, *, flush: bool = True, timeout: typing.Union[int, float, None] = None):
         if self._closed:
-            raise TopicWriterError("writer is closed")
+            return
 
-        return asyncio.run_coroutine_threadsafe(coro, self._loop)
+        self._closed = True
 
-    def _call_sync(self, coro: Coroutine, timeout):
-        f = self._call(coro)
-        try:
-            return f.result(timeout)
-        except TimeoutError:
-            f.cancel()
-            raise
+        self._caller.safe_call_with_result(self._async_writer.close(flush=flush), timeout)
 
-    def close(self, flush: bool = True):
+    def _check_closed(self):
         if self._closed:
-            return
+            raise TopicWriterClosedError()
 
-        self._closed = True
+    def async_flush(self) -> Future:
+        self._check_closed()
 
-        # for no call self._call_sync on closed object
-        asyncio.run_coroutine_threadsafe(
-            self._async_writer.close(flush=flush), self._loop
-        ).result()
+        return self._caller.unsafe_call_with_future(self._async_writer.flush())
 
-    def async_flush(self) -> Future:
-        if self._closed:
-            raise TopicWriterError("writer is closed")
-        return self._call(self._async_writer.flush())
+    def flush(self, *, timeout=None):
+        self._check_closed()
 
-    def flush(self, timeout=None):
-        self._call_sync(self._async_writer.flush(), timeout)
+        return self._caller.unsafe_call_with_result(self._async_writer.flush(), timeout)
 
     def async_wait_init(self) -> Future[PublicWriterInitInfo]:
-        return self._call(self._async_writer.wait_init())
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_future(self._async_writer.wait_init())
 
-    def wait_init(self, timeout: Optional[TimeoutType] = None) -> PublicWriterInitInfo:
-        return self._call_sync(self._async_writer.wait_init(), timeout)
+    def wait_init(self, *, timeout: TimeoutType = None) -> PublicWriterInitInfo:
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_writer.wait_init(), timeout)
 
     def write(
         self,
         messages: Union[Message, List[Message]],
-        timeout: Union[float, None] = None,
+        timeout: TimeoutType = None,
     ):
-        self._call_sync(self._async_writer.write(messages), timeout=timeout)
+        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]],
     ) -> Future[Union[PublicWriteResult, List[PublicWriteResult]]]:
-        return self._call(self._async_writer.write_with_ack(messages))
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_future(self._async_writer.write_with_ack(messages))
 
     def write_with_ack(
         self,
         messages: Union[Message, List[Message]],
         timeout: Union[float, None] = None,
     ) -> Union[PublicWriteResult, List[PublicWriteResult]]:
-        return self._call_sync(
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(
             self._async_writer.write_with_ack(messages), timeout=timeout
         )