flagsmith-common 2.2.4__py3-none-any.whl

task_processor/tasks.py

@@ -0,0 +1,71 @@
+import logging
+import typing
+from datetime import timedelta
+
+from django.conf import settings
+from django.db.models import Q
+from django.utils import timezone
+
+from task_processor.decorators import (
+    register_recurring_task,
+    register_task_handler,
+)
+from task_processor.models import HealthCheckModel, RecurringTaskRun, Task
+
+if typing.TYPE_CHECKING:
+    # ugh https://github.com/typeddjango/django-stubs/issues/1744
+    # TODO maybe switch to https://github.com/getsentry/sentry-forked-django-stubs
+    HealthCheckModel.objects = HealthCheckModel._default_manager
+    RecurringTaskRun.objects = RecurringTaskRun._default_manager
+
+
+logger = logging.getLogger(__name__)
+
+
+@register_task_handler()
+def create_health_check_model(health_check_model_uuid: str) -> None:
+    logger.info("Creating health check model.")
+    HealthCheckModel.objects.create(uuid=health_check_model_uuid)
+
+
+@register_recurring_task(
+    run_every=settings.TASK_DELETE_RUN_EVERY,
+    first_run_time=settings.TASK_DELETE_RUN_TIME,
+)
+def clean_up_old_tasks() -> None:
+    if not settings.ENABLE_CLEAN_UP_OLD_TASKS:
+        return
+
+    now = timezone.now()
+    delete_before = now - timedelta(days=settings.TASK_DELETE_RETENTION_DAYS)
+
+    # build the query
+    query = Q(completed=True)
+    if settings.TASK_DELETE_INCLUDE_FAILED_TASKS:
+        query = query | Q(num_failures__gte=3)
+    query = Q(scheduled_for__lt=delete_before) & query
+
+    # TODO: validate if deleting in batches is more / less impactful on the DB
+    while True:
+        # delete in batches of settings.TASK_DELETE_BATCH_SIZE
+        num_tasks_deleted, _ = Task.objects.filter(
+            pk__in=Task.objects.filter(query).values_list("id", flat=True)[
+                0 : settings.TASK_DELETE_BATCH_SIZE  # noqa:E203
+            ]
+        ).delete()
+        if num_tasks_deleted == 0:
+            break
+
+
+@register_recurring_task(
+    run_every=settings.TASK_DELETE_RUN_EVERY,
+    first_run_time=settings.TASK_DELETE_RUN_TIME,
+)
+def clean_up_old_recurring_task_runs() -> None:
+    if not settings.ENABLE_CLEAN_UP_OLD_TASKS:
+        return
+
+    now = timezone.now()
+    delete_before = now - timedelta(days=settings.RECURRING_TASK_RUN_RETENTION_DAYS)
+
+    RecurringTaskRun.objects.filter(finished_at__lt=delete_before).delete()

task_processor/threads.py

@@ -0,0 +1,128 @@
+import logging
+import time
+import typing
+from datetime import datetime, timedelta
+from threading import Thread
+
+from django.conf import settings
+from django.db import close_old_connections
+from django.utils import timezone
+
+from task_processor.processor import run_recurring_tasks, run_tasks
+from task_processor.task_registry import initialise
+from task_processor.types import TaskProcessorConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TaskRunnerCoordinator(Thread):
+    def __init__(
+        self,
+        *args: typing.Any,
+        config: TaskProcessorConfig,
+        **kwargs: typing.Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.config = config
+        self._threads: list[TaskRunner] = []
+        self._monitor_threads = True
+
+    def run(self) -> None:
+        initialise()
+
+        logger.info("Processor starting")
+
+        for _ in range(self.config.num_threads):
+            self._threads.append(
+                task := TaskRunner(
+                    sleep_interval_millis=self.config.sleep_interval_ms,
+                    queue_pop_size=self.config.queue_pop_size,
+                )
+            )
+            task.start()
+
+        ms_before_unhealthy = (
+            self.config.grace_period_ms + self.config.sleep_interval_ms
+        )
+        while self._monitor_threads:
+            time.sleep(1)
+            unhealthy_threads = self._get_unhealthy_threads(
+                ms_before_unhealthy=ms_before_unhealthy
+            )
+            if unhealthy_threads:
+                logger.warning("%d unhealthy threads detected", len(unhealthy_threads))
+
+        for thread in self._threads:
+            thread.join()
+
+    def _get_unhealthy_threads(self, ms_before_unhealthy: int) -> list["TaskRunner"]:
+        unhealthy_threads = []
+        healthy_threshold = timezone.now() - timedelta(milliseconds=ms_before_unhealthy)
+
+        for thread in self._threads:
+            if (
+                not thread.is_alive()
+                or not thread.last_checked_for_tasks
+                or thread.last_checked_for_tasks < healthy_threshold
+            ):
+                unhealthy_threads.append(thread)
+        return unhealthy_threads
+
+    def stop(self) -> None:
+        self._monitor_threads = False
+        for t in self._threads:
+            t.stop()
+
+
+class TaskRunner(Thread):
+    def __init__(
+        self,
+        *args: typing.Any,
+        sleep_interval_millis: int = 2000,
+        queue_pop_size: int = 1,
+        **kwargs: typing.Any,
+    ):
+        super(TaskRunner, self).__init__(*args, **kwargs)
+        self.sleep_interval_millis = sleep_interval_millis
+        self.queue_pop_size = queue_pop_size
+        self.last_checked_for_tasks: datetime | None = None
+
+        self._stopped = False
+
+    def run(self) -> None:
+        while not self._stopped:
+            self.last_checked_for_tasks = timezone.now()
+            self.run_iteration()
+            time.sleep(self.sleep_interval_millis / 1000)
+
+    def run_iteration(self) -> None:
+        """
+        Consume and execute tasks from the queue, and run recurring tasks
+
+        This method tries to consume tasks from multiple databases as to ensure
+        that any remaining tasks are processed after opting in or out of a
+        separate database setup.
+        """
+        database_is_separate = "task_processor" in settings.TASK_PROCESSOR_DATABASES
+
+        for database in settings.TASK_PROCESSOR_DATABASES:
+            try:
+                run_tasks(database, self.queue_pop_size)
+
+                # Recurring tasks are only run on one database
+                if (database == "default") ^ database_is_separate:
+                    run_recurring_tasks(database)
+            except Exception as exception:
+                # To prevent task threads from dying if they get an error retrieving the tasks from the
+                # database this will allow the thread to continue trying to retrieve tasks if it can
+                # successfully re-establish a connection to the database.
+                exception_repr = f"{exception.__class__.__module__}.{repr(exception)}"
+                logger.error(
+                    f"Error handling tasks from database '{database}': {exception_repr}",
+                    exc_info=exception,
+                )
+
+                close_old_connections()
+
+    def stop(self) -> None:
+        self._stopped = True

task_processor/types.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import Callable, ParamSpec, TypeAlias, TypedDict
+
+TaskParameters = ParamSpec("TaskParameters")
+
+TaskCallable: TypeAlias = Callable[TaskParameters, None]
+
+
+@dataclass
+class TaskProcessorConfig:
+    num_threads: int
+    sleep_interval_ms: int
+    grace_period_ms: int
+    queue_pop_size: int
+
+
+class MonitoringInfo(TypedDict):
+    waiting: int

task_processor/urls.py

@@ -0,0 +1,5 @@
+from django.urls import path
+
+from task_processor.views import monitoring
+
+urlpatterns = [path("monitoring/", monitoring)]

task_processor/utils.py

@@ -0,0 +1,71 @@
+import argparse
+import inspect
+import logging
+from contextlib import contextmanager
+from typing import Any, Generator
+
+from task_processor.threads import TaskRunnerCoordinator
+from task_processor.types import TaskCallable, TaskProcessorConfig
+
+logger = logging.getLogger(__name__)
+
+
+def get_task_identifier_from_function(
+    function: TaskCallable[Any],
+    task_name: str | None,
+) -> str:
+    module = inspect.getmodule(function)
+    assert module
+    return f"{module.__name__.rsplit('.')[-1]}.{task_name or function.__name__}"
+
+
+def add_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--numthreads",
+        type=int,
+        help="Number of worker threads to run.",
+        default=5,
+    )
+    parser.add_argument(
+        "--sleepintervalms",
+        type=int,
+        help="Number of millis each worker waits before checking for new tasks",
+        default=2000,
+    )
+    parser.add_argument(
+        "--graceperiodms",
+        type=int,
+        help="Number of millis before running task is considered 'stuck'.",
+        default=20000,
+    )
+    parser.add_argument(
+        "--queuepopsize",
+        type=int,
+        help="Number of tasks each worker will pop from the queue on each cycle.",
+        default=10,
+    )
+
+
+@contextmanager
+def start_task_processor(
+    options: dict[str, Any],
+) -> Generator[
+    TaskRunnerCoordinator,
+    None,
+    None,
+]:
+    config = TaskProcessorConfig(
+        num_threads=options["numthreads"],
+        sleep_interval_ms=options["sleepintervalms"],
+        grace_period_ms=options["graceperiodms"],
+        queue_pop_size=options["queuepopsize"],
+    )
+
+    logger.debug("Config: %s", config)
+
+    coordinator = TaskRunnerCoordinator(config=config)
+    coordinator.start()
+    try:
+        yield coordinator
+    finally:
+        coordinator.stop()

task_processor/views.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+from drf_yasg.utils import swagger_auto_schema  # type: ignore[import-untyped]
+from rest_framework.decorators import api_view, permission_classes
+from rest_framework.permissions import IsAdminUser, IsAuthenticated
+from rest_framework.request import Request
+from rest_framework.response import Response
+
+from task_processor.monitoring import get_num_waiting_tasks
+from task_processor.serializers import MonitoringSerializer
+
+
+@swagger_auto_schema(method="GET", responses={200: MonitoringSerializer()})  # type: ignore[misc]
+@api_view(http_method_names=["GET"])
+@permission_classes([IsAuthenticated, IsAdminUser])
+def monitoring(request: Request, /, **kwargs: Any) -> Response:
+    return Response(
+        data={"waiting": get_num_waiting_tasks()},
+        content_type="application/json",
+    )