fix logs when task fails, make more options configurable

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

authentik/lib/default.yml

@@ -160,6 +160,8 @@ worker:
   processes: 2
   threads: 1
   consumer_listen_timeout: "seconds=30"
+  task_max_retries: 20
+  task_default_time_limit: "minutes=10"
   task_purge_interval: "days=1"
   task_expiration: "days=30"
   scheduler_interval: "seconds=60"

authentik/root/settings.py

@@ -380,23 +380,27 @@ DRAMATIQ = {
         CONFIG.get("worker.scheduler_interval")
     ).total_seconds(),
     "middlewares": (
-        ("django_dramatiq_postgres.middleware.SchedulerMiddleware", {}),
+        # ("django_dramatiq_postgres.middleware.SchedulerMiddleware", {}),
         ("django_dramatiq_postgres.middleware.FullyQualifiedActorName", {}),
         # TODO: fixme
         # ("dramatiq.middleware.prometheus.Prometheus", {}),
         ("django_dramatiq_postgres.middleware.DbConnectionMiddleware", {}),
         ("dramatiq.middleware.age_limit.AgeLimit", {}),
         (
-            # 5 minutes task timeout by default for all tasks, in ms
             "dramatiq.middleware.time_limit.TimeLimit",
-            {"time_limit": 600_000},
+            {
+                "time_limit": timedelta_from_string(
+                    CONFIG.get("worker.task_default_time_limit")
+                ).total_seconds()
+                * 1000
+            },
         ),
         ("dramatiq.middleware.shutdown.ShutdownNotifications", {}),
         ("dramatiq.middleware.callbacks.Callbacks", {}),
         ("dramatiq.middleware.pipelines.Pipelines", {}),
         (
             "dramatiq.middleware.retries.Retries",
-            {"max_retries": 20 if not TEST else 0},
+            {"max_retries": CONFIG.get_int("worker.task_max_retries") if not TEST else 0},
         ),
         # TODO: results
         ("django_dramatiq_postgres.middleware.CurrentTask", {}),

authentik/tasks/middleware.py

@@ -1,5 +1,6 @@
 from typing import Any
 
+from dramatiq import get_logger
 from dramatiq.broker import Broker
 from dramatiq.message import Message
 from dramatiq.middleware import Middleware
@@ -35,21 +36,23 @@ class RelObjMiddleware(Middleware):
 
 
 class LoggingMiddleware(Middleware):
-    def before_enqueue(self, broker: Broker, message: Message, delay: int):
-        message.options["model_defaults"]["_messages"] = [
+    def after_enqueue(self, broker: Broker, message: Message, delay: int):
+        task: Task = message.options["task"]
+        task_created: bool = message.options["task_created"]
+        task._messages.append(
             Task._make_message(
                 str(type(self)),
                 TaskStatus.INFO,
-                "Task is being queued",
+                "Task is being queued" if task_created else "Task is being retried",
                 delay=delay,
             )
-        ]
+        )
+        task.save(update_fields=("_messages",))
 
     def before_process_message(self, broker: Broker, message: Message):
         task: Task = message.options["task"]
         task.log(str(type(self)), TaskStatus.INFO, "Task is being processed")
 
-    # TODO: also after_skip_message
     def after_process_message(
         self,
         broker: Broker,
@@ -74,6 +77,10 @@ class LoggingMiddleware(Middleware):
                 "{exception_to_string(exception)}",
             ).save()
 
+    def after_skip_message(self, broker: Broker, message: Message):
+        task: Task = message.options["task"]
+        task.log(str(type(self)), TaskStatus.INFO, "Task has been skipped")
+
 
 class DescriptionMiddleware(Middleware):
     @property

packages/django-dramatiq-postgres/django_dramatiq_postgres/broker.py

@@ -14,6 +14,7 @@ from django.db import (
     InterfaceError,
     OperationalError,
     connections,
+    transaction,
 )
 from django.db.backends.postgresql.base import DatabaseWrapper
 from django.db.models import QuerySet
@@ -152,23 +153,26 @@ class PostgresBroker(Broker):
         message.options["model_defaults"] = self.model_defaults(message)
         self.emit_before("enqueue", message, delay)
 
-        query = {
-            "message_id": message.message_id,
-        }
-        defaults = message.options["model_defaults"]
-        del message.options["model_defaults"]
-        create_defaults = {
-            **query,
-            **defaults,
-        }
+        with transaction.atomic(using=self.db_alias):
+            query = {
+                "message_id": message.message_id,
+            }
+            defaults = message.options["model_defaults"]
+            del message.options["model_defaults"]
+            create_defaults = {
+                **query,
+                **defaults,
+            }
 
-        self.query_set.update_or_create(
-            **query,
-            defaults=defaults,
-            create_defaults=create_defaults,
-        )
+            task, created = self.query_set.update_or_create(
+                **query,
+                defaults=defaults,
+                create_defaults=create_defaults,
+            )
+            message.options["task"] = task
+            message.options["task_created"] = created
 
-        self.emit_after("enqueue", message, delay)
+            self.emit_after("enqueue", message, delay)
         return message
 
     def get_declared_queues(self) -> set[str]:

packages/django-dramatiq-postgres/django_dramatiq_postgres/middleware.py

@@ -1,6 +1,6 @@
 import contextvars
 from threading import Event
-from typing import Any
+from typing import Any, override
 
 from django.core.exceptions import ImproperlyConfigured
 from django.db import (
@@ -58,6 +58,9 @@ class CurrentTask(Middleware):
             raise RuntimeError("CurrentTask.get_task() can only be called in a running task")
         return task[-1]
 
+    def before_enqueue(self, broker: Broker, message: Message, delay: int):
+        self.after_process_message(broker, message)
+
     def before_process_message(self, broker: Broker, message: Message):
         tasks = self._TASKS.get()
         if tasks is None:
@@ -75,29 +78,31 @@ class CurrentTask(Middleware):
     ):
         tasks: list[TaskBase] | None = self._TASKS.get()
         if tasks is None or len(tasks) == 0:
-            self.logger.warning("Task was None, not saving. This should not happen.")
             return
-        else:
-            task = tasks[-1]
-            fields_to_exclude = {
-                "message_id",
-                "queue_name",
-                "actor_name",
-                "message",
-                "state",
-                "mtime",
-                "result",
-                "result_expiry",
-            }
-            fields_to_update = [
-                f.name
-                for f in task._meta.get_fields()
-                if f.name not in fields_to_exclude and not f.auto_created and f.column
-            ]
-            if fields_to_update:
-                tasks[-1].save(update_fields=fields_to_update)
+
+        task = tasks[-1]
+        fields_to_exclude = {
+            "message_id",
+            "queue_name",
+            "actor_name",
+            "message",
+            "state",
+            "mtime",
+            "result",
+            "result_expiry",
+        }
+        fields_to_update = [
+            f.name
+            for f in task._meta.get_fields()
+            if f.name not in fields_to_exclude and not f.auto_created and f.column
+        ]
+        if fields_to_update:
+            task.save(update_fields=fields_to_update)
         self._TASKS.set(tasks[:-1])
 
+    def after_skip_message(self, broker: Broker, message: Message):
+        self.after_process_message(broker, message)
+
 
 class SchedulerMiddleware(Middleware):
     def __init__(self):