Merge pull request #227 V3 fix sync call timeout

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):
@@ -60,3 +60,88 @@ def start_event_loop():
 
         _shared_event_loop = event_loop_set_done.result()
         return _shared_event_loop
+
+
+class CallFromSyncToAsync:
+    _loop: asyncio.AbstractEventLoop
+
+    def __init__(self, loop: asyncio.AbstractEventLoop):
+        self._loop = loop
+
+    def unsafe_call_with_future(
+        self, coro: typing.Coroutine
+    ) -> concurrent.futures.Future:
+        """
+        returned result from coro may be lost
+        """
+        return asyncio.run_coroutine_threadsafe(coro, self._loop)
+
+    def unsafe_call_with_result(self, coro: typing.Coroutine, timeout: TimeoutType):
+        """
+        returned result from coro may be lost by race future cancel by timeout and return value from coroutine
+        """
+        f = self.unsafe_call_with_future(coro)
+        try:
+            return f.result(timeout)
+        except concurrent.futures.TimeoutError:
+            raise TimeoutError()
+        finally:
+            if not f.done():
+                f.cancel()
+
+    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.
+        """
+
+        if timeout is not None and timeout <= 0:
+            return self._safe_call_fast(coro)
+
+        async def call_coro():
+            task = self._loop.create_task(coro)
+            try:
+                res = await asyncio.wait_for(task, timeout)
+                return res
+            except asyncio.TimeoutError:
+                try:
+                    res = await task
+                    return res
+                except asyncio.CancelledError:
+                    pass
+
+                # return builtin TimeoutError instead of asyncio.TimeoutError
+                raise TimeoutError()
+
+        return asyncio.run_coroutine_threadsafe(call_coro(), self._loop).result()
+
+    def _safe_call_fast(self, coro: typing.Coroutine):
+        """
+        no lost returned value from coro, but may be slower especially timeout latency - it wait coroutine cancelation.
+        Wait coroutine result only one loop.
+        """
+        res = concurrent.futures.Future()
+
+        async def call_coro():
+            try:
+                res.set_result(await coro)
+            except asyncio.CancelledError:
+                res.set_exception(TimeoutError())
+
+        coro_future = asyncio.run_coroutine_threadsafe(call_coro(), self._loop)
+        asyncio.run_coroutine_threadsafe(asyncio.sleep(0), self._loop).result()
+        coro_future.cancel()
+        return res.result()
+
+    def call_sync(self, callback: typing.Callable[[], typing.Any]) -> typing.Any:
+        result = concurrent.futures.Future()
+
+        def call_callback():
+            try:
+                res = callback()
+                result.set_result(res)
+            except BaseException as err:
+                result.set_exception(err)
+
+        self._loop.call_soon_threadsafe(call_callback)
+
+        return result.result()

ydb/_topic_common/common_test.py

@@ -1,9 +1,12 @@
 import asyncio
+import threading
+import time
 import typing
 
 import grpc
 import pytest
 
+from .common import CallFromSyncToAsync
 from .._grpc.grpcwrapper.common_utils import (
     GrpcWrapperAsyncIO,
     ServerStatus,
@@ -25,6 +28,23 @@
     )
 
 
+@pytest.fixture()
+def separate_loop():
+    loop = asyncio.new_event_loop()
+
+    def run_loop():
+        loop.run_forever()
+        pass
+
+    t = threading.Thread(target=run_loop, name="test separate loop")
+    t.start()
+
+    yield loop
+
+    loop.call_soon_threadsafe(lambda: loop.stop())
+    t.join()
+
+
 @pytest.mark.asyncio
 class Test:
     async def test_callback_from_asyncio(self):
@@ -111,3 +131,160 @@ def test_failed(self):
         assert not status.is_success()
         with pytest.raises(issues.Overloaded):
             issues._process_response(status)
+
+
+@pytest.mark.asyncio
+class TestCallFromSyncToAsync:
+    @pytest.fixture()
+    def caller(self, separate_loop):
+        return CallFromSyncToAsync(separate_loop)
+
+    def test_unsafe_call_with_future(self, separate_loop, caller):
+        callback_loop = None
+
+        async def callback():
+            nonlocal callback_loop
+            callback_loop = asyncio.get_running_loop()
+            return 1
+
+        f = caller.unsafe_call_with_future(callback())
+
+        assert f.result() == 1
+        assert callback_loop is separate_loop
+
+    def test_unsafe_call_with_result_ok(self, separate_loop, caller):
+        callback_loop = None
+
+        async def callback():
+            nonlocal callback_loop
+            callback_loop = asyncio.get_running_loop()
+            return 1
+
+        res = caller.unsafe_call_with_result(callback(), None)
+
+        assert res == 1
+        assert callback_loop is separate_loop
+
+    def test_unsafe_call_with_result_timeout(self, separate_loop, caller):
+        timeout = 0.01
+        callback_loop = None
+
+        async def callback():
+            nonlocal callback_loop
+            callback_loop = asyncio.get_running_loop()
+            await asyncio.sleep(1)
+            return 1
+
+        start = time.monotonic()
+        with pytest.raises(TimeoutError):
+            caller.unsafe_call_with_result(callback(), timeout)
+        finished = time.monotonic()
+
+        assert callback_loop is separate_loop
+        assert finished - start > timeout
+
+    def test_safe_call_with_result_ok(self, separate_loop, caller):
+        callback_loop = None
+
+        async def callback():
+            nonlocal callback_loop
+            callback_loop = asyncio.get_running_loop()
+            return 1
+
+        res = caller.safe_call_with_result(callback(), 1)
+
+        assert res == 1
+        assert callback_loop is separate_loop
+
+    def test_safe_call_with_result_timeout(self, separate_loop, caller):
+        timeout = 0.01
+        callback_loop = None
+        cancelled = False
+
+        async def callback():
+            nonlocal callback_loop, cancelled
+            callback_loop = asyncio.get_running_loop()
+            try:
+                await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                cancelled = True
+                raise
+
+            return 1
+
+        start = time.monotonic()
+        with pytest.raises(TimeoutError):
+            caller.safe_call_with_result(callback(), timeout)
+        finished = time.monotonic()
+
+        # wait one loop for handle task cancelation
+        asyncio.run_coroutine_threadsafe(asyncio.sleep(0), separate_loop)
+
+        assert callback_loop is separate_loop
+        assert finished - start > timeout
+        assert cancelled
+
+    def test_safe_callback_with_0_timeout_ok(self, separate_loop, caller):
+        callback_loop = None
+
+        async def f1():
+            return 1
+
+        async def f2():
+            return await f1()
+
+        async def callback():
+            nonlocal callback_loop
+            callback_loop = asyncio.get_running_loop()
+            return await f2()
+
+        res = caller.safe_call_with_result(callback(), 0)
+        assert callback_loop is separate_loop
+        assert res == 1
+
+    def test_safe_callback_with_0_timeout_timeout(self, separate_loop, caller):
+        callback_loop = None
+        cancelled = False
+
+        async def callback():
+            try:
+                nonlocal callback_loop, cancelled
+
+                callback_loop = asyncio.get_running_loop()
+                await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                cancelled = True
+                raise
+
+        with pytest.raises(TimeoutError):
+            caller.safe_call_with_result(callback(), 0)
+
+        assert callback_loop is separate_loop
+        assert cancelled
+
+    def test_call_sync_ok(self, separate_loop, caller):
+        callback_eventloop = None
+
+        def callback():
+            nonlocal callback_eventloop
+            callback_eventloop = asyncio.get_running_loop()
+            return 1
+
+        res = caller.call_sync(callback)
+        assert callback_eventloop is separate_loop
+        assert res == 1
+
+    def test_call_sync_error(self, separate_loop, caller):
+        callback_eventloop = None
+
+        class TestError(RuntimeError):
+            pass
+
+        def callback():
+            nonlocal callback_eventloop
+            callback_eventloop = asyncio.get_running_loop()
+            raise TestError
+
+        with pytest.raises(TestError):
+            caller.call_sync(callback)
+        assert callback_eventloop is separate_loop