add bucket options

.gitignore

@@ -15,3 +15,4 @@ __pycache__
 htmlcov/
 dive.log
 .vscode
+venv

celery_exporter/__main__.py

@@ -7,10 +7,12 @@
 
 import click
 from .core import CeleryExporter
+from prometheus_client import Histogram
 
 __VERSION__ = (2, 0, 0)
 
 LOG_FORMAT = "[%(asctime)s] %(name)s:%(levelname)s: %(message)s"
+DEFAULT_BUCKETS = ",".join(map(str, Histogram.DEFAULT_BUCKETS))
 
 
 @click.command(context_settings={"auto_envvar_prefix": "CELERY_EXPORTER"})
@@ -62,6 +64,22 @@
     allow_from_autoenv=False,
     help="Periodically enable Celery events.",
 )
+@click.option(
+    "--runtime-histogram-bucket",
+    type=str,
+    show_default=True,
+    show_envvar=True,
+    default=DEFAULT_BUCKETS,
+    help="Default buckets of runtime historgram in seconds",
+)
+@click.option(
+    "--latency-histogram-bucket",
+    type=str,
+    show_default=True,
+    show_envvar=True,
+    default=DEFAULT_BUCKETS,
+    help="Default buckets of latency historgram in seconds.",
+)
 @click.option(
     "--tz", type=str, allow_from_autoenv=False, help="Timezone used by the celery app."
 )
@@ -76,6 +94,8 @@ def main(
     namespace,
     transport_options,
     enable_events,
+    runtime_histogram_bucket,
+    latency_histogram_bucket,
     tz,
     verbose,
 ):  # pragma: no cover
@@ -101,13 +121,40 @@ def main(
             )
             sys.exit(1)
 
+    def decode_buckets(buckets_list):
+        return tuple(map(float, buckets_list.split(",")))
+
+    try:
+        runtime_histogram_bucket = decode_buckets(runtime_histogram_bucket)
+    except ValueError:
+        print(
+            "Error parsing runtime_histogram_bucket options '{}' must be a comma sperated list of numbers".format(
+                runtime_histogram_bucket
+            ),
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    try:
+        latency_histogram_bucket = decode_buckets(latency_histogram_bucket)
+    except ValueError:
+        print(
+            "Error parsing latency_histogram_bucket options '{}' must be a comma sperated list of numbers".format(
+                latency_histogram_bucket
+            ),
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
     celery_exporter = CeleryExporter(
         broker_url,
         listen_address,
         max_tasks,
         namespace,
         transport_options,
         enable_events,
+        runtime_histogram_bucket,
+        latency_histogram_bucket,
     )
     celery_exporter.start()
 

celery_exporter/core.py

@@ -20,25 +20,36 @@ def __init__(
         namespace="celery",
         transport_options=None,
         enable_events=False,
+        runtime_histogram_bucket=prometheus_client.Histogram.DEFAULT_BUCKETS,
+        latency_histogram_bucket=prometheus_client.Histogram.DEFAULT_BUCKETS,
     ):
         self._listen_address = listen_address
         self._max_tasks = max_tasks
         self._namespace = namespace
         self._enable_events = enable_events
+        self._runtime_histogram_bucket = runtime_histogram_bucket
+        self._latency_histogram_bucket = latency_histogram_bucket
 
         self._app = celery.Celery(broker=broker_url)
         self._app.conf.broker_transport_options = transport_options or {}
 
     def start(self):
 
-        setup_metrics(self._app, self._namespace)
+        setup_metrics(
+            self._app,
+            self._namespace,
+            self._runtime_histogram_bucket,
+            self._latency_histogram_bucket,
+        )
 
         self._start_httpd()
 
         t = TaskThread(
             app=self._app,
             namespace=self._namespace,
             max_tasks_in_memory=self._max_tasks,
+            runtime_histogram_bucket=self._runtime_histogram_bucket,
+            latency_histogram_bucket=self._latency_histogram_bucket,
         )
         t.daemon = True
         t.start()

celery_exporter/monitor.py

@@ -6,23 +6,39 @@
 
 import celery
 import celery.states
+import prometheus_client
 
-from .metrics import TASKS, TASKS_RUNTIME, LATENCY, WORKERS
 from celery_state import CeleryState
 from .utils import get_config
 
+TASKS = None
+WORKERS = None
+TASKS_RUNTIME = None
+LATENCY = None
+
 
 class TaskThread(threading.Thread):
     """
     MonitorThread is the thread that will collect the data that is later
     exposed from Celery using its eventing system.
     """
 
-    def __init__(self, app, namespace, max_tasks_in_memory, *args, **kwargs):
+    def __init__(
+        self,
+        app,
+        namespace,
+        max_tasks_in_memory,
+        runtime_histogram_bucket,
+        latency_histogram_bucket,
+        *args,
+        **kwargs
+    ):
         self._app = app
         self._namespace = namespace
         self.log = logging.getLogger("task-thread")
         self._state = CeleryState(max_tasks_in_memory=max_tasks_in_memory)
+        self._runtime_histogram_bucket = runtime_histogram_bucket
+        self._latency_histogram_bucket = latency_histogram_bucket
         self._known_states = set()
         self._known_states_names = set()
         self._tasks_started = dict()
@@ -56,12 +72,22 @@ def _monitor(self):  # pragma: no cover
                     recv = self._app.events.Receiver(
                         conn, handlers={"*": self._process_event}
                     )
-                    setup_metrics(self._app, self._namespace)
+                    setup_metrics(
+                        self._app,
+                        self._namespace,
+                        self._runtime_histogram_bucket,
+                        self._latency_histogram_bucket,
+                    )
                     self.log.info("Start capturing events...")
                     recv.capture(limit=None, timeout=None, wakeup=True)
             except Exception:
                 self.log.exception("Connection failed")
-                setup_metrics(self._app, self._namespace)
+                setup_metrics(
+                    self._app,
+                    self._namespace,
+                    self._runtime_histogram_bucket,
+                    self._latency_histogram_bucket,
+                )
                 time.sleep(5)
 
 
@@ -109,11 +135,41 @@ def enable_events(self):
         self._app.control.enable_events()
 
 
-def setup_metrics(app, namespace):
+def setup_metrics(app, namespace, task_buckets, latency_buckets):
     """
     This initializes the available metrics with default values so that
     even before the first event is received, data can be exposed.
     """
+    global TASKS
+    global TASKS_RUNTIME
+    global LATENCY
+    global WORKERS
+
+    if TASKS == None:
+        TASKS = prometheus_client.Counter(
+            "celery_tasks_total",
+            "Number of task events.",
+            ["namespace", "name", "state", "queue"],
+        )
+    if TASKS_RUNTIME == None:
+        TASKS_RUNTIME = prometheus_client.Histogram(
+            "celery_tasks_runtime_seconds",
+            "Task runtime.",
+            ["namespace", "name", "queue"],
+            buckets=task_buckets,
+        )
+    if LATENCY == None:
+        LATENCY = prometheus_client.Histogram(
+            "celery_tasks_latency_seconds",
+            "Time between a task is received and started.",
+            ["namespace", "name", "queue"],
+            buckets=latency_buckets,
+        )
+    if WORKERS == None:
+        WORKERS = prometheus_client.Gauge(
+            "celery_workers", "Number of alive workers", ["namespace"]
+        )
+
     WORKERS.labels(namespace=namespace)
     config = get_config(app)
 

test/test_func.py

@@ -4,6 +4,7 @@
 import celery
 import celery_exporter.monitor
 from celery_exporter.core import CeleryExporter
+from prometheus_client import Histogram
 
 prom_http_server_mock = MagicMock(return_value=None)
 setup_metrics_mock = MagicMock(return_value=None)
@@ -33,7 +34,10 @@ def setUp(self):
     def test_setup_metrics(self):
         self.cel_exp.start()
         setup_metrics_mock.assert_called_with(
-            self.cel_exp._app, TestCeleryExporter.namespace
+            self.cel_exp._app,
+            TestCeleryExporter.namespace,
+            Histogram.DEFAULT_BUCKETS,
+            Histogram.DEFAULT_BUCKETS,
         )
 
     def test_http_server(self):
@@ -46,6 +50,8 @@ def test_task_thread(self):
             self.cel_exp._app,
             TestCeleryExporter.namespace,
             TestCeleryExporter.max_tasks,
+            runtime_histogram_bucket=Histogram.DEFAULT_BUCKETS,
+            latency_histogram_bucket=Histogram.DEFAULT_BUCKETS,
         )
 
     def test_worker_thread(self):

test/test_unit.py

@@ -5,7 +5,7 @@
 
 from celery.events import Event
 from celery.utils import uuid
-from prometheus_client import REGISTRY
+from prometheus_client import REGISTRY, Histogram
 from unittest import TestCase
 from unittest.mock import patch
 
@@ -36,7 +36,12 @@ def setUp(self):
                     "celery@12311847jsa2": {},
                 }
                 registered.return_value = {"celery@d6f95e9e24fc": [self.task, "trial"]}
-                setup_metrics(self.app, self.namespace)  # reset metrics
+                setup_metrics(
+                    self.app,
+                    self.namespace,
+                    task_buckets=Histogram.DEFAULT_BUCKETS,
+                    latency_buckets=Histogram.DEFAULT_BUCKETS,
+                )  # reset metrics
 
     def test_initial_metric_values(self):
         self._assert_task_states(celery.states.ALL_STATES, 0)
@@ -116,7 +121,11 @@ def test_tasks_events(self):
         runtime = 234.5
 
         m = TaskThread(
-            app=self.app, namespace=self.namespace, max_tasks_in_memory=self.max_tasks
+            app=self.app,
+            namespace=self.namespace,
+            max_tasks_in_memory=self.max_tasks,
+            runtime_histogram_bucket=Histogram.DEFAULT_BUCKETS,
+            latency_histogram_bucket=Histogram.DEFAULT_BUCKETS,
         )
 
         self._assert_task_states(celery.states.ALL_STATES, 0)