Consume tasks from multiple databases

Author: emyller

poetry.lock

@@ -62,3 +62,5 @@
 
 # Avoid models.W042 warnings
 DEFAULT_AUTO_FIELD = "django.db.models.AutoField"
+
+TASK_PROCESSOR_DATABASES = ["default"]

settings/dev.py

@@ -31,8 +31,7 @@ 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 = list(Task.objects.get_tasks_to_process(database, num_tasks))
-
+    tasks = Task.objects.get_tasks_to_process(database, num_tasks)
     if tasks:
         logger.debug(f"Running {len(tasks)} task(s) from database '{database}'")
 
@@ -55,7 +54,7 @@ def run_tasks(database: str, 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) from database '{database}'",
+                f"Finished running {len(task_runs)} task(s) from database '{database}'"
             )
 
         return task_runs
@@ -69,9 +68,7 @@ def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
     # a problem for now, but we should be mindful of this limitation
     tasks = RecurringTask.objects.get_tasks_to_process(database)
     if tasks:
-        logger.debug(
-            f"Running {len(tasks)} recurring task(s) from database '{database}'",
-        )
+        logger.debug(f"Running {len(tasks)} recurring task(s)")
 
         task_runs = []
 
@@ -99,9 +96,7 @@ def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
 
         if task_runs:
             RecurringTaskRun.objects.bulk_create(task_runs)
-            logger.debug(
-                f"Finished running {len(task_runs)} recurring task(s) from database '{database}'",
-            )
+            logger.debug(f"Finished running {len(task_runs)} recurring task(s)")
 
         return task_runs
 

src/task_processor/processor.py

@@ -4,6 +4,7 @@
 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
 
@@ -95,23 +96,33 @@ def run(self) -> None:
             time.sleep(self.sleep_interval_millis / 1000)
 
     def run_iteration(self) -> None:
-        for database in ["default", "task_processor"]:
+        """
+        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)
-                run_recurring_tasks(database)
+
+                # 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.
                 # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
                 #  tasks here?
-
+                exception_repr = f"{exception.__class__.__module__}.{repr(exception)}"
                 logger.error(
-                    "Received error retrieving tasks from database '%s': %s.",
-                    database,
-                    repr(exception),
+                    f"Error handling tasks from database '{database}': {exception_repr}",
                     exc_info=exception,
                 )
+
                 close_old_connections()
 
     def stop(self) -> None:

src/task_processor/threads.py

@@ -432,7 +432,7 @@ def test_run_task_runs_task_and_creates_task_run_object_when_failure(
     assert not task.completed
 
     expected_log_records = [
-        ("DEBUG", "Running 1 task(s)"),
+        ("DEBUG", "Running 1 task(s) from database 'default'"),
         (
             "DEBUG",
             f"Running task {task.task_identifier} id={task.id} args={task.args} kwargs={task.kwargs}",
@@ -441,7 +441,7 @@ def test_run_task_runs_task_and_creates_task_run_object_when_failure(
             "ERROR",
             f"Failed to execute task '{task.task_identifier}', with id {task.id}. Exception: {msg}",
         ),
-        ("DEBUG", "Finished running 1 task(s)"),
+        ("DEBUG", "Finished running 1 task(s) from database 'default'"),
     ]
 
     assert expected_log_records == [
@@ -698,7 +698,7 @@ def test_run_tasks_skips_locked_tasks(
 
     # When
     # we spawn a new thread to run the first task (configured to just sleep)
-    task_runner_thread = Thread(target=run_tasks)
+    task_runner_thread = Thread(target=run_tasks, args=("default",))
     task_runner_thread.start()
 
     # and subsequently attempt to run another task in the main thread