flagsmith-common 2.2.4__py3-none-any.whl

task_processor/migrations/sql/0013_get_recurringtasks_to_process.sql

@@ -0,0 +1,33 @@
+CREATE OR REPLACE FUNCTION get_recurringtasks_to_process()
+RETURNS SETOF task_processor_recurringtask AS $$
+DECLARE
+    row_to_return task_processor_recurringtask;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_recurringtask
+        -- Add one minute to the timeout as a grace period for overhead
+        WHERE is_locked = FALSE OR (locked_at IS NOT NULL AND locked_at < NOW() - timeout + INTERVAL '1 minute')
+        ORDER BY last_picked_at NULLS FIRST
+        LIMIT 1
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_recurringtask
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE, locked_at = NOW(), last_picked_at = NOW()
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the columns here, the client will receive a row
+        -- that is locked but still shows `is_locked` as `False` and `locked_at` as `None`.
+        row_to_return.is_locked := TRUE;
+        row_to_return.locked_at := NOW();
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

task_processor/models.py

@@ -0,0 +1,237 @@
+import typing
+import uuid
+from datetime import datetime, timedelta
+
+import simplejson as json
+from django.core.serializers.json import DjangoJSONEncoder
+from django.db import models
+from django.utils import timezone
+
+from task_processor.exceptions import TaskQueueFullError
+from task_processor.managers import RecurringTaskManager, TaskManager
+from task_processor.task_registry import get_task, registered_tasks
+from task_processor.types import TaskCallable
+
+_django_json_encoder_default = DjangoJSONEncoder().default
+
+
+class TaskPriority(models.IntegerChoices):
+    LOWER = 100
+    LOW = 75
+    NORMAL = 50
+    HIGH = 25
+    HIGHEST = 0
+
+
+class AbstractBaseTask(models.Model):
+    uuid = models.UUIDField(unique=True, default=uuid.uuid4)
+    created_at = models.DateTimeField(auto_now_add=True)
+    task_identifier = models.CharField(max_length=200)
+    serialized_args = models.TextField(blank=True, null=True)
+    serialized_kwargs = models.TextField(blank=True, null=True)
+    is_locked = models.BooleanField(default=False)
+    timeout = models.DurationField(blank=True, null=True)
+
+    class Meta:
+        abstract = True
+
+    @property
+    def args(self) -> tuple[typing.Any, ...]:
+        if self.serialized_args:
+            args = self.deserialize_data(self.serialized_args)
+            return tuple(args)
+        return ()
+
+    @property
+    def kwargs(self) -> typing.Dict[str, typing.Any]:
+        if self.serialized_kwargs:
+            kwargs = self.deserialize_data(self.serialized_kwargs)
+            if typing.TYPE_CHECKING:
+                assert isinstance(kwargs, dict)
+            return kwargs
+        return {}
+
+    @staticmethod
+    def serialize_data(data: typing.Any) -> str:
+        return json.dumps(data, default=_django_json_encoder_default)
+
+    @staticmethod
+    def deserialize_data(data: str) -> typing.Any:
+        return json.loads(data)
+
+    def mark_failure(self) -> None:
+        self.unlock()
+
+    def mark_success(self) -> None:
+        self.unlock()
+
+    def unlock(self) -> None:
+        self.is_locked = False
+
+    def run(self) -> None:
+        return self.callable(*self.args, **self.kwargs)
+
+    @property
+    def callable(self) -> TaskCallable[typing.Any]:
+        task = get_task(self.task_identifier)
+        return task.task_function
+
+
+class Task(AbstractBaseTask):
+    scheduled_for = models.DateTimeField(blank=True, null=True, default=timezone.now)
+
+    timeout = models.DurationField(blank=True, null=True)
+
+    # denormalise failures and completion so that we can use select_for_update
+    num_failures = models.IntegerField(default=0)
+    completed = models.BooleanField(default=False)
+    objects: TaskManager = TaskManager()
+    priority = models.SmallIntegerField(
+        default=None, null=True, choices=TaskPriority.choices
+    )
+
+    class Meta:
+        # We have customised the migration in 0004 to only apply this change to postgres databases
+        # TODO: work out how to index the taskprocessor_task table for Oracle and MySQL
+        indexes = [
+            models.Index(
+                name="incomplete_tasks_idx",
+                fields=["scheduled_for"],
+                condition=models.Q(completed=False, num_failures__lt=3),
+            )
+        ]
+
+    @classmethod
+    def create(
+        cls,
+        task_identifier: str,
+        scheduled_for: datetime,
+        priority: TaskPriority = TaskPriority.NORMAL,
+        queue_size: int | None = None,
+        *,
+        args: typing.Tuple[typing.Any, ...] | None = None,
+        kwargs: typing.Dict[str, typing.Any] | None = None,
+        timeout: timedelta | None = timedelta(seconds=60),
+    ) -> "Task":
+        if queue_size and cls._is_queue_full(task_identifier, queue_size):
+            raise TaskQueueFullError(
+                f"Queue for task {task_identifier} is full. "
+                f"Max queue size is {queue_size}"
+            )
+        return Task(
+            task_identifier=task_identifier,
+            scheduled_for=scheduled_for,
+            priority=priority,
+            serialized_args=cls.serialize_data(args or tuple()),
+            serialized_kwargs=cls.serialize_data(kwargs or dict()),
+            timeout=timeout,
+        )
+
+    @classmethod
+    def _is_queue_full(cls, task_identifier: str, queue_size: int) -> bool:
+        return (
+            cls.objects.filter(
+                task_identifier=task_identifier,
+                completed=False,
+                num_failures__lt=3,
+            ).count()
+            > queue_size
+        )
+
+    def mark_failure(self) -> None:
+        super().mark_failure()
+        self.num_failures += 1
+
+    def mark_success(self) -> None:
+        super().mark_success()
+        self.completed = True
+
+
+class RecurringTask(AbstractBaseTask):
+    run_every = models.DurationField()
+    first_run_time = models.TimeField(blank=True, null=True)
+
+    locked_at = models.DateTimeField(blank=True, null=True)
+    timeout = models.DurationField(default=timedelta(minutes=30))
+
+    last_picked_at = models.DateTimeField(blank=True, null=True)
+    objects: RecurringTaskManager = RecurringTaskManager()
+
+    class Meta:
+        constraints = [
+            models.UniqueConstraint(
+                fields=["task_identifier", "run_every"],
+                name="unique_run_every_tasks",
+            ),
+        ]
+
+    def unlock(self) -> None:
+        self.is_locked = False
+        self.locked_at = None
+
+    @property
+    def should_execute(self) -> bool:
+        now = timezone.now()
+        last_task_run = (
+            self.task_runs.order_by("-started_at").first() if self.pk else None
+        )
+
+        if not last_task_run:
+            # If we have never run this task, then we should execute it only if
+            # the time has passed after which we want to ensure this task runs.
+            # This allows us to control when intensive tasks should be run.
+            return not (self.first_run_time and self.first_run_time > now.time())
+
+        # if the last run was at t- run_every, then we should execute it
+        if (timezone.now() - last_task_run.started_at) >= self.run_every:
+            return True
+
+        # if the last run was not a success and we do not have
+        # more than 3 failures in t- run_every, then we should execute it
+        if (
+            last_task_run.result != TaskResult.SUCCESS.name
+            and self.task_runs.filter(started_at__gte=(now - self.run_every)).count()
+            <= 3
+        ):
+            return True
+        # otherwise, we should not execute it
+        return False
+
+    @property
+    def is_task_registered(self) -> bool:
+        return self.task_identifier in registered_tasks
+
+
+class TaskResult(models.Choices):
+    SUCCESS = "SUCCESS"
+    FAILURE = "FAILURE"
+
+
+class AbstractTaskRun(models.Model):
+    started_at = models.DateTimeField()
+    finished_at = models.DateTimeField(blank=True, null=True)
+    result = models.CharField(
+        max_length=50, choices=TaskResult.choices, blank=True, null=True, db_index=True
+    )
+    error_details = models.TextField(blank=True, null=True)
+    task = models.ForeignKey(
+        AbstractBaseTask, on_delete=models.CASCADE, related_name="task_runs"
+    )
+
+    class Meta:
+        abstract = True
+
+
+class TaskRun(AbstractTaskRun):
+    task = models.ForeignKey(Task, on_delete=models.CASCADE, related_name="task_runs")
+
+
+class RecurringTaskRun(AbstractTaskRun):
+    task = models.ForeignKey(
+        RecurringTask, on_delete=models.CASCADE, related_name="task_runs"
+    )
+
+
+class HealthCheckModel(models.Model):
+    created_at = models.DateTimeField(auto_now_add=True)
+    uuid = models.UUIDField(unique=True, blank=False, null=False)

task_processor/monitoring.py

@@ -0,0 +1,12 @@
+from django.utils import timezone
+
+from task_processor.models import Task
+
+
+def get_num_waiting_tasks() -> int:
+    return Task.objects.filter(
+        num_failures__lt=3,
+        completed=False,
+        scheduled_for__lt=timezone.now(),
+        is_locked=False,
+    ).count()

task_processor/processor.py

@@ -0,0 +1,202 @@
+import logging
+import traceback
+import typing
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from datetime import timedelta
+
+from django.conf import settings
+from django.utils import timezone
+
+from task_processor import metrics
+from task_processor.exceptions import TaskBackoffError
+from task_processor.managers import RecurringTaskManager, TaskManager
+from task_processor.models import (
+    AbstractBaseTask,
+    RecurringTask,
+    RecurringTaskRun,
+    Task,
+    TaskResult,
+    TaskRun,
+)
+from task_processor.task_registry import TaskType, get_task
+
+T = typing.TypeVar("T", bound=AbstractBaseTask)
+AnyTaskRun = TaskRun | RecurringTaskRun
+
+logger = logging.getLogger(__name__)
+
+UNREGISTERED_RECURRING_TASK_GRACE_PERIOD = timedelta(minutes=30)
+
+
+def run_tasks(database: str, num_tasks: int = 1) -> list[TaskRun]:
+    if num_tasks < 1:
+        raise ValueError("Number of tasks to process must be at least one")
+
+    task_manager: TaskManager = Task.objects.db_manager(database)
+    tasks = task_manager.get_tasks_to_process(num_tasks)
+    if tasks:
+        logger.debug(f"Running {len(tasks)} task(s) from database '{database}'")
+
+        executed_tasks = []
+        task_runs = []
+
+        for task in tasks:
+            task, task_run = _run_task(task)
+
+            executed_tasks.append(task)
+            assert isinstance(task_run, TaskRun)
+            task_runs.append(task_run)
+
+        if executed_tasks:
+            Task.objects.using(database).bulk_update(
+                executed_tasks,
+                fields=["completed", "num_failures", "is_locked", "scheduled_for"],
+            )
+
+        if task_runs:
+            TaskRun.objects.using(database).bulk_create(task_runs)
+            logger.debug(
+                f"Finished running {len(task_runs)} task(s) from database '{database}'"
+            )
+
+        return task_runs
+
+    return []
+
+
+def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
+    # NOTE: We will probably see a lot of delay in the execution of recurring tasks
+    # if the tasks take longer then `run_every` to execute. This is not
+    # a problem for now, but we should be mindful of this limitation
+    task_manager: RecurringTaskManager = RecurringTask.objects.db_manager(database)
+    tasks = task_manager.get_tasks_to_process()
+    if tasks:
+        logger.debug(f"Running {len(tasks)} recurring task(s)")
+
+        task_runs = []
+
+        for task in tasks:
+            if not task.is_task_registered:
+                # This is necessary to ensure that old instances of the task processor,
+                # which may still be running during deployment, do not remove tasks added by new instances.
+                # Reference: https://github.com/Flagsmith/flagsmith/issues/2551
+                task_age = timezone.now() - task.created_at
+                if task_age > UNREGISTERED_RECURRING_TASK_GRACE_PERIOD:
+                    task.delete(using=database)
+                continue
+
+            if task.should_execute:
+                task, task_run = _run_task(task)
+                assert isinstance(task_run, RecurringTaskRun)
+                task_runs.append(task_run)
+            else:
+                task.unlock()
+
+        # update all tasks that were not deleted
+        to_update = [task for task in tasks if task.id]
+        RecurringTask.objects.using(database).bulk_update(
+            to_update,
+            fields=["is_locked", "locked_at"],
+        )
+
+        if task_runs:
+            RecurringTaskRun.objects.using(database).bulk_create(task_runs)
+            logger.debug(f"Finished running {len(task_runs)} recurring task(s)")
+
+        return task_runs
+
+    return []
+
+
+def _run_task(
+    task: T,
+) -> typing.Tuple[T, AnyTaskRun]:
+    assert settings.TASK_PROCESSOR_MODE, (
+        "Attempt to run tasks in a non-task-processor environment"
+    )
+
+    ctx = ExitStack()
+    timer = metrics.flagsmith_task_processor_task_duration_seconds.time()
+    ctx.enter_context(timer)
+
+    task_identifier = task.task_identifier
+    registered_task = get_task(task_identifier)
+
+    logger.debug(
+        f"Running task {task_identifier} id={task.pk} args={task.args} kwargs={task.kwargs}"
+    )
+    task_run: AnyTaskRun = task.task_runs.model(started_at=timezone.now(), task=task)  # type: ignore[attr-defined]
+    result: str
+    executor = None
+
+    try:
+        # Use explicit executor management to avoid blocking on shutdown
+        # when tasks timeout but continue running in worker threads.
+        # The default context manager behavior (wait=True) would block
+        # the TaskRunner thread indefinitely waiting for stuck workers.
+        executor = ThreadPoolExecutor(max_workers=1)
+        future = executor.submit(task.run)
+        timeout = task.timeout.total_seconds() if task.timeout else None
+        future.result(timeout=timeout)  # Wait for completion or timeout
+
+        task_run.result = result = TaskResult.SUCCESS.value
+        task_run.finished_at = timezone.now()
+        task.mark_success()
+
+        logger.debug(f"Task {task_identifier} id={task.pk} completed")
+
+    except Exception as e:
+        # For errors that don't include a default message (e.g., TimeoutError),
+        # fall back to using repr.
+        err_msg = str(e) or repr(e)
+
+        task.mark_failure()
+
+        task_run.result = result = TaskResult.FAILURE.value
+        task_run.error_details = str(traceback.format_exc())
+
+        logger.error(
+            "Failed to execute task '%s', with id %d. Exception: %s",
+            task_identifier,
+            task.pk,
+            err_msg,
+            exc_info=True,
+        )
+
+        if isinstance(e, TaskBackoffError):
+            assert registered_task.task_type == TaskType.STANDARD, (
+                "Attempt to back off a recurring task (currently not supported)"
+            )
+            if typing.TYPE_CHECKING:
+                assert isinstance(task, Task)
+            if task.num_failures <= 3:
+                delay_until = e.delay_until or timezone.now() + timedelta(
+                    seconds=settings.TASK_BACKOFF_DEFAULT_DELAY_SECONDS,
+                )
+                task.scheduled_for = delay_until
+                logger.info(
+                    "Backoff requested. Task '%s' set to retry at %s",
+                    task_identifier,
+                    delay_until,
+                )
+
+    finally:
+        # Always shutdown the executor without waiting for worker threads.
+        # This prevents the TaskRunner thread from blocking indefinitely
+        # when a task times out but continues running in a worker thread.
+        if executor is not None:
+            executor.shutdown(wait=False)
+
+    labels = {
+        "task_identifier": task_identifier,
+        "task_type": registered_task.task_type.value.lower(),
+        "result": result.lower(),
+    }
+
+    timer.labels(**labels)  # type: ignore[no-untyped-call]
+    ctx.close()
+
+    metrics.flagsmith_task_processor_finished_tasks_total.labels(**labels).inc()
+
+    return task, task_run

task_processor/routers.py

@@ -0,0 +1,55 @@
+from django.db.models import Model
+
+
+class TaskProcessorRouter:
+    """
+    Routing of database operations for task processor models
+    """
+
+    route_app_labels = ["task_processor"]
+
+    def db_for_read(self, model: type[Model], **hints: None) -> str | None:
+        if model._meta.app_label in self.route_app_labels:
+            return "task_processor"
+
+        return None
+
+    def db_for_write(self, model: type[Model], **hints: None) -> str | None:
+        if model._meta.app_label in self.route_app_labels:
+            return "task_processor"
+
+        return None
+
+    def allow_relation(self, obj1: Model, obj2: Model, **hints: None) -> bool | None:
+        both_objects_from_task_processor = (
+            obj1._meta.app_label in self.route_app_labels
+            and obj2._meta.app_label in self.route_app_labels
+        )
+
+        if both_objects_from_task_processor:
+            return True
+
+        return None
+
+    def allow_migrate(
+        self,
+        db: str,
+        app_label: str,
+        **hints: None,
+    ) -> bool | None:
+        """
+        Allow migrations to hit BOTH databases
+
+        NOTE: We run migrations on both databases because:
+
+        - The `task_processor` separate database was only introduced later in
+          history, and migrating to it does not delete old data from `default`.
+          We'd rather keep data in `default` consistent across time rather than
+          leaving behind possibly inconsistent data.
+        - We want to make it easier to migrate to the new database, _or back_
+          to a single database setup if needed. Running DDL consistently helps.
+        """
+        if app_label in self.route_app_labels:
+            return db in ["default", "task_processor"]
+
+        return None

task_processor/serializers.py

@@ -0,0 +1,7 @@
+from rest_framework import serializers
+
+from task_processor.types import MonitoringInfo
+
+
+class MonitoringSerializer(serializers.Serializer[MonitoringInfo]):
+    waiting = serializers.IntegerField(read_only=True)

task_processor/task_registry.py

@@ -0,0 +1,90 @@
+import enum
+import logging
+import typing
+from dataclasses import dataclass
+
+from task_processor.exceptions import TaskProcessingError
+from task_processor.types import TaskCallable
+
+logger = logging.getLogger(__name__)
+
+
+class TaskType(enum.Enum):
+    STANDARD = "STANDARD"
+    RECURRING = "RECURRING"
+
+
+@dataclass
+class RegisteredTask:
+    task_identifier: str
+    task_function: TaskCallable[typing.Any]
+    task_type: TaskType = TaskType.STANDARD
+    task_kwargs: dict[str, typing.Any] | None = None
+
+
+registered_tasks: dict[str, RegisteredTask] = {}
+
+
+def initialise() -> None:
+    global registered_tasks
+
+    from task_processor.models import RecurringTask
+
+    for task_identifier, registered_task in registered_tasks.items():
+        logger.debug("Initialising task '%s'", task_identifier)
+
+        if registered_task.task_type == TaskType.RECURRING:
+            logger.debug("Persisting recurring task '%s'", task_identifier)
+            RecurringTask.objects.update_or_create(
+                task_identifier=task_identifier,
+                defaults=registered_task.task_kwargs,
+            )
+
+
+def get_task(task_identifier: str) -> RegisteredTask:
+    global registered_tasks
+
+    try:
+        return registered_tasks[task_identifier]
+    except KeyError:
+        raise TaskProcessingError(
+            "No task registered with identifier '%s'. Ensure your task is "
+            "decorated with @register_task_handler.",
+            task_identifier,
+        )
+
+
+def register_task(
+    task_identifier: str,
+    callable_: TaskCallable[typing.Any],
+) -> None:
+    global registered_tasks
+
+    registered_task = RegisteredTask(
+        task_identifier=task_identifier,
+        task_function=callable_,
+    )
+    registered_tasks[task_identifier] = registered_task
+
+
+def register_recurring_task(
+    task_identifier: str,
+    callable_: TaskCallable[typing.Any],
+    **task_kwargs: typing.Any,
+) -> None:
+    global registered_tasks
+
+    logger.debug("Registering recurring task '%s'", task_identifier)
+
+    registered_task = RegisteredTask(
+        task_identifier=task_identifier,
+        task_function=callable_,
+        task_type=TaskType.RECURRING,
+        task_kwargs=task_kwargs,
+    )
+    registered_tasks[task_identifier] = registered_task
+
+    logger.debug(
+        "Registered tasks now has the following tasks registered: %s",
+        list(registered_tasks.keys()),
+    )

task_processor/task_run_method.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class TaskRunMethod(Enum):
+    SYNCHRONOUSLY = "SYNCHRONOUSLY"
+    SEPARATE_THREAD = "SEPARATE_THREAD"
+    TASK_PROCESSOR = "TASK_PROCESSOR"