address comments

Author: dianfu

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

@@ -126,7 +126,7 @@ def timeout(self, value: Row):
         # note that we use assert_equals instead of assert_equals_sorted
         self.assert_equals(expected, results)
 
-    def test_complete_async_function_with_non_iterable_result(self):
+    def test_non_iterable_result(self):
         self.env.set_parallelism(1)
         ds = self.env.from_collection(
             [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
@@ -147,12 +147,12 @@ def timeout(self, value: Row):
         ds.add_sink(self.test_sink)
         try:
             self.env.execute()
+            self.fail()
         except Exception as e:
             message = str(e)
-            self.assertTrue("The 'result_future' of AsyncFunction should be completed with data of "
-                            "list type" in message)
+            self.assertTrue("The result of AsyncFunction should be of list type" in message)
 
-    def test_complete_async_function_with_exception(self):
+    def test_none_result(self):
         self.env.set_parallelism(1)
         ds = self.env.from_collection(
             [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)],
@@ -162,20 +162,21 @@ def test_complete_async_function_with_exception(self):
         class MyAsyncFunction(AsyncFunction):
 
             async def async_invoke(self, value: Row):
-                raise Exception("encountered an exception")
+                await asyncio.sleep(10)
+                return None
 
             def timeout(self, value: Row):
-                # raise the same exception to make sure test case is stable in all cases
-                raise Exception("encountered an exception")
+                return None
 
         ds = AsyncDataStream.unordered_wait(
-            ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
+            ds, MyAsyncFunction(), Time.seconds(1), 2, Types.INT())
         ds.add_sink(self.test_sink)
         try:
             self.env.execute()
+            self.fail()
         except Exception as e:
             message = str(e)
-            self.assertTrue("Could not complete the element" in message)
+            self.assertTrue("The result of AsyncFunction cannot be none" in message)
 
     def test_raise_exception_in_async_invoke(self):
         self.env.set_parallelism(1)
@@ -198,6 +199,7 @@ def timeout(self, value: Row):
         ds.add_sink(self.test_sink)
         try:
             self.env.execute()
+            self.fail()
         except Exception as e:
             message = str(e)
             self.assertTrue("encountered an exception" in message)
@@ -223,6 +225,7 @@ def timeout(self, value: Row):
         ds.add_sink(self.test_sink)
         try:
             self.env.execute()
+            self.fail()
         except Exception as e:
             message = str(e)
             self.assertTrue("encountered an exception" in message)
@@ -319,6 +322,7 @@ async def async_invoke(self, value: Row):
         try:
             AsyncDataStream.unordered_wait(
                 ds, MyAsyncFunction(), Time.seconds(5), 2, Types.INT())
+            self.fail()
         except Exception as e:
             message = str(e)
             self.assertTrue("AsyncFunction is still not supported for 'thread' mode" in message)

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

@@ -85,23 +85,37 @@ def complete(self, result: List[OUT]):
         if not self._completed.compare_and_set(False, True):
             return
 
+        self._complete_internal(result)
+
+    def _complete_internal(self, result: List[OUT]):
         if isinstance(result, Iterable):
             self._process_results(result)
         else:
+            if result is None:
+                self._exception_handler(
+                    RuntimeError("The result of AsyncFunction cannot be none, "
+                                 "please check the methods 'async_invoke' and "
+                                 "'timeout' of class '%s'." % self._classname))
+            else:
+                self._exception_handler(
+                    RuntimeError("The result of AsyncFunction should be of list type, "
+                                 "please check the methods 'async_invoke' and "
+                                 "'timeout' of class '%s'." % self._classname))
+
             # complete with empty result, so that we remove timer and move ahead processing
             self._process_results([])
 
-            raise RuntimeError("The 'result_future' of AsyncFunction should be completed with "
-                               "data of list type, please check the methods 'async_invoke' and "
-                               "'timeout' of class '%s'." % self._classname)
-
     def complete_exceptionally(self, error: Exception):
         # already completed, so ignore exception
         if not self._completed.compare_and_set(False, True):
             return
 
-        self._exception_handler(
-            Exception("Could not complete the element:" + str(self._record), error))
+        self._complete_exceptionally_internal(error)
+
+    def _complete_exceptionally_internal(self, error: Exception):
+        self._exception_handler(Exception(
+            "Error happens inside the class '%s' during handling input '%s'"
+            % (self._classname, str(self._record)), error))
 
         #  complete with empty result, so that we remove timer and move ahead processing
         self._process_results([])
@@ -114,12 +128,14 @@ def _process_results(self, result: List[OUT]):
         self._result_future.complete(result)
 
     def _timer_triggered(self):
-        if not self._completed.get():
-            try:
-                result = self._timeout_func(self._record)
-                self._result_future.complete(result)
-            except Exception as e:
-                self._result_future.complete_exceptionally(e)
+        if not self._completed.compare_and_set(False, True):
+            return
+
+        try:
+            result = self._timeout_func(self._record)
+            self._complete_internal(result)
+        except Exception as error:
+            self._complete_exceptionally_internal(error)
 
 
 class RetryableResultHandler(ResultFuture, Generic[IN, OUT]):
@@ -166,7 +182,7 @@ def _process_retry(self, result: Optional[List[OUT]], error: Optional[Exception]
             self._delayed_retry_timer = threading.Timer(next_backoff_time_sec, self._do_retry)
             self._delayed_retry_timer.start()
         else:
-            if result is not None:
+            if error is None:
                 self._result_handler.complete(result)
             else:
                 self._result_handler.complete_exceptionally(error)
@@ -189,18 +205,20 @@ def _timer_triggered(self):
         """
         Rewrite the timeout process to deal with retry state.
         """
-        if not self._result_handler._completed.get():
-            # cancel delayed retry timer first
-            self._cancel_retry_timer()
+        if not self._result_handler._completed.compare_and_set(False, True):
+            return
 
-            # force reset _retry_awaiting to prevent the handler to trigger retry unnecessarily
-            self._retry_awaiting.set(False)
+        # cancel delayed retry timer first
+        self._cancel_retry_timer()
 
-            try:
-                result = self._result_handler._timeout_func(self._result_handler._record)
-                self._result_handler.complete(result)
-            except Exception as e:
-                self._result_handler.complete_exceptionally(e)
+        # force reset _retry_awaiting to prevent the handler to trigger retry unnecessarily
+        self._retry_awaiting.set(False)
+
+        try:
+            result = self._result_handler._timeout_func(self._result_handler._record)
+            self._result_handler._complete_internal(result)
+        except Exception as e:
+            self._result_handler._complete_exceptionally_internal(e)
 
 
 class Emitter(threading.Thread):