feat: add support for polling in FlexConnect server

Add a polling mechanism to the FlexConnect functions so that
long-running tasks can be polled for and canceled.

The TaskExecutor now supports returning a timestamp of when
a particular task was submitted. This is to keep track of the
call deadline breaches.

JIRA: CQ-1124
risk: low

Author: no23reason

gooddata-flexconnect/gooddata_flexconnect/function/flight_methods.py

@@ -1,4 +1,5 @@
 #  (C) 2024 GoodData Corporation
+import time
 from collections.abc import Generator
 from typing import Optional
 
@@ -12,7 +13,6 @@
     FlightServerMethods,
     ServerContext,
     TaskExecutionResult,
-    TaskWaitTimeoutError,
     flight_server_methods,
 )
 
@@ -23,11 +23,26 @@
 _LOGGER = structlog.get_logger("gooddata_flexconnect.rpc")
 
 
+def _prepare_poll_error(task_id: str) -> pyarrow.flight.FlightError:
+    return ErrorInfo.poll(
+        flight_info=None,
+        cancel_descriptor=pyarrow.flight.FlightDescriptor.for_command(f"c:{task_id}".encode()),
+        retry_descriptor=pyarrow.flight.FlightDescriptor.for_command(f"r:{task_id}".encode()),
+    )
+
+
 class _FlexConnectServerMethods(FlightServerMethods):
-    def __init__(self, ctx: ServerContext, registry: FlexConnectFunctionRegistry, call_deadline_ms: float) -> None:
+    def __init__(
+        self,
+        ctx: ServerContext,
+        registry: FlexConnectFunctionRegistry,
+        call_deadline_ms: float,
+        poll_interval_ms: float,
+    ) -> None:
         self._ctx = ctx
         self._registry = registry
         self._call_deadline = call_deadline_ms / 1000
+        self._poll_interval = poll_interval_ms / 1000
 
     @staticmethod
     def _create_descriptor(fun_name: str, metadata: Optional[dict]) -> pyarrow.flight.FlightDescriptor:
@@ -140,39 +155,53 @@ def get_flight_info(
         descriptor: pyarrow.flight.FlightDescriptor,
     ) -> pyarrow.flight.FlightInfo:
         structlog.contextvars.bind_contextvars(peer=context.peer())
-        task: Optional[FlexConnectFunctionTask] = None
 
-        try:
-            task = self._prepare_task(context, descriptor)
-            self._ctx.task_executor.submit(task)
+        # first, check if the descriptor is a cancel descriptor
+        if descriptor.command is None or not len(descriptor.command):
+            raise ErrorInfo.bad_argument(
+                "Incorrect FlexConnect function invocation. Flight descriptor must contain command "
+                "with the invocation payload."
+            )
 
-            try:
-                # XXX: this should be enhanced to implement polling
-                task_result = self._ctx.task_executor.wait_for_result(task.task_id, self._call_deadline)
-            except TaskWaitTimeoutError:
-                cancelled = self._ctx.task_executor.cancel(task.task_id)
-                _LOGGER.warning(
-                    "flexconnect_fun_call_timeout", task_id=task.task_id, fun=task.fun_name, cancelled=cancelled
-                )
+        task_id: str
+        fun_name: Optional[str] = None
 
+        if descriptor.command.startswith(b"c:"):
+            # cancel descriptor: just cancel the given task and raise cancellation exception
+            task_id = descriptor.command[2:].decode()
+            self._ctx.task_executor.cancel(task_id)
+            raise ErrorInfo.for_reason(
+                ErrorCode.COMMAND_CANCELLED, "FlexConnect function invocation was cancelled."
+            ).to_cancelled_error()
+        elif descriptor.command.startswith(b"r:"):
+            # retry descriptor: extract the task_id, do not submit it again and do one polling iteration
+            task_id = descriptor.command[2:].decode()
+            # 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.task_id}."
+                    ErrorCode.TIMEOUT, f"GetFlightInfo timed out while waiting for task {task_id}."
                 ).to_timeout_error()
+        else:
+            # 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
+            try:
+                task = self._prepare_task(context, descriptor)
+                self._ctx.task_executor.submit(task)
+                task_id = task.task_id
+                fun_name = task.fun_name
+            except Exception:
+                _LOGGER.error("flexconnect_fun_submit_failed", exc_info=True)
+                raise
 
-            # if this bombs then there must be something really wrong because the task
-            # was clearly submitted and code was waiting for its completion. this invariant
-            # should not happen in this particular code path. The None return value may
-            # be applicable one day when polling is in use and a request comes to check whether
-            # particular task id finished
-            assert task_result is not None
-
+        try:
+            task_result = self._ctx.task_executor.wait_for_result(task_id, timeout=self._poll_interval)
             return self._prepare_flight_info(task_result)
+        except TimeoutError:
+            raise _prepare_poll_error(task_id)
         except Exception:
-            if task is not None:
-                _LOGGER.error("get_flight_info_failed", task_id=task.task_id, fun=task.fun_name, exc_info=True)
-            else:
-                _LOGGER.error("flexconnect_fun_submit_failed", exc_info=True)
-
+            _LOGGER.error("get_flight_info_failed", task_id=task_id, fun=fun_name, exc_info=True)
             raise
 
     def do_get(
@@ -201,7 +230,9 @@ def do_get(
 _FLEX_CONNECT_CONFIG_SECTION = "flexconnect"
 _FLEX_CONNECT_FUNCTION_LIST = "functions"
 _FLEX_CONNECT_CALL_DEADLINE_MS = "call_deadline_ms"
+_FLEX_CONNECT_POLLING_INTERVAL_MS = "polling_interval_ms"
 _DEFAULT_FLEX_CONNECT_CALL_DEADLINE_MS = 180_000
+_DEFAULT_FLEX_CONNECT_POLLING_INTERVAL_MS = 2000
 
 
 def _read_call_deadline_ms(ctx: ServerContext) -> int:
@@ -223,6 +254,24 @@ def _read_call_deadline_ms(ctx: ServerContext) -> int:
         )
 
 
+def _read_polling_interval_ms(ctx: ServerContext) -> int:
+    polling_interval = ctx.settings.get(f"{_FLEX_CONNECT_CONFIG_SECTION}.{_FLEX_CONNECT_POLLING_INTERVAL_MS}")
+    if polling_interval is None:
+        return _DEFAULT_FLEX_CONNECT_POLLING_INTERVAL_MS
+
+    try:
+        polling_interval = int(polling_interval)
+        if polling_interval <= 0:
+            raise ValueError()
+        return polling_interval
+    except ValueError:
+        raise ValueError(
+            f"Value of {_FLEX_CONNECT_CONFIG_SECTION}.{_FLEX_CONNECT_POLLING_INTERVAL_MS} must "
+            f"be a positive number - duration, in milliseconds, that FlexConnect function "
+            f"waits for the result during one polling iteration."
+        )
+
+
 @flight_server_methods
 def create_flexconnect_flight_methods(ctx: ServerContext) -> FlightServerMethods:
     """
@@ -236,8 +285,9 @@ def create_flexconnect_flight_methods(ctx: ServerContext) -> FlightServerMethods
     """
     modules = list(ctx.settings.get(f"{_FLEX_CONNECT_CONFIG_SECTION}.{_FLEX_CONNECT_FUNCTION_LIST}") or [])
     call_deadline_ms = _read_call_deadline_ms(ctx)
+    polling_interval_ms = _read_polling_interval_ms(ctx)
 
     _LOGGER.info("flexconnect_init", modules=modules)
     registry = FlexConnectFunctionRegistry().load(ctx, modules)
 
-    return _FlexConnectServerMethods(ctx, registry, call_deadline_ms)
+    return _FlexConnectServerMethods(ctx, registry, call_deadline_ms, polling_interval_ms)

gooddata-flexconnect/tests/server/conftest.py

@@ -80,7 +80,8 @@ def flexconnect_server(
     funs = f"[{funs}]"
 
     os.environ["GOODDATA_FLIGHT_FLEXCONNECT__FUNCTIONS"] = funs
-    os.environ["GOODDATA_FLIGHT_FLEXCONNECT__CALL_DEADLINE_MS"] = "500"
+    os.environ["GOODDATA_FLIGHT_FLEXCONNECT__CALL_DEADLINE_MS"] = "1000"
+    os.environ["GOODDATA_FLIGHT_FLEXCONNECT__POLLING_INTERVAL_MS"] = "500"
 
     with server(create_flexconnect_flight_methods, tls, mtls) as s:
         yield 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)
+        time.sleep(1.5)
 
         return pyarrow.table(
             data={

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

@@ -0,0 +1,39 @@
+#  (C) 2024 GoodData Corporation
+import time
+from typing import Optional
+
+import pyarrow
+from gooddata_flexconnect.function.function import FlexConnectFunction
+from gooddata_flight_server import ArrowData
+
+_DATA: Optional[pyarrow.Table] = None
+
+
+class _PollableFun(FlexConnectFunction):
+    Name = "PollableFun"
+    Schema = pyarrow.schema(
+        fields=[
+            pyarrow.field("col1", pyarrow.int64()),
+            pyarrow.field("col2", pyarrow.string()),
+            pyarrow.field("col3", pyarrow.bool_()),
+        ]
+    )
+
+    def call(
+        self,
+        parameters: dict,
+        columns: tuple[str, ...],
+        headers: dict[str, list[str]],
+    ) -> ArrowData:
+        # sleep is intentionally setup to be longer than one polling interval
+        # (see conftest.py // flexconnect_server fixture)
+        time.sleep(0.7)
+
+        return pyarrow.table(
+            data={
+                "col1": [1, 2, 3],
+                "col2": ["a", "b", "c"],
+                "col3": [True, False, True],
+            },
+            schema=self.Schema,
+        )

gooddata-flexconnect/tests/server/test_flexconnect_server.py

@@ -1,8 +1,9 @@
 #  (C) 2024 GoodData Corporation
+
 import orjson
 import pyarrow.flight
 import pytest
-from gooddata_flight_server import ErrorCode
+from gooddata_flight_server import ErrorCode, ErrorInfo, RetryInfo
 
 from tests.assert_error_info import assert_error_code
 from tests.server.conftest import flexconnect_server
@@ -89,6 +90,43 @@ def test_basic_function_tls(tls_ca_cert):
         assert data.column_names == ["col1", "col2", "col3"]
 
 
+def test_function_with_polling():
+    """
+    Flight RPC implementation that invokes FlexConnect can return a polling info.
+
+    This way, the client can poll for results that take longer to complete.
+    """
+    with flexconnect_server(["tests.server.funs.fun4"]) as s:
+        c = pyarrow.flight.FlightClient(s.location)
+        descriptor = pyarrow.flight.FlightDescriptor.for_command(
+            orjson.dumps(
+                {
+                    "functionName": "PollableFun",
+                    "parameters": {"test1": 1, "test2": 2, "test3": 3},
+                }
+            )
+        )
+
+        # the function is set to sleep a bit longer than the polling interval,
+        # so the first iteration returns retry info in the exception
+        with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
+            c.get_flight_info(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)
+
+        # use the retry info to poll again for the result,
+        # now it should be ready and returned normally
+        info = c.get_flight_info(retry_info.retry_descriptor)
+        data: pyarrow.Table = c.do_get(info.endpoints[0].ticket).read_all()
+
+        assert len(data) == 3
+        assert data.column_names == ["col1", "col2", "col3"]
+
+
 def test_function_with_call_deadline():
     """
     Flight RPC implementation that invokes FlexConnect can be setup with
@@ -115,4 +153,54 @@ def test_function_with_call_deadline():
         with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
             c.get_flight_info(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)
+
+        # poll twice to reach the call deadline
+        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)
+
+        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():
+    """
+    Run a long-running function and cancel it after one poll iteration.
+    """
+    with flexconnect_server(["tests.server.funs.fun3"]) as s:
+        c = pyarrow.flight.FlightClient(s.location)
+        descriptor = pyarrow.flight.FlightDescriptor.for_command(
+            orjson.dumps(
+                {
+                    "functionName": "LongRunningFun",
+                    "parameters": {"test1": 1, "test2": 2, "test3": 3},
+                }
+            )
+        )
+
+        with pytest.raises(pyarrow.flight.FlightTimedOutError) as e:
+            c.get_flight_info(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)
+
+        with pytest.raises(pyarrow.flight.FlightCancelledError) as e:
+            c.get_flight_info(retry_info.cancel_descriptor)
+
+        assert e.value is not None

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

@@ -54,6 +54,15 @@ def submit(
         """
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def get_task_submitted_timestamp(self, task_id: str) -> Optional[float]:
+        """
+        Returns the timestamp of when the task with the given id was submitted.
+        :param task_id: task id to get the timestamp for
+        :return: Timestamp in seconds since epoch of when the task was submitted or None if there is no such task
+        """
+        raise NotImplementedError
+
     @abc.abstractmethod
     def wait_for_result(self, task_id: str, timeout: Optional[float] = None) -> Optional[TaskExecutionResult]:
         """

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

@@ -565,6 +565,14 @@ def submit(
         execution.start()
         self._metrics.queue_size.set(self._queue_size)
 
+    def get_task_submitted_timestamp(self, task_id: str) -> Optional[float]:
+        with self._task_lock:
+            execution = self._executions.get(task_id)
+
+        if execution is not None:
+            return execution.stats.created
+        return None
+
     def wait_for_result(self, task_id: str, timeout: Optional[float] = None) -> Optional[TaskExecutionResult]:
         with self._task_lock:
             execution = self._executions.get(task_id)