From 7a943eac31b97a4fcb7f98b1caeb5f80e5a798b4 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens@goauthentik.io>
Date: Sat, 13 Jan 2024 22:59:55 +0100
Subject: [PATCH] improve metrics

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
---
 authentik/events/apps.py    | 13 ++++++++++++-
 authentik/events/models.py  | 13 ++++++++++---
 authentik/events/signals.py |  2 ++
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/authentik/events/apps.py b/authentik/events/apps.py
index cd79b54bd..95cc65ab5 100644
--- a/authentik/events/apps.py
+++ b/authentik/events/apps.py
@@ -1,15 +1,26 @@
 """authentik events app"""
-from prometheus_client import Gauge
+from prometheus_client import Gauge, Histogram
 
 from authentik.blueprints.apps import ManagedAppConfig
 from authentik.lib.config import CONFIG, ENV_PREFIX
 
+# TODO: Deprecated metric - remove in 2024.2 or later
 GAUGE_TASKS = Gauge(
     "authentik_system_tasks",
     "System tasks and their status",
     ["task_name", "task_uid", "status"],
 )
 
+SYSTEM_TASK_TIME = Histogram(
+    "authentik_system_tasks_time_seconds",
+    "Runtime of system tasks",
+)
+SYSTEM_TASK_STATUS = Gauge(
+    "authentik_system_tasks_status",
+    "System task status",
+    ["task_name", "task_uid", "status"],
+)
+
 
 class AuthentikEventsConfig(ManagedAppConfig):
     """authentik events app"""
diff --git a/authentik/events/models.py b/authentik/events/models.py
index 91d36c5a0..b50db8162 100644
--- a/authentik/events/models.py
+++ b/authentik/events/models.py
@@ -27,7 +27,7 @@ from authentik.core.middleware import (
     SESSION_KEY_IMPERSONATE_USER,
 )
 from authentik.core.models import ExpiringModel, Group, PropertyMapping, User
-from authentik.events.apps import GAUGE_TASKS
+from authentik.events.apps import GAUGE_TASKS, SYSTEM_TASK_STATUS, SYSTEM_TASK_TIME
 from authentik.events.context_processors.base import get_context_processors
 from authentik.events.utils import (
     cleanse_dict,
@@ -622,12 +622,19 @@ class SystemTask(SerializerModel, ExpiringModel):
 
     def update_metrics(self):
         """Update prometheus metrics"""
-        duration = max(self.finish_timestamp.timestamp() - self.start_timestamp.timestamp(), 0)
+        duration = max(self.finish_timestamp - self.start_timestamp, 0)
+        # TODO: Deprecated metric - remove in 2024.2 or later
         GAUGE_TASKS.labels(
             task_name=self.name.split(":")[0],
             task_uid=self.uid or "",
-            status=self.status.name.lower(),
+            status=self.status.lower(),
         ).set(duration)
+        SYSTEM_TASK_TIME.observe(duration)
+        SYSTEM_TASK_STATUS.labels(
+            task_name=self.name.split(":")[0],
+            task_uid=self.uid or "",
+            status=self.status.lower(),
+        ).inc()
 
     def __str__(self) -> str:
         return f"System Task {self.name}"
diff --git a/authentik/events/signals.py b/authentik/events/signals.py
index dcec5eb42..bc5a38556 100644
--- a/authentik/events/signals.py
+++ b/authentik/events/signals.py
@@ -8,6 +8,7 @@ from django.http import HttpRequest
 
 from authentik.core.models import User
 from authentik.core.signals import login_failed, password_changed
+from authentik.events.apps import SYSTEM_TASK_STATUS
 from authentik.events.models import Event, EventAction, SystemTask
 from authentik.events.tasks import event_notification_handler, gdpr_cleanup
 from authentik.flows.models import Stage
@@ -106,5 +107,6 @@ def event_user_pre_delete_cleanup(sender, instance: User, **_):
 @receiver(monitoring_set)
 def monitoring_system_task(sender, **_):
     """Update metrics when task is saved"""
+    SYSTEM_TASK_STATUS.clear()
     for task in SystemTask.objects.all():
         task.update_metrics()