feat(suse,broker): Move away from session-level lock into transaction-level locks

authentik/lib/logging.py

@@ -33,21 +33,41 @@ def get_log_level():
 
 def structlog_configure():
     """Configure structlog itself"""
+    is_debug_enabled = CONFIG.get_bool("debug")
+
+    processors = [
+        structlog.stdlib.add_log_level,
+        structlog.stdlib.add_logger_name,
+        structlog.contextvars.merge_contextvars,
+        add_process_id,
+        add_tenant_information,
+        structlog.stdlib.PositionalArgumentsFormatter(),
+        structlog.processors.TimeStamper(fmt="iso", utc=False),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.ExceptionRenderer(
+            structlog.tracebacks.ExceptionDictTransformer(show_locals=is_debug_enabled)
+        ),
+    ]
+
+    if get_log_level() == "DEBUG" or is_debug_enabled:
+        # Add call site references to log lines
+        processors.append(
+            structlog.processors.CallsiteParameterAdder(
+                {
+                    structlog.processors.CallsiteParameter.PATHNAME,
+                    structlog.processors.CallsiteParameter.FUNC_NAME,
+                    structlog.processors.CallsiteParameter.LINENO,
+                    structlog.processors.CallsiteParameter.THREAD_NAME,
+                }
+            )
+        )
+
+    # this processor must be last in the chain
+    # https://www.structlog.org/en/stable/api.html#structlog.stdlib.ProcessorFormatter.wrap_for_formatter
+    processors.append(structlog.stdlib.ProcessorFormatter.wrap_for_formatter)
+
     structlog.configure_once(
-        processors=[
-            structlog.stdlib.add_log_level,
-            structlog.stdlib.add_logger_name,
-            structlog.contextvars.merge_contextvars,
-            add_process_id,
-            add_tenant_information,
-            structlog.stdlib.PositionalArgumentsFormatter(),
-            structlog.processors.TimeStamper(fmt="iso", utc=False),
-            structlog.processors.StackInfoRenderer(),
-            structlog.processors.ExceptionRenderer(
-                structlog.tracebacks.ExceptionDictTransformer(show_locals=CONFIG.get_bool("debug"))
-            ),
-            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
-        ],
+        processors=processors,
         logger_factory=structlog.stdlib.LoggerFactory(),
         wrapper_class=structlog.make_filtering_bound_logger(
             getattr(logging, get_log_level(), logging.WARNING)

authentik/suse/worker/base_broker.py

@@ -74,7 +74,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
 # Django representation of pg_try_advisory_lock
 class PGTryLock(Func):
     arity = 1
-    function = "pg_try_advisory_lock"
+    function = "pg_try_advisory_xact_lock"
     output_field = BooleanField()
     conditional = True
 
@@ -290,7 +290,6 @@ def __init__(
         self.db_alias = db_alias
         self.queue_name = queue_name
         self.timeout = timeout // 1000
-        self.to_unlock: set[str] = set()
         self.in_processing: set[str] = set()
         self.prefetch = prefetch
         self.misses = 0
@@ -437,8 +436,7 @@ def _consume_one(self, message_id: str) -> Message[Any] | None:
             cursor.execute(sql, params)
 
             if cursor.rowcount != 1:
-                # Lock was not successful, mark it for unlock on next iteration
-                self._unlock_message(message_id)
+                # Lock was not successful, other consumer stole it from us
                 return None
 
         # Go on normally
@@ -461,7 +459,6 @@ def __next__(self) -> MessageProxy | None:
 
         # Run required processes first
         self._scheduler()
-        self._purge_locks()
 
         while True:
             # Try getting a message_id out of the in-flight set
@@ -513,26 +510,12 @@ def __next__(self) -> MessageProxy | None:
                     return None
                 # else: retry above, self.pending is guaranteed to have content
 
-    def _unlock_message(self, message_id: str) -> bool:
-        self.logger.debug("Unlocking message", message_id=message_id)
-        try:
-            with self.locks_connection.cursor() as cursor:
-                cursor.execute(
-                    "SELECT pg_advisory_unlock(%s)",
-                    (self._get_message_lock_id(message_id),),
-                )
-            return True
-        except DATABASE_ERRORS:
-            self.to_unlock.add(str(message_id))
-            return False
-
     def _post_process_message(self, message: Message[Any], state: TaskState) -> None:
         self.logger.debug("Post-processing message", message=message.message_id, state=state)
         self.logger.debug(
             "Removing from the in-memory queue", message=message.message_id, state=state
         )
         self.discard_in_processing(str(message.message_id))
-        self.to_unlock.add(str(message.message_id))
 
         self.logger.debug("Marking it in the DB", message=message.message_id, state=state)
         self.query_set.filter(
@@ -560,7 +543,6 @@ def requeue(self, messages: Iterable[Message[Any]]) -> None:
             state=TaskState.QUEUED,
         )
         for message in messages:
-            self.to_unlock.add(str(message.message_id))
             self.discard_in_processing(str(message.message_id))
 
     def _scheduler(self) -> None:
@@ -571,15 +553,6 @@ def _scheduler(self) -> None:
         self.scheduler.run()
         self.scheduler_last_run = timezone.now()
 
-    def _purge_locks(self) -> None:
-        while True:
-            try:
-                message_id = self.to_unlock.pop()
-            except KeyError:
-                break
-            if not self._unlock_message(str(message_id)):
-                return
-
     def _auto_purge(self) -> None:
         if timezone.now() - self.task_purge_last_run < self.task_purge_interval:
             return
@@ -600,16 +573,6 @@ def _auto_purge(self) -> None:
     @raise_connection_error
     def close(self) -> None:
         try:
-            self._purge_locks()
-        finally:
-            if self._locks_connection is not None:
-                conn = self._locks_connection
-                self._locks_connection = None
-                try:
-                    conn.close()
-                except DATABASE_ERRORS:
-                    pass
-            try:
-                connections.close_all()
-            except DATABASE_ERRORS:
-                pass
+            connections.close_all()
+        except DATABASE_ERRORS:
+            pass