rewrite reader to sync caller

Author: rekby

ydb/_topic_reader/topic_reader_sync.py

@@ -4,7 +4,8 @@
 from typing import List, Union, Iterable, Optional, Coroutine
 
 from ydb._grpc.grpcwrapper.common_utils import SupportedDriverType
-from ydb._topic_common.common import _get_shared_event_loop
+from ydb._topic_common.common import _get_shared_event_loop, CallFromSyncToAsync
+from ydb._topic_reader import datatypes
 from ydb._topic_reader.datatypes import PublicMessage, PublicBatch, ICommittable
 from ydb._topic_reader.topic_reader import (
     PublicReaderSettings,
@@ -18,29 +19,31 @@
 
 
 class TopicReaderSync:
-    _loop: asyncio.AbstractEventLoop
+    _caller: CallFromSyncToAsync
     _async_reader: PublicAsyncIOReader
     _closed: bool
 
     def __init__(
-        self,
-        driver: SupportedDriverType,
-        settings: PublicReaderSettings,
-        *,
-        eventloop: Optional[asyncio.AbstractEventLoop] = None,
+            self,
+            driver: SupportedDriverType,
+            settings: PublicReaderSettings,
+            *,
+            eventloop: Optional[asyncio.AbstractEventLoop] = None,
     ):
         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_reader():
             return PublicAsyncIOReader(driver, settings)
 
         self._async_reader = asyncio.run_coroutine_threadsafe(
-            create_reader(), self._loop
+            create_reader(), loop
         ).result()
 
     def __del__(self):
@@ -52,26 +55,6 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def _call(self, coro) -> concurrent.futures.Future:
-        """
-        Call async function and return future fow wait result
-        """
-        if self._closed:
-            raise TopicReaderClosedError()
-
-        return asyncio.run_coroutine_threadsafe(coro, self._loop)
-
-    def _call_sync(self, coro: Coroutine, timeout):
-        """
-        Call async function, wait and return result
-        """
-        f = self._call(coro)
-        try:
-            return f.result(timeout)
-        except TimeoutError:
-            f.cancel()
-            raise
-
     def async_sessions_stat(self) -> concurrent.futures.Future:
         """
         Receive stat from the server, return feature.
@@ -87,7 +70,7 @@ async def sessions_stat(self) -> List[SessionStat]:
         raise NotImplementedError()
 
     def messages(
-        self, *, timeout: Union[float, None] = None
+            self, *, timeout: Union[float, None] = None
     ) -> Iterable[PublicMessage]:
         """
         todo?
@@ -121,11 +104,11 @@ def async_wait_message(self) -> concurrent.futures.Future:
         raise NotImplementedError()
 
     def batches(
-        self,
-        *,
-        max_messages: Union[int, None] = None,
-        max_bytes: Union[int, None] = None,
-        timeout: Union[float, None] = None,
+            self,
+            *,
+            max_messages: Union[int, None] = None,
+            max_bytes: Union[int, None] = None,
+            timeout: Union[float, None] = None,
     ) -> Iterable[PublicBatch]:
         """
         Block until receive new batch.
@@ -137,11 +120,11 @@ def batches(
         raise NotImplementedError()
 
     def receive_batch(
-        self,
-        *,
-        max_messages: typing.Union[int, None] = None,
-        max_bytes: typing.Union[int, None] = None,
-        timeout: Union[float, None] = None,
+            self,
+            *,
+            max_messages: typing.Union[int, None] = None,
+            max_bytes: typing.Union[int, None] = None,
+            timeout: Union[float, None] = None,
     ) -> Union[PublicBatch, None]:
         """
         Get one messages batch from reader
@@ -150,37 +133,42 @@ def receive_batch(
         if no new message in timeout seconds (default - infinite): raise TimeoutError()
         if timeout <= 0 - it will fast non block method, get messages from internal buffer only.
         """
-        return self._call_sync(
-            self._async_reader.receive_batch(
-                max_messages=max_messages, max_bytes=max_bytes
-            ),
-            timeout,
-        )
+        self._check_closed()
 
-    def commit(self, mess: ICommittable):
+        return self._caller.safe_call_with_result(
+            self._async_reader.receive_batch(max_messages=max_messages, max_bytes=max_bytes),
+            timeout)
+
+    def commit(self, mess: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
         """
         Put commit message to internal buffer.
 
         For the method no way check the commit result
         (for example if lost connection - commits will not re-send and committed messages will receive again)
         """
-        self._call_sync(self._async_reader.commit(mess), None)
+        self._check_closed()
+
+        self._caller.call_sync(self._async_reader.commit(mess))
 
     def commit_with_ack(
-        self, mess: ICommittable
+            self, mess: ICommittable, timeout: typing.Union[int, float, None] = None
     ) -> Union[CommitResult, List[CommitResult]]:
         """
         write commit message to a buffer and wait ack from the server.
 
         if receive in timeout seconds (default - infinite): raise TimeoutError()
         """
-        return self._call_sync(self._async_reader.commit_with_ack(mess), None)
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_reader.commit_with_ack(mess), timeout)
 
-    def async_commit_with_ack(self, mess: ICommittable) -> concurrent.futures.Future:
+    def async_commit_with_ack(self, mess: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]) -> concurrent.futures.Future:
         """
         write commit message to a buffer and return Future for wait result.
         """
-        return self._call(self._async_reader.commit_with_ack(mess), None)
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_future(self._async_reader.commit_with_ack(mess))
 
     def async_flush(self) -> concurrent.futures.Future:
         """
@@ -194,12 +182,14 @@ def flush(self):
         """
         raise NotImplementedError()
 
-    def close(self):
+    def close(self, *, timeout: typing.Union[int, float, None] = None):
         if self._closed:
             return
+
         self._closed = True
 
-        # for no call self._call_sync on closed object
-        asyncio.run_coroutine_threadsafe(
-            self._async_reader.close(), self._loop
-        ).result()
+        self._caller.safe_call_with_result(self._async_reader.close(), timeout)
+
+    def _check_closed(self):
+        if self._closed:
+            raise TopicReaderClosedError()