WIP: New PoC for processing tasks from multiple databases

Author: emyller

src/task_processor/managers.py

@@ -3,16 +3,25 @@
 from django.db.models import Manager
 
 if typing.TYPE_CHECKING:
-    from django.db.models.query import RawQuerySet
-
     from task_processor.models import RecurringTask, Task
 
 
 class TaskManager(Manager["Task"]):
-    def get_tasks_to_process(self, num_tasks: int) -> "RawQuerySet[Task]":
-        return self.raw("SELECT * FROM get_tasks_to_process(%s)", [num_tasks])
+    def get_tasks_to_process(
+        self,
+        database: str,
+        num_tasks: int,
+    ) -> typing.List["Task"]:
+        return list(
+            self.using(database).raw(
+                "SELECT * FROM get_tasks_to_process(%s)",
+                [num_tasks],
+            ),
+        )
 
 
 class RecurringTaskManager(Manager["RecurringTask"]):
-    def get_tasks_to_process(self) -> "RawQuerySet[RecurringTask]":
-        return self.raw("SELECT * FROM get_recurringtasks_to_process()")
+    def get_tasks_to_process(self, database: str) -> typing.List["RecurringTask"]:
+        return list(
+            self.using(database).raw("SELECT * FROM get_recurringtasks_to_process()"),
+        )

src/task_processor/processor.py

@@ -27,14 +27,14 @@
 UNREGISTERED_RECURRING_TASK_GRACE_PERIOD = timedelta(minutes=30)
 
 
-def run_tasks(num_tasks: int = 1) -> list[TaskRun]:
+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")
 
-    tasks = Task.objects.get_tasks_to_process(num_tasks)
+    tasks = list(Task.objects.get_tasks_to_process(database, num_tasks))
 
     if tasks:
-        logger.debug(f"Running {len(tasks)} task(s)")
+        logger.debug(f"Running {len(tasks)} task(s) from database '{database}'")
 
         executed_tasks = []
         task_runs = []
@@ -54,20 +54,24 @@ def run_tasks(num_tasks: int = 1) -> list[TaskRun]:
 
         if task_runs:
             TaskRun.objects.bulk_create(task_runs)
-            logger.debug(f"Finished running {len(task_runs)} task(s)")
+            logger.debug(
+                f"Finished running {len(task_runs)} task(s) from database '{database}'",
+            )
 
         return task_runs
 
     return []
 
 
-def run_recurring_tasks() -> list[RecurringTaskRun]:
+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
-    tasks = RecurringTask.objects.get_tasks_to_process()
+    tasks = RecurringTask.objects.get_tasks_to_process(database)
     if tasks:
-        logger.debug(f"Running {len(tasks)} recurring task(s)")
+        logger.debug(
+            f"Running {len(tasks)} recurring task(s) from database '{database}'",
+        )
 
         task_runs = []
 
@@ -95,7 +99,9 @@ def run_recurring_tasks() -> list[RecurringTaskRun]:
 
         if task_runs:
             RecurringTaskRun.objects.bulk_create(task_runs)
-            logger.debug(f"Finished running {len(task_runs)} recurring task(s)")
+            logger.debug(
+                f"Finished running {len(task_runs)} recurring task(s) from database '{database}'",
+            )
 
         return task_runs
 

src/task_processor/threads.py

@@ -95,18 +95,24 @@ def run(self) -> None:
             time.sleep(self.sleep_interval_millis / 1000)
 
     def run_iteration(self) -> None:
-        try:
-            run_tasks(self.queue_pop_size)
-            run_recurring_tasks()
-        except Exception as e:
-            # 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.
-            # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
-            #  tasks here?
-
-            logger.error("Received error retrieving tasks: %s.", e, exc_info=e)
-            close_old_connections()
+        for database in ["default", "task_processor"]:
+            try:
+                run_tasks(database, self.queue_pop_size)
+                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.
+                # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
+                #  tasks here?
+
+                logger.error(
+                    "Received error retrieving tasks from database '%s': %s.",
+                    database,
+                    repr(exception),
+                    exc_info=exception,
+                )
+                close_old_connections()
 
     def stop(self) -> None:
         self._stopped = True

tests/unit/task_processor/test_unit_task_processor_processor.py

@@ -85,7 +85,7 @@ def test_run_task_runs_task_and_creates_task_run_object_when_success(
     task.save()
 
     # When
-    task_runs = run_tasks()
+    task_runs = run_tasks("default")
 
     # Then
     assert cache.get(DEFAULT_CACHE_KEY)
@@ -116,7 +116,7 @@ def test_run_task_kills_task_after_timeout(
     task.save()
 
     # When
-    task_runs = run_tasks()
+    task_runs = run_tasks("default")
 
     # Then
     assert len(task_runs) == TaskRun.objects.filter(task=task).count() == 1
@@ -157,7 +157,7 @@ def _dummy_recurring_task() -> None:
         task_identifier="test_unit_task_processor_processor._dummy_recurring_task",
     )
     # When
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert len(task_runs) == RecurringTaskRun.objects.filter(task=task).count() == 1
@@ -195,7 +195,7 @@ def _dummy_recurring_task() -> None:
         task_identifier="test_unit_task_processor_processor._dummy_recurring_task",
     )
     # When
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert cache.get(DEFAULT_CACHE_KEY)
@@ -227,7 +227,7 @@ def _dummy_recurring_task() -> None:
 
     # When
     assert cache.get(DEFAULT_CACHE_KEY) is None
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert cache.get(DEFAULT_CACHE_KEY) == DEFAULT_CACHE_VALUE
@@ -261,16 +261,16 @@ def _dummy_recurring_task() -> None:
     )
 
     # When
-    first_task_runs = run_recurring_tasks()
+    first_task_runs = run_recurring_tasks("default")
 
     # run the process again before the task is scheduled to run again to ensure
     # that tasks are unlocked when they are picked up by the task processor but
     # not executed.
-    no_task_runs = run_recurring_tasks()
+    no_task_runs = run_recurring_tasks("default")
 
     time.sleep(0.3)
 
-    second_task_runs = run_recurring_tasks()
+    second_task_runs = run_recurring_tasks("default")
 
     # Then
     assert len(first_task_runs) == 1
@@ -310,7 +310,7 @@ def _dummy_recurring_task_3() -> None:
 
     # When, we call run_recurring_tasks in a loop few times
     for _ in range(4):
-        run_recurring_tasks()
+        run_recurring_tasks("default")
 
     # Then - we should have exactly one RecurringTaskRun for each task
     for i in range(1, 4):
@@ -339,8 +339,8 @@ def _dummy_recurring_task() -> None:
     )
 
     # When - we call run_recurring_tasks twice
-    run_recurring_tasks()
-    run_recurring_tasks()
+    run_recurring_tasks("default")
+    run_recurring_tasks("default")
 
     # Then - we expect the task to have been run once
 
@@ -367,7 +367,7 @@ def _a_task() -> None:
     registered_tasks.pop(task_identifier)
 
     # When
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert len(task_runs) == 0
@@ -395,7 +395,7 @@ def _a_task() -> None:
     registered_tasks.pop(task_identifier)
 
     # When
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert len(task_runs) == 0
@@ -418,7 +418,7 @@ def test_run_task_runs_task_and_creates_task_run_object_when_failure(
     task.save()
 
     # When
-    task_runs = run_tasks()
+    task_runs = run_tasks("default")
 
     # Then
     assert len(task_runs) == TaskRun.objects.filter(task=task).count() == 1
@@ -460,10 +460,10 @@ def test_run_task_runs_failed_task_again(
     task.save()
 
     # When
-    first_task_runs = run_tasks()
+    first_task_runs = run_tasks("default")
 
     # Now, let's run the task again
-    second_task_runs = run_tasks()
+    second_task_runs = run_tasks("default")
 
     # Then
     task_runs = first_task_runs + second_task_runs
@@ -498,7 +498,7 @@ def _raise_exception(organisation_name: str) -> None:
     task = RecurringTask.objects.get(task_identifier=task_identifier)
 
     # When
-    task_runs = run_recurring_tasks()
+    task_runs = run_recurring_tasks("default")
 
     # Then
     assert len(task_runs) == RecurringTaskRun.objects.filter(task=task).count() == 1
@@ -513,7 +513,7 @@ def _raise_exception(organisation_name: str) -> None:
 def test_run_task_does_nothing_if_no_tasks() -> None:
     # Given - no tasks
     # When
-    result = run_tasks()
+    result = run_tasks("default")
     # Then
     assert result == []
     assert not TaskRun.objects.exists()
@@ -551,9 +551,9 @@ def test_run_task_runs_tasks_in_correct_priority(
     task_3.save()
 
     # When
-    task_runs_1 = run_tasks()
-    task_runs_2 = run_tasks()
-    task_runs_3 = run_tasks()
+    task_runs_1 = run_tasks("default")
+    task_runs_2 = run_tasks("default")
+    task_runs_3 = run_tasks("default")
 
     # Then
     assert task_runs_1[0].task == task_3
@@ -573,7 +573,7 @@ def test_run_tasks__fails_if_not_in_task_processor_mode(
 
     # When
     with pytest.raises(AssertionError):
-        run_tasks()
+        run_tasks("default")
 
 
 @pytest.mark.django_db(transaction=True)
@@ -609,8 +609,8 @@ def _fake_recurring_task() -> None:
     ).save()
 
     # When
-    run_tasks(2)
-    run_recurring_tasks()
+    run_tasks("default", 2)
+    run_recurring_tasks("default")
 
     # Then
     assert_metric(
@@ -703,7 +703,7 @@ def test_run_tasks_skips_locked_tasks(
 
     # and subsequently attempt to run another task in the main thread
     time.sleep(1)  # wait for the thread to start and hold the task
-    task_runs = run_tasks()
+    task_runs = run_tasks("default")
 
     # Then
     # the second task is run while the 1st task is held
@@ -730,7 +730,7 @@ def test_run_more_than_one_task(dummy_task: TaskHandler[[str, str]]) -> None:
     Task.objects.bulk_create(tasks)
 
     # When
-    task_runs = run_tasks(5)
+    task_runs = run_tasks("default", 5)
 
     # Then
     assert len(task_runs) == num_tasks
@@ -772,7 +772,7 @@ def my_task() -> None:
     )
 
     # When
-    run_recurring_tasks()
+    run_recurring_tasks("default")
 
     # Then
     recurring_task.refresh_from_db()