[FLINK-38615][python] Optimize AsyncFunction to remove ResultFuture

Author: dianfu

flink-python/pyflink/datastream/__init__.py

@@ -272,7 +272,7 @@
                                           KeyedCoProcessFunction, AggregateFunction, WindowFunction,
                                           ProcessWindowFunction, BroadcastProcessFunction,
                                           KeyedBroadcastProcessFunction, AsyncFunction,
-                                          ResultFuture, AsyncRetryPredicate, AsyncRetryStrategy)
+                                          AsyncRetryPredicate, AsyncRetryStrategy)
 from pyflink.datastream.slot_sharing_group import SlotSharingGroup, MemorySize
 from pyflink.datastream.state_backend import (StateBackend, CustomStateBackend,
                                               PredefinedOptions, HashMapStateBackend,
@@ -345,7 +345,6 @@
     'SlotSharingGroup',
     'MemorySize',
     'OutputTag',
-    'ResultFuture',
     'AsyncRetryPredicate',
     'AsyncRetryStrategy'
 ]

flink-python/pyflink/datastream/functions.py

@@ -56,7 +56,6 @@
     'KeyedBroadcastProcessFunction',
     'AsyncFunction',
     'AsyncFunctionDescriptor',
-    'ResultFuture',
     'AsyncRetryPredicate',
     'AsyncRetryStrategy',
 ]
@@ -979,68 +978,32 @@ def exponential_backoff(
             result_predicate, exception_predicate)
 
 
-class ResultFuture(Generic[OUT]):
-    """
-    Collects data / error in user codes while processing async i/o.
-    """
-
-    @abstractmethod
-    def complete(self, result: List[OUT]):
-        """
-        Completes the result future with a collection of result objects.
-
-        Note that it should be called for exactly one time in the user code. Calling this function
-        for multiple times will cause data lose.
-
-        Put all results in a collection and then emit output.
-
-        :param result: A list of results.
-        """
-        pass
-
-    @abstractmethod
-    def complete_exceptionally(self, error: Exception):
-        """
-        Completes the result future exceptionally with an exception.
-
-        :param error: An Exception object.
-        """
-        pass
-
-
 class AsyncFunction(Function, Generic[IN, OUT]):
     """
     A function to trigger Async I/O operation.
 
-    For each #async_invoke, an async io operation can be triggered, and once it has been done, the
-    result can be collected by calling :func:`~ResultFuture.complete`. For each async operation, its
+    For each #async_invoke, an async io operation can be triggered. For each async operation, its
     context is stored in the operator immediately after invoking #async_invoke, avoiding blocking
     for each stream input as long as the internal buffer is not full.
-
-    :class:`~ResultFuture` can be passed into callbacks or futures to collect the result data. An
-    error can also be propagated to the async IO operator by
-    :func:`~ResultFuture.complete_exceptionally`.
     """
 
     @abstractmethod
-    async def async_invoke(self, value: IN, result_future: ResultFuture[OUT]):
+    async def async_invoke(self, value: IN) -> List[OUT]:
         """
         Trigger async operation for each stream input.
         In case of a user code error. You can raise an exception to make the task fail and
         trigger fail-over process.
 
         :param value: Input element coming from an upstream task.
-        :param result_future: A future to be completed with the result data.
         """
         pass
 
-    def timeout(self, value: IN, result_future: ResultFuture[OUT]):
+    def timeout(self, value: IN) -> List[OUT]:
         """
-        In case :func:`~ResultFuture.async_invoke` timeout occurred. By default, the result future
-        is exceptionally completed with a timeout exception.
+        In case :func:`~ResultFuture.async_invoke` timeout occurred. By default, it raises
+        a timeout exception.
         """
-        result_future.complete_exceptionally(
-            TimeoutError("Async function call has timed out for input: " + str(value)))
+        raise TimeoutError("Async function call has timed out for input: " + str(value))
 
 
 class AsyncFunctionDescriptor(object):

flink-python/pyflink/datastream/tests/test_async_function.py

@@ -20,7 +20,7 @@
 from typing import List
 
 from pyflink.common import Types, Row, Time, Configuration, WatermarkStrategy
-from pyflink.datastream import AsyncDataStream, AsyncFunction, ResultFuture, \
+from pyflink.datastream import AsyncDataStream, AsyncFunction, \
     StreamExecutionEnvironment, AsyncRetryStrategy
 from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction, \
     SecondColumnTimestampAssigner
@@ -54,12 +54,12 @@ def test_unordered_mode(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(2)
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete([value[0] + value[1]])
+            def timeout(self, value: Row):
+                return [value[0] + value[1]]
 
         ds = AsyncDataStream.unordered_wait(
             ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
@@ -78,12 +78,12 @@ def test_ordered_mode(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(random.randint(1, 2))
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete([value[0] + value[1]])
+            def timeout(self, value: Row):
+                return [value[0] + value[1]]
 
         ds = AsyncDataStream.ordered_wait(
             ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
@@ -110,12 +110,12 @@ def test_watermark(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(random.randint(1, 3))
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete([value[0] + value[1]])
+            def timeout(self, value: Row):
+                return [value[0] + value[1]]
 
         ds = AsyncDataStream.unordered_wait(
             ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
@@ -135,12 +135,12 @@ def test_complete_async_function_with_non_iterable_result(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(2)
-                result_future.complete(value[0] + value[1])
+                return value[0] + value[1]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete(value[0] + value[1])
+            def timeout(self, value: Row):
+                return value[0] + value[1]
 
         ds = AsyncDataStream.unordered_wait(
             ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
@@ -161,12 +161,12 @@ def test_complete_async_function_with_exception(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete_exceptionally(Exception("encountered an exception"))
+            async def async_invoke(self, value: Row):
+                raise Exception("encountered an exception")
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
+            def timeout(self, value: Row):
                 # raise the same exception to make sure test case is stable in all cases
-                result_future.complete_exceptionally(Exception("encountered an exception"))
+                raise Exception("encountered an exception")
 
         ds = AsyncDataStream.unordered_wait(
             ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
@@ -186,10 +186,10 @@ def test_raise_exception_in_async_invoke(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 raise Exception("encountered an exception")
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
+            def timeout(self, value: Row):
                 # raise the same exception to make sure test case is stable in all cases
                 raise Exception("encountered an exception")
 
@@ -211,11 +211,11 @@ def test_raise_exception_in_timeout(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(10)
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
+            def timeout(self, value: Row):
                 raise Exception("encountered an exception")
 
         ds = AsyncDataStream.unordered_wait(
@@ -236,12 +236,12 @@ def test_processing_timeout(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(10)
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete([value[0] - value[1]])
+            def timeout(self, value: Row):
+                return [value[0] - value[1]]
 
         ds = AsyncDataStream.unordered_wait(
             ds, MyAsyncFunction(), Time.seconds(1), 2, Types.INT())
@@ -264,19 +264,19 @@ def __init__(self):
                 self.retries_1 = {}
                 self.retries_2 = {}
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(1)
                 if value in self.retries_2:
-                    result_future.complete([value[0] + value[1]])
+                    return [value[0] + value[1]]
                 elif value in self.retries_1:
                     self.retries_2[value] = True
-                    result_future.complete([value[0] + value[1] + 1])
+                    return [value[0] + value[1] + 1]
                 else:
                     self.retries_1[value] = True
-                    result_future.complete_exceptionally(ValueError("failed the first time"))
+                    raise ValueError("failed the first time")
 
-            def timeout(self, value: Row, result_future: ResultFuture[int]):
-                result_future.complete([value[0] + value[1]])
+            def timeout(self, value: Row):
+                return [value[0] + value[1]]
 
         def result_predicate(result: List[int]):
             return result[0] % 2 == 1
@@ -312,9 +312,9 @@ def test_run_async_function_in_thread_mode(self):
 
         class MyAsyncFunction(AsyncFunction):
 
-            async def async_invoke(self, value: Row, result_future: ResultFuture[int]):
+            async def async_invoke(self, value: Row):
                 await asyncio.sleep(2)
-                result_future.complete([value[0] + value[1]])
+                return [value[0] + value[1]]
 
         try:
             AsyncDataStream.unordered_wait(

flink-python/pyflink/fn_execution/datastream/process/async_function/operation.py

@@ -21,10 +21,10 @@
 from datetime import datetime
 from typing import TypeVar, Generic, List, Iterable, Callable, Optional
 
-from pyflink.datastream import RuntimeContext, ResultFuture
+from pyflink.datastream import RuntimeContext
 from pyflink.datastream.functions import AsyncFunctionDescriptor, AsyncRetryStrategy
 from pyflink.fn_execution.datastream.process.async_function.queue import \
-    UnorderedStreamElementQueue, StreamElementQueue, OrderedStreamElementQueue
+    UnorderedStreamElementQueue, StreamElementQueue, OrderedStreamElementQueue, ResultFuture
 from pyflink.fn_execution.datastream.process.operations import Operation
 from pyflink.fn_execution.datastream.process.runtime_context import StreamingRuntimeContext
 
@@ -63,7 +63,7 @@ class ResultHandler(ResultFuture, Generic[IN, OUT]):
 
     def __init__(self,
                  classname: str,
-                 timeout_func: Callable[[IN, ResultFuture[OUT]], None],
+                 timeout_func: Callable[[IN], List[OUT]],
                  exception_handler: Callable[[Exception], None],
                  record: IN,
                  result_future: ResultFuture[OUT]):
@@ -115,7 +115,11 @@ def _process_results(self, result: List[OUT]):
 
     def _timer_triggered(self):
         if not self._completed.get():
-            self._timeout_func(self._record, self)
+            try:
+                result = self._timeout_func(self._record)
+                self._result_future.complete(result)
+            except Exception as e:
+                self._result_future.complete_exceptionally(e)
 
 
 class RetryableResultHandler(ResultFuture, Generic[IN, OUT]):
@@ -192,7 +196,11 @@ def _timer_triggered(self):
             # force reset _retry_awaiting to prevent the handler to trigger retry unnecessarily
             self._retry_awaiting.set(False)
 
-            self._result_handler._timeout_func(self._result_handler._record, self)
+            try:
+                result = self._result_handler._timeout_func(self._result_handler._record)
+                self._result_handler._result_future.complete(result)
+            except Exception as e:
+                self._result_handler._result_future.complete_exceptionally(e)
 
 
 class Emitter(threading.Thread):
@@ -223,9 +231,8 @@ def stop(self):
 
 
 class AsyncFunctionRunner(threading.Thread):
-    def __init__(self, exception_handler: Callable[[Exception], None]):
+    def __init__(self):
         super().__init__()
-        self._exception_handler = exception_handler
         self._loop = None
         self._ready = threading.Event()
 
@@ -251,14 +258,15 @@ def stop(self):
             self._loop.call_soon_threadsafe(self._loop.stop)
             self.join(timeout=1.0)
 
-    async def exception_handler_wrapper(self, async_function, *arg):
+    async def exception_handler_wrapper(self, async_function, record, result_handler):
         try:
-            await async_function(*arg)
+            result = await async_function(record)
+            result_handler.complete(result)
         except Exception as e:
-            self._exception_handler(e)
+            result_handler.complete_exceptionally(e)
 
-    def run_async(self, async_function, *arg):
-        wrapped_function = self.exception_handler_wrapper(async_function, *arg)
+    def run_async(self, async_function, record, result_handler):
+        wrapped_function = self.exception_handler_wrapper(async_function, record, result_handler)
         asyncio.run_coroutine_threadsafe(wrapped_function, self._loop)
 
 
@@ -306,7 +314,7 @@ def open(self):
         self._emitter.daemon = True
         self._emitter.start()
 
-        self._async_function_runner = AsyncFunctionRunner(self._mark_exception)
+        self._async_function_runner = AsyncFunctionRunner()
         self._async_function_runner.daemon = True
         self._async_function_runner.start()
         self._async_function_runner.wait_ready()