fix: handle cancel return values in FlexConnect

When cancelling, be more transparent about whether the cancellation
actually happened.
Also handle missing tasks: that now yields BAD_ARGUMENT exceptions.

JIRA: CQ-1124
risk: low

Author: no23reason

gooddata-flexconnect/gooddata_flexconnect/function/flight_methods.py

@@ -89,14 +89,21 @@ def _prepare_task(
             cmd=submit_invocation.command,
         )
 
-    def _prepare_flight_info(self, task_result: TaskExecutionResult) -> pyarrow.flight.FlightInfo:
+    def _prepare_flight_info(
+        self, task_id: str, task_result: Optional[TaskExecutionResult]
+    ) -> pyarrow.flight.FlightInfo:
+        if task_result is None:
+            raise ErrorInfo.for_reason(
+                ErrorCode.BAD_ARGUMENT, f"Task with id '{task_id}' does not exist."
+            ).to_user_error()
+
         if task_result.error is not None:
             raise task_result.error.as_flight_error()
 
         if task_result.cancelled:
             raise ErrorInfo.for_reason(
                 ErrorCode.COMMAND_CANCELLED,
-                f"FlexConnect function invocation was cancelled. Invocation task was: '{task_result.task_id}'.",
+                f"FlexConnect function invocation was cancelled. Invocation task was: '{task_id}'.",
             ).to_server_error()
 
         result = task_result.result
@@ -107,7 +114,7 @@ def _prepare_flight_info(self, task_result: TaskExecutionResult) -> pyarrow.flig
             descriptor=pyarrow.flight.FlightDescriptor.for_command(task_result.cmd),
             endpoints=[
                 pyarrow.flight.FlightEndpoint(
-                    ticket=pyarrow.flight.Ticket(ticket=orjson.dumps({"task_id": task_result.task_id})),
+                    ticket=pyarrow.flight.Ticket(ticket=orjson.dumps({"task_id": task_id})),
                     locations=[self._ctx.location],
                 )
             ],
@@ -140,20 +147,16 @@ def get_flight_info(
 
         if isinstance(invocation, CancelInvocation):
             # cancel the given task and raise cancellation exception
-            self._ctx.task_executor.cancel(invocation.task_id)
+            if self._ctx.task_executor.cancel(invocation.task_id):
+                raise ErrorInfo.for_reason(
+                    ErrorCode.COMMAND_CANCELLED, "FlexConnect function invocation was cancelled."
+                ).to_cancelled_error()
             raise ErrorInfo.for_reason(
-                ErrorCode.COMMAND_CANCELLED, "FlexConnect function invocation was cancelled."
+                ErrorCode.COMMAND_CANCEL_NOT_POSSIBLE, "FlexConnect function invocation could not be cancelled."
             ).to_cancelled_error()
         elif isinstance(invocation, RetryInvocation):
             # retry descriptor: extract the task_id, do not submit it again and do one polling iteration
             task_id = invocation.task_id
-            # for retries, we also need to check the call deadline for the whole call duration
-            task_timestamp = self._ctx.task_executor.get_task_submitted_timestamp(task_id)
-            if task_timestamp is not None and time.perf_counter() - task_timestamp > self._call_deadline:
-                self._ctx.task_executor.cancel(task_id)
-                raise ErrorInfo.for_reason(
-                    ErrorCode.TIMEOUT, f"GetFlightInfo timed out while waiting for task {task_id}."
-                ).to_timeout_error()
         elif isinstance(invocation, SubmitInvocation):
             # basic first-time submit: submit the task and do one polling iteration.
             # do not check call deadline to give it a chance to wait for the result at least once
@@ -171,8 +174,18 @@ def get_flight_info(
 
         try:
             task_result = self._ctx.task_executor.wait_for_result(task_id, timeout=self._poll_interval)
-            return self._prepare_flight_info(task_result)
+            return self._prepare_flight_info(task_id, task_result)
         except TimeoutError:
+            # first, check the call deadline for the whole call duration
+            task_timestamp = self._ctx.task_executor.get_task_submitted_timestamp(task_id)
+            if task_timestamp is not None and time.perf_counter() - task_timestamp > self._call_deadline:
+                self._ctx.task_executor.cancel(task_id)
+                raise ErrorInfo.for_reason(
+                    ErrorCode.TIMEOUT, f"GetFlightInfo timed out while waiting for task {task_id}."
+                ).to_timeout_error()
+
+            # if the result is not ready, and we still have time, indicate to the client
+            # how to poll for the results
             raise _prepare_poll_error(task_id)
         except Exception:
             _LOGGER.error("get_flight_info_failed", task_id=task_id, fun=fun_name, exc_info=True)

gooddata-flexconnect/tests/server/conftest.py

@@ -80,7 +80,7 @@ def flexconnect_server(
     funs = f"[{funs}]"
 
     os.environ["GOODDATA_FLIGHT_FLEXCONNECT__FUNCTIONS"] = funs
-    os.environ["GOODDATA_FLIGHT_FLEXCONNECT__CALL_DEADLINE_MS"] = "1000"
+    os.environ["GOODDATA_FLIGHT_FLEXCONNECT__CALL_DEADLINE_MS"] = "1200"
     os.environ["GOODDATA_FLIGHT_FLEXCONNECT__POLLING_INTERVAL_MS"] = "500"
 
     with server(create_flexconnect_flight_methods, tls, mtls) as s:

gooddata-flexconnect/tests/server/funs/fun3.py

@@ -27,7 +27,7 @@ def call(
     ) -> ArrowData:
         # sleep is intentionally setup to be longer than the deadline for
         # the function invocation (see conftest.py // flexconnect_server fixture)
-        time.sleep(1.5)
+        time.sleep(2)
 
         return pyarrow.table(
             data={

gooddata-flexconnect/tests/server/test_flexconnect_server.py

@@ -10,6 +10,10 @@
 
 
 def test_basic_function():
+    """
+    This function should return immediately when called, no polling necessary.
+    :return:
+    """
     with flexconnect_server(["tests.server.funs.fun1"]) as s:
         c = pyarrow.flight.FlightClient(s.location)
         fun_infos = list(c.list_flights())
@@ -126,6 +130,13 @@ def test_function_with_polling():
         assert len(data) == 3
         assert data.column_names == ["col1", "col2", "col3"]
 
+        # also check that trying to cancel already completed task results in cancelled with correct code
+        with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
+            c.get_flight_info(retry_info.cancel_descriptor)
+
+        assert e.value is not None
+        assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)
+
 
 def test_function_with_call_deadline():
     """
@@ -150,6 +161,7 @@ def test_function_with_call_deadline():
             )
         )
 
+        # the initial submit returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
             c.get_flight_info(descriptor)
 
@@ -159,24 +171,21 @@ def test_function_with_call_deadline():
         error_info = ErrorInfo.from_bytes(e.value.extra_info)
         retry_info = RetryInfo.from_bytes(error_info.body)
 
-        # poll twice to reach the call deadline
+        # the next poll still returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
             c.get_flight_info(retry_info.retry_descriptor)
 
         assert e.value is not None
         assert_error_code(ErrorCode.POLL, e.value)
 
-        error_info = ErrorInfo.from_bytes(e.value.extra_info)
-        retry_info = RetryInfo.from_bytes(error_info.body)
-
+        # the third one reaches the deadline so the Timeout code is returned instead
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
             c.get_flight_info(retry_info.retry_descriptor)
 
-        # and then ensure the timeout error is returned
         assert_error_code(ErrorCode.TIMEOUT, e.value)
 
 
-def test_function_with_cancelation():
+def test_function_with_cancellation():
     """
     Run a long-running function and cancel it after one poll iteration.
     """
@@ -191,6 +200,7 @@ def test_function_with_cancelation():
             )
         )
 
+        # the initial submit returns polling info
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
             c.get_flight_info(descriptor)
 
@@ -200,7 +210,16 @@ def test_function_with_cancelation():
         error_info = ErrorInfo.from_bytes(e.value.extra_info)
         retry_info = RetryInfo.from_bytes(error_info.body)
 
+        # use the poll info to cancel the task
+        with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
+            c.get_flight_info(retry_info.cancel_descriptor)
+
+        assert e.value is not None
+        assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)
+
+        # even multiple cancellations return the same error
         with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
             c.get_flight_info(retry_info.cancel_descriptor)
 
         assert e.value is not None
+        assert_error_code(ErrorCode.COMMAND_CANCELLED, e.value)

gooddata-flight-server/gooddata_flight_server/tasks/thread_task_executor.py

@@ -605,8 +605,8 @@ def cancel(self, task_id: str) -> bool:
             return True
 
         if execution is None:
-            # the task was not and is not running - cancel not possible
-            return False
+            # the task was not and is not running - cancel not necessary
+            return True
 
         return execution.cancel()