gooddata-flight-server 1.34.1.dev1__py3-none-any.whl

gooddata_flight_server/tasks/metrics.py

@@ -0,0 +1,115 @@
+#  (C) 2024 GoodData Corporation
+import threading
+from typing import Callable, TypeVar
+
+from prometheus_client import Counter, Gauge, Summary
+from prometheus_client.metrics import MetricWrapperBase
+
+_TMetric = TypeVar("_TMetric", bound=MetricWrapperBase)
+
+
+class TaskExecutorMetrics:
+    """
+    Facade to access prometheus metrics that the task executor maintains.
+
+    Note that this is somewhat more convoluted because:
+
+    1. The TaskExecutor can produce metrics for various task types so metric names have to
+       be variable (based on prefix)
+
+    2. Prometheus does not like double-registration of metrics - this is something that
+       definitely happens during various tests.
+
+    Thus, for each metric, the class maintains a static mapping (prefix -> actual instance) and
+    every time the class is instantiated with particular prefix, the constructor will get existing
+    or create new instances.
+    """
+
+    _QueueSize: dict[str, Gauge] = {}
+    _CloseQueueSize: dict[str, Gauge] = {}
+    _WaitTime: dict[str, Summary] = {}
+    _TaskE2EDuration: dict[str, Summary] = {}
+    _TaskDuration: dict[str, Summary] = {}
+    _TaskErrors: dict[str, Counter] = {}
+    _TaskCancelled: dict[str, Counter] = {}
+    _TaskCompleted: dict[str, Counter] = {}
+    _MapLock = threading.Lock()
+
+    @staticmethod
+    def _get_or_create(d: dict[str, _TMetric], prefix: str, create_fun: Callable[[], _TMetric]) -> _TMetric:
+        with TaskExecutorMetrics._MapLock:
+            existing = d.get(prefix)
+            if existing is not None:
+                return existing
+
+            new = create_fun()
+            d[prefix] = new
+            return new
+
+    def __init__(self, prefix: str) -> None:
+        self._prefix = prefix
+
+        self.queue_size = self._get_or_create(
+            TaskExecutorMetrics._QueueSize,
+            prefix,
+            lambda: Gauge(f"{prefix}_task_queue", "Number of tasks waiting in queue."),
+        )
+
+        self.close_queue_size = self._get_or_create(
+            TaskExecutorMetrics._CloseQueueSize,
+            prefix,
+            lambda: Gauge(
+                f"{prefix}_close_queue",
+                "Number of task execution results waiting in the queue to be closed and cleaned up.",
+            ),
+        )
+
+        self.wait_time = self._get_or_create(
+            TaskExecutorMetrics._WaitTime,
+            prefix,
+            lambda: Summary(
+                f"{prefix}_task_wait",
+                "Time a task spends waiting in queue before it is executed.",
+            ),
+        )
+
+        self.task_duration = self._get_or_create(
+            TaskExecutorMetrics._TaskDuration,
+            prefix,
+            lambda: Summary(
+                f"{prefix}_task_duration",
+                "Duration of task run itself (does not include wait or prerequisite resolution duration).",
+            ),
+        )
+
+        self.task_e2e_duration = self._get_or_create(
+            TaskExecutorMetrics._TaskE2EDuration,
+            prefix,
+            lambda: Summary(
+                f"{prefix}_task_e2e_duration",
+                "End-to-end duration of the task execution. Includes prerequisite resolution duration and "
+                "time spent in queue. This is the duration as observed by the callers.",
+            ),
+        )
+
+        self.task_errors = self._get_or_create(
+            TaskExecutorMetrics._TaskErrors,
+            prefix,
+            lambda: Counter(f"{prefix}_task_error", "Number of failed tasks."),
+        )
+
+        self.task_cancelled = self._get_or_create(
+            TaskExecutorMetrics._TaskCancelled,
+            prefix,
+            lambda: Counter(f"{prefix}_task_cancelled", "Number of cancelled tasks."),
+        )
+
+        self.task_completed = self._get_or_create(
+            TaskExecutorMetrics._TaskCompleted,
+            prefix,
+            lambda: Counter(
+                f"{prefix}_task_completed",
+                "Number of completed tasks - this includes all tasks regardless "
+                "of how their execution completed (success, failure, cancel).",
+            ),
+        )

gooddata_flight_server/tasks/task.py

@@ -0,0 +1,193 @@
+#  (C) 2024 GoodData Corporation
+import abc
+import threading
+import uuid
+from concurrent.futures import CancelledError
+from typing import Optional, Union, final
+
+from gooddata_flight_server.tasks.task_error import TaskError
+from gooddata_flight_server.tasks.task_result import TaskResult
+
+
+class Task(abc.ABC):
+    """
+    Abstract base class for executable tasks.
+
+    This class provides the essential boilerplate and declares a single `run` method which
+    should be implemented by subclasses. The `run` does the actual work and returns its
+    result.
+
+    The Task design is such that it allows for runtime cancel-ability:
+
+    - task can be flagged as cancellable or not (default is True)
+    - when cancellable, the `cancel` method can be used to indicate that the task
+      should cancel
+    - this turns on the cancelled indicator
+
+    A cancellable task should test the cancelled indicator after each significant
+    step and bail-out by raising CancelledError
+
+    If the `run` method is entering a point of no return (e.g. cancel / rollback is
+    no longer feasible), then it must first switch the task to be non-cancellable
+    using the `switch_non_cancellable` - this may raise CancelledError if the `run`
+    was raced and someone cancelled the task.
+    """
+
+    __slots__ = (
+        "_task_id",
+        "_cmd",
+        "_cancel_lock",
+        "_cancelled",
+        "_cancellable",
+        "_triggers",
+    )
+
+    def __init__(
+        self,
+        cmd: bytes,
+        cancellable: bool = True,
+        task_id: Optional[str] = None,
+    ):
+        self._task_id = task_id or uuid.uuid4().hex
+        self._cmd = cmd
+        self._cancel_lock = threading.Lock()
+        self._cancelled = False
+        self._cancellable = cancellable
+
+    @final
+    @property
+    def task_id(self) -> str:
+        return self._task_id
+
+    @final
+    @property
+    def cmd(self) -> bytes:
+        return self._cmd
+
+    @final
+    @property
+    def cancelled(self) -> bool:
+        """
+        :return: true if the running task was cancelled
+        """
+        with self._cancel_lock:
+            return self._cancelled
+
+    def check_cancelled(self) -> None:
+        """
+        Checks whether task got cancelled - if so, raises CancelledError.
+
+        This is utility method that may be used in Task.run() to perform
+        cancellation checks.
+
+        :return: nothing
+        """
+        if self.cancelled:
+            raise CancelledError()
+
+    @final
+    def cancel(self) -> bool:
+        """
+        Try to cancel an *already running* task. Depending on the state of the task,
+        this may or may not be possible.
+
+        If the cancel succeeds, it is guaranteed that the task has no side-effects on
+        the rest of the system - it is as if it never run.
+
+        :return: True if cancel was successful, False if not
+        """
+        with self._cancel_lock:
+            if not self._cancellable:
+                return False
+
+            first_cancel = not self._cancelled
+            self._cancelled = True
+
+            if first_cancel:
+                try:
+                    self.on_task_cancel()
+                except Exception:
+                    pass
+
+            return True
+
+    @final
+    def switch_non_cancellable(self) -> None:
+        """
+        Switch the task to non-cancellable state.
+
+        If the task got cancelled, raises CancelledError() at this point.
+        Otherwise, sets the non-cancellable flag and returns.
+
+        :return: nothing
+        :raises: CancelledError if the switch is not possible because the task got cancelled already
+        """
+        with self._cancel_lock:
+            if self._cancelled:
+                raise CancelledError()
+
+            self._cancellable = False
+
+    def on_task_cancel(self) -> None:
+        """
+        This method will be called when a task is cancelled. That is, when it is still in
+        cancellable state and someone calls the cancel() for the first time.
+
+        The concrete implementation may optionally override this method to do something
+        special on cancellation - like cascading the cancellation to further sub-components.
+
+        Important: this method should not block.
+
+        :return: nothing
+        """
+        return
+
+    def on_task_error(self, error: TaskError) -> Optional[TaskError]:
+        """
+        This method will be called when a task fails with and raises an exception. It
+        will be called after executor creates an instance of TaskError from the
+        exception, and BEFORE it performs logging / tracking of the exception.
+
+        The concrete implementation may optionally override this method to do something
+        with the TaskError that was created by the executor. For example:
+
+        - intercept automatically generated TaskError and replace / modify it
+          (for example categorize client errors)
+
+        - do custom logging / tracking of the error
+
+        For convenience, if this method returns None, the executor will use the original
+        task error instance as-is.
+
+        :param error: TaskError as categorized by the executor
+        :return: None if the original `error` should be used, otherwise an instance of
+         TaskError
+        """
+        return error
+
+    @abc.abstractmethod
+    def run(self) -> Union[TaskResult, TaskError]:
+        """
+        Runs the task.
+
+        This method should be implemented by subclasses and do work according to payload
+        included in the `cmd`. Upon successful completion, the method should return a
+        TaskResult - either FlightPathTaskResult (when task produced a flight path) or
+        FlightDataTaskResult (when task created a live result).
+
+        Upon failure, the task has two options - use whichever is more convenient:
+
+        - either raise an exception: in this case the TaskExecutor will analyze and
+          convert the exception to TaskError (with error codes and everything) using
+          the built-in logic; the Task's `on_task_error` method will be called with
+          the TaskError created using the standard error handling logic
+
+        - return TaskError: this will be used by the TaskExecutor as-is. This option
+          is useful in situations when the task wants to do more elaborate
+          error handling / logging / reporting.
+
+        :return: result of the task
+        :raise Exception
+        :raise CancelledError: when the task's run was cancelled
+        """
+        raise NotImplementedError

gooddata_flight_server/tasks/task_error.py

@@ -0,0 +1,60 @@
+#  (C) 2024 GoodData Corporation
+import dataclasses
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import pyarrow.flight
+
+from gooddata_flight_server.errors.error_info import ErrorInfo
+
+
+@dataclass(frozen=True)
+class TaskError:
+    """
+    Detail about failed task execution. The original Exception that was raised and
+    failed the task is intentionally _not_ stored here.
+
+    That is because by storing the exception and the included traceback, the code
+    would also hold onto stack frames and all local variables bound to them - which
+    can in return hold onto _a lot_ of memory.
+
+    See: https://cosmicpercolator.com/2016/01/13/exception-leaks-in-python-2-and-3/
+    See also: https://github.com/apache/arrow/issues/36540
+
+    Also note, that clearing the traceback as hinted in above article helps somewhat
+    but is not ideal when working with FlightErrors. While testing and measuring, I
+    have found that even a freshly constructed FlightError (for example a freshly constructed
+    copy of the original exception's message + extra_info) has some non-trivial overhead.
+    Unsure why is that, and I'm not going to spend more time to investigate :)
+
+    Therefore, I have converged to this approach where the task error contains a
+    ErrorInfo (all essential detail) and a factory function to create the
+    actual FlightError.
+
+    The code that is supposed to raise the actual exception to the client will
+    instantiate the exception when needed.
+    """
+
+    error_info: ErrorInfo
+    error_factory: Callable[[str, Optional[bytes]], pyarrow.flight.FlightError]
+
+    client_error: bool = False
+    """
+    indicates whether the task failed because client provided invalid input.
+
+    this will be used purely for logging / tracking purposes. e.g. tasks failed due to client
+    providing bad input are logged as info and the task executor does not bump error counter
+    metrics
+    """
+
+    def as_flight_error(self) -> pyarrow.flight.FlightError:
+        """
+        :return: new instance of FlightError that should be raised
+        """
+        return self.error_info.to_flight_error(self.error_factory)
+
+    def to_client_error(self) -> "TaskError":
+        """
+        :return: creates a copy of this instance with `client_error` indicator set to True
+        """
+        return dataclasses.replace(self, client_error=True)

gooddata_flight_server/tasks/task_executor.py

@@ -0,0 +1,96 @@
+#  (C) 2024 GoodData Corporation
+import abc
+from typing import Optional
+
+from gooddata_flight_server.tasks.task import Task
+from gooddata_flight_server.tasks.task_result import TaskExecutionResult
+
+
+class TaskAttributes:
+    TaskId = "gooddata_flight_server.task_id"
+    TaskCancelled = "gooddata_flight_server.task_cancelled"
+    TaskError = "gooddata_flight_server.task_error"
+    TaskErrorCode = "gooddata_flight_server.task_error.code"
+    TaskErrorMsg = "gooddata_flight_server.task_error.msg"
+    TaskErrorDetail = "gooddata_flight_server.task_error.detail"
+
+
+class TaskExecutor(abc.ABC):
+    """
+    Declares interface for Task Executors. These allow asynchronous execution
+    of tasks which 'somehow' generate flight data.
+
+    The methods on this interface are designed to support a pollable
+    GetFlightInfo -> DoGet flows. A task can be submitted, polled for
+    completion or cancelled.
+
+    Once the task finishes (with any outcome), a TaskExecutionResult is available
+    and describes the outcome. On success, the execution result contains
+    a reference to task's actual result.
+
+    The execution result and the task's actual result are retained in the
+    executor for a limited (configurable) amount of time.
+
+    The actual task result is either:
+
+    - FlightDataTaskResult - which represents data that was 'generated' by the task
+      somehow and is available for single or repeated reads
+    - FlightPathTaskResult - which represent data stored under some flight path;
+      typically, flight commands include option to sink result under a flight path and
+      if that is the case the task will generate data and store it and then return
+      the pointer
+    """
+
+    @abc.abstractmethod
+    def submit(
+        self,
+        task: Task,
+    ) -> None:
+        """
+        Submit a new task that will perform all work as described in the provided command.
+
+        :param task: task to run
+        :return: nothing, task is always submitted
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def wait_for_result(self, task_id: str, timeout: Optional[float] = None) -> Optional[TaskExecutionResult]:
+        """
+        Wait for the task with the provided task id to finish.
+
+        If a task already finished and this service still has records describing it's result, then
+        the saved result is returned immediately.
+
+        Otherwise, if a task is pending, the method will block (optionally with timeout) until
+        the task completes. If it does not complete in given timeframe, the code will raise TimeoutError.
+
+        Note: if the task was cancelled, the result will have 'cancelled' indicator flag set to True.
+
+        :param task_id: task id to wait for
+        :param timeout: time to wait for completion
+        :raise TaskWaitTimeoutError: if the wait for task completion timed out
+        :return: result or None if there is no such task
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def cancel(self, task_id: str) -> bool:
+        """
+        Try to cancel a task - either by dropping it from the task queue or by cancelling
+        a running task or dropping a result of already finished task.
+
+        :param task_id: task id to cancel
+        :return: true if cancelled, false if cancel not possible (no such task or task not cancellable anymore)
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close_result(self, task_id: str) -> bool:
+        """
+        Try to close result of a previously completed task.
+
+        :param task_id: task id, whose result to close
+        :return: true if result closed, false if no result for that task id
+        """
+        raise NotImplementedError