Basics of multi-queue working

Author: AlanCoding

ansible_base/lib/dynamic_config/settings_logic.py

@@ -336,4 +336,12 @@ def get_dab_settings(
         # if you want to use a local queue instead, that would be more performant
         dab_data['DAB_TASK_ADMIN_QUEUE'] = 'dab_broadcast'
 
+        # Queues that the local node will listen to
+        # useful to manage different pools of work and types of nodes
+        # this should probably include the admin queue
+        dab_data['DAB_TASK_LISTEN_QUEUES'] = ['dab_broadcast']
+
+        # Dispatcher spawns python subprocesses, so you can control the number of them here
+        dab_data['DAB_TASK_MAX_WORKERS'] = 4
+
     return dab_data

ansible_base/lib/utils/db.py

@@ -160,6 +160,11 @@ def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
 dj_db_dict = dict[str, Union[str, int]]
 
 
+def psycopg_connection_from_django(**kwargs) -> psycopg.Connection:
+    "Compatibility with dispatcher connection factory, just returns the Django connection"
+    return connection.connection
+
+
 def psycopg_kwargs_from_settings_dict(settings_dict: dj_db_dict) -> dict:
     """Return psycopg connection creation kwargs given Django db settings info
 

ansible_base/task/config.py

@@ -11,8 +11,12 @@ def get_config() -> dict:
         "version": 2,
         "brokers": {
             "pg_notify": {
-                "config": psycopg_params
-                # TODO: hook in synchronous connection factory
+                "config": psycopg_params,  # used to create the service async connection
+                "sync_connection_factory": "ansible_base.lib.utils.db.psycopg_connection_from_django",
+                "channels": settings.DAB_TASK_LISTEN_QUEUES,
+                # The default publish channel allows using .delay without more arguments
+                # this may still have task-specific overrides
+                "default_publish_channel": settings.DAB_TASK_ADMIN_QUEUE
             }
         },
         "producers": {
@@ -23,5 +27,5 @@ def get_config() -> dict:
                 }
             },
         },
-        "pool": {"max_workers": 4},  # TODO: to settings
+        "pool": {"max_workers": settings.DAB_TASK_MAX_WORKERS},
     }

ansible_base/task/migrations/0004_task_queue_alter_task_wrapper_uuid.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.16 on 2025-02-25 21:32
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dab_tas', '0003_task_args_task_kwargs'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='task',
+            name='queue',
+            field=models.TextField(default='', help_text='The work queue this task is submitted to, only nodes listening to this queue may process the task'),
+        ),
+        migrations.AlterField(
+            model_name='task',
+            name='wrapper_uuid',
+            field=models.CharField(blank=True, default='', editable=False, help_text='Adopted UUID of DAB task app wrapper, attached when background service starts task', max_length=40),
+        ),
+    ]

ansible_base/task/models.py

@@ -25,14 +25,17 @@ class Task(models.Model):
     """
 
     uuid = models.UUIDField(primary_key=True, default=uuid4, editable=False, help_text=_('Original dispatcher UUID generated at time of task publishing'))
-    wrapper_uuid = models.UUIDField(null=True, default=None, editable=False, help_text=_('Adopted UUID of DAB task app wrapper, attached when background service starts task'))
+    # TODO: change back to UUIDField, workaround dispatcher using non-UUID uuids for now
+    # null=True
+    wrapper_uuid = models.CharField(blank=True, default='', max_length=40, editable=False, help_text=_('Adopted UUID of DAB task app wrapper, attached when background service starts task'))
     state = models.CharField(
         choices=[(s, s.title()) for s in sorted(vars(TASK_STATES).values())],
         default=TASK_STATES.WAITING,
         max_length=15,
         help_text=_('Choices of this field track with acknowledgement and completion of a task'),
     )
     name = models.TextField(help_text=_('Importable path for class or method'))
+    queue = models.TextField(default='', help_text=_('The work queue this task is submitted to, only nodes listening to this queue may process the task'))
 
     args = models.JSONField(
         default=dict, null=False, blank=True, help_text=_("The arguments for the task run.")

ansible_base/task/publish.py

@@ -1,5 +1,6 @@
 from dispatcher.utils import serialize_task
 from django.db import transaction
+from django.conf import settings
 
 from ansible_base.task.models import Task
 from ansible_base.task.tasks import run_task_from_queue
@@ -8,38 +9,44 @@
 
 
 class TaskPublisher:
-    def __init__(self, fn):
+    def __init__(self, fn, queue=None):
         self.fn = fn
+        self.queue = queue
 
     @property
     def task_name(self):
         return serialize_task(self.fn)
 
+    def submit_wrapper_task(self):
+        run_task_from_queue.apply_async(queue=self.queue)
+
     def apply_async(self, args=None, kwargs=None):
         # this function may allow additional arguments in the future, but not now
-        Task.objects.create(name=self.task_name, args=args, kwargs=kwargs)
+        if self.queue is None:
+            queue = settings.DAB_TASK_ADMIN_QUEUE
+        else:
+            queue = self.queue
+        Task.objects.create(name=self.task_name, args=args, kwargs=kwargs, queue=queue)
 
         # pg_notify message (probably) to wake up
-        transaction.on_commit(run_task_from_queue.delay)
+        transaction.on_commit(self.submit_wrapper_task)
 
     def delay(self, *args, **kwargs):
         return self.apply_async(args=args, kwargs=kwargs)
 
 
 class TaskDecorator:
-    def __init__(self, *args, **kwargs):
-        # TODO: this will process a task timeout, just not yet set up
-        self.args = args
-        self.kwargs = kwargs
+    def __init__(self, queue=None):
+        self.queue = queue
 
     def __call__(self, fn):
-        publisher = TaskPublisher(fn)
+        publisher = TaskPublisher(fn, queue=self.queue)
 
         setattr(fn, 'apply_async', publisher.apply_async)
         setattr(fn, 'delay', publisher.delay)
 
         return fn
 
 
-def durable_task():
-    return TaskDecorator()
+def durable_task(queue=None):
+    return TaskDecorator(queue=queue)

ansible_base/task/tasks.py

@@ -2,14 +2,12 @@
 
 from dispatcher.utils import resolve_callable
 from dispatcher.publish import task
-from dispatcher.control import Control
+from dispatcher.factories import get_control_from_settings
 
 from django.db import transaction
 from django.conf import settings
 from django.utils.timezone import now, timedelta
 
-from ansible_base.lib.utils.db import get_pg_notify_params
-
 from ansible_base.task.models import TASK_STATES, Task
 
 logger = logging.getLogger(__name__)
@@ -18,11 +16,11 @@
 @task(queue=settings.DAB_TASK_ADMIN_QUEUE, bind=True)
 def run_task_from_queue(dispatcher):
     with transaction.atomic():
-        task = Task.objects.filter(state=TASK_STATES.WAITING).select_for_update().first()
+        task = Task.objects.filter(state=TASK_STATES.WAITING, queue__in=settings.DAB_TASK_LISTEN_QUEUES).select_for_update().first()
         if task:
             task.state = TASK_STATES.RUNNING
             task.started_at = now()
-            task.wrapper_uuid = dispatcher.uuid
+            task.wrapper_uuid = str(dispatcher.uuid)
             task.save(update_fields=['state', 'started_at', 'wrapper_uuid'])
         else:
             return
@@ -35,21 +33,16 @@ def run_task_from_queue(dispatcher):
         task_callable = resolve_callable(task.name)
         task_callable(*task.args, **task.kwargs)
     except Exception:
-        logger.traceback(f'Failed to run and complete {task.name}')
+        logger.exception(f'Failed to run and complete {task.name}')
 
     task.delete()
 
 
 @task(queue=settings.DAB_TASK_ADMIN_QUEUE)
 def manage_lost_tasks(grace_period: int = 10):
     cutoff_time = now() - timedelta(minutes=grace_period)
-    for task in Task.objects.filter(state=TASK_STATES.RUNNING, started_at__lt=cutoff_time).iterator():
-        psycopg_params = get_pg_notify_params()
-        psycopg_params.pop('autocommit')  # dispatcher automatically adds this, causes error, TODO: need pre-check
-        psycopg_params.pop('cursor_factory')
-        psycopg_params.pop('context')  # TODO: remove in inner method, makes non-async, not good
-
-        ctl = Control(settings.DAB_TASK_ADMIN_QUEUE, config=psycopg_params)
+    for task in Task.objects.filter(state=TASK_STATES.RUNNING, started_at__lt=cutoff_time, queue__in=settings.DAB_TASK_LISTEN_QUEUES).iterator():
+        ctl = get_control_from_settings(default_publish_channel=settings.DAB_TASK_ADMIN_QUEUE)
 
         running_tasks = ctl.control_with_reply('running', data={'uuid': str(task.wrapper_uuid)})
 
@@ -66,7 +59,7 @@ def manage_lost_tasks(grace_period: int = 10):
             # TODO: feature of retry policy
             try:
                 with transaction.atomic():
-                    task = Task.objects.get(uuid=task.uuid).select_for_update()
+                    task = Task.objects.select_for_update().get(uuid=task.uuid)
                     logger.warning(f'Could not find task {task.name} {task.wrapper_uuid}, deleting entry')
                     task.delete()
             except Task.DoesNotExist:

test_app/settings.py

@@ -196,6 +196,8 @@
 ALLOW_SHARED_RESOURCE_CUSTOM_ROLES = True  # Allow making custom roles with org change permission, for example
 ALLOW_LOCAL_ASSIGNING_JWT_ROLES = False
 
+DAB_TASK_LISTEN_QUEUES = ['dab_broadcast', 'dab_task_tasks']
+
 ANSIBLE_BASE_USER_VIEWSET = 'test_app.views.UserViewSet'
 
 LOGIN_URL = "/login/login"

test_app/tasks.py

@@ -11,6 +11,11 @@ def hello_world():
     print('hello world')
 
 
+@durable_task(queue='dab_task_tasks')
+def hello_world_other_queue():
+    print('hello world, sent from test_app_tasks queue')
+
+
 @durable_task()
 def create_uuid_entry(uuid: UUID):
     UUIDModel.objects.create(id=uuid)