Prometheus: fix offline worker metrics

flower/api/prometheus_metrics.py

@@ -0,0 +1,77 @@
+import logging
+from typing import List, Tuple, Set
+
+from prometheus_client import Counter as PrometheusCounter, Histogram, Gauge
+from prometheus_client.metrics import MetricWrapperBase
+
+logger = logging.getLogger(__name__)
+
+
+class LabelNames:
+    WORKER = 'worker'
+    TYPE = 'type'
+    TASK = 'task'
+
+
+class PrometheusMetrics(object):
+    events = PrometheusCounter(
+        'flower_events_total', "Number of events", [LabelNames.WORKER, LabelNames.TYPE, LabelNames.TASK]
+    )
+    runtime = Histogram('flower_task_runtime_seconds', "Task runtime", [LabelNames.WORKER, LabelNames.TASK])
+    prefetch_time = Gauge(
+        'flower_task_prefetch_time_seconds',
+        "The time the task spent waiting at the celery worker to be executed.",
+        [LabelNames.WORKER, LabelNames.TASK]
+    )
+    number_of_prefetched_tasks = Gauge(
+        'flower_worker_prefetched_tasks',
+        'Number of tasks of given type prefetched at a worker',
+        [LabelNames.WORKER, LabelNames.TASK]
+    )
+    worker_online = Gauge('flower_worker_online', "Worker online status", [LabelNames.WORKER])
+    worker_number_of_currently_executing_tasks = Gauge(
+        'flower_worker_number_of_currently_executing_tasks',
+        "Number of tasks currently executing at a worker",
+        [LabelNames.WORKER]
+    )
+
+    @property
+    def transient_metrics(self) -> List[MetricWrapperBase]:
+        return [
+            self.events,
+            self.runtime,
+            self.prefetch_time,
+            self.number_of_prefetched_tasks,
+            self.worker_online,
+            self.worker_number_of_currently_executing_tasks
+        ]
+
+    def remove_metrics_for_offline_workers(self, offline_workers: Set[str]):
+        for metric in self.transient_metrics:
+            labels_sets_for_offline_workers = self._get_label_sets_for_offline_workers(
+                metric=metric, offline_workers=offline_workers
+            )
+            for label_set in labels_sets_for_offline_workers:
+                try:
+                    metric.remove(*label_set)
+                    logger.debug('Removed label set: %s for metric %s', label_set, metric)
+                except KeyError:
+                    pass
+
+    @staticmethod
+    def _get_label_sets_for_offline_workers(
+            metric: MetricWrapperBase, offline_workers: Set[str]
+    ) -> Set[Tuple[str, ...]]:
+        sampled_metrics = metric.collect()
+
+        label_sets_for_offline_workers = set()
+        for sampled_metric in sampled_metrics:
+            for sample in sampled_metric.samples:
+                labels = sample.labels
+                worker = labels.get(LabelNames.WORKER)
+                if worker is None or worker not in offline_workers:
+                    continue
+
+                label_sets_for_offline_workers.add(tuple([labels[label_name] for label_name in metric._labelnames]))
+
+        return label_sets_for_offline_workers

flower/events.py

@@ -5,6 +5,7 @@
 import collections
 
 from functools import partial
+from typing import Dict, Any, Set
 
 from tornado.ioloop import IOLoop
 from tornado.ioloop import PeriodicCallback
@@ -13,35 +14,14 @@
 from celery.events.state import State
 
 from . import api
+from .api.prometheus_metrics import PrometheusMetrics
+from .options import options
 
 from collections import Counter
 
-from prometheus_client import Counter as PrometheusCounter, Histogram, Gauge
-
 logger = logging.getLogger(__name__)
 
 
-class PrometheusMetrics(object):
-    events = PrometheusCounter('flower_events_total', "Number of events", ['worker', 'type', 'task'])
-    runtime = Histogram('flower_task_runtime_seconds', "Task runtime", ['worker', 'task'])
-    prefetch_time = Gauge(
-        'flower_task_prefetch_time_seconds',
-        "The time the task spent waiting at the celery worker to be executed.",
-        ['worker', 'task']
-    )
-    number_of_prefetched_tasks = Gauge(
-        'flower_worker_prefetched_tasks',
-        'Number of tasks of given type prefetched at a worker',
-        ['worker', 'task']
-    )
-    worker_online = Gauge('flower_worker_online', "Worker online status", ['worker'])
-    worker_number_of_currently_executing_tasks = Gauge(
-        'flower_worker_number_of_currently_executing_tasks',
-        "Number of tasks currently executing at a worker",
-        ['worker']
-    )
-
-
 class EventsState(State):
     # EventsState object is created and accessed only from ioloop thread
 
@@ -103,6 +83,84 @@ def event(self, event):
         if cls:
             cls.send_message(event)
 
+    def remove_metrics_for_offline_workers(self):
+        if options.purge_offline_workers is not None:
+            offline_workers = self.get_offline_workers(workers=self.get_workers())
+            if not offline_workers:
+                return
+
+            self.metrics.remove_metrics_for_offline_workers(
+                offline_workers=offline_workers
+            )
+
+    def get_online_workers(self) -> Dict[str, Any]:
+        workers = self.get_workers()
+        if options.purge_offline_workers is not None:
+            for name in self.get_offline_workers(workers=workers):
+                workers.pop(name)
+
+        return workers
+
+    def get_workers(self) -> Dict[str, Any]:
+        workers = {}
+        for name, values in self.counter.items():
+            if name not in self.workers:
+                continue
+            worker = self.workers[name]
+            info = dict(values)
+            info.update(self._as_dict(worker))
+            info.update(status=worker.alive)
+            workers[name] = info
+
+        return workers
+
+    @staticmethod
+    def get_offline_workers(workers: Dict[str, Any]) -> Set[str]:
+        timestamp = int(time.time())
+        offline_workers = set()
+        for name, info in workers.items():
+            if info.get('status', True):
+                continue
+
+            heartbeats = info.get('heartbeats', [])
+            last_heartbeat = int(max(heartbeats)) if heartbeats else None
+            if not last_heartbeat or timestamp - last_heartbeat > options.purge_offline_workers:
+                offline_workers.add(name)
+                logger.debug('Found offline worker: %s', name)
+
+        return offline_workers
+
+    @classmethod
+    def _as_dict(cls, worker) -> Dict[str, Any]:
+        if hasattr(worker, '_fields'):
+            return dict((k, worker.__getattribute__(k)) for k in worker._fields)
+        else:
+            return cls._info(worker)
+
+    @classmethod
+    def _info(cls, worker) -> Dict[str, Any]:
+        _fields = (
+            'hostname',
+            'pid',
+            'freq',
+            'heartbeats',
+            'clock',
+            'active',
+            'processed',
+            'loadavg',
+            'sw_ident',
+            'sw_ver',
+            'sw_sys'
+        )
+
+        def _keys():
+            for key in _fields:
+                value = getattr(worker, key, None)
+                if value is not None:
+                    yield key, value
+
+        return dict(_keys())
+
 
 class Events(threading.Thread):
     events_enable_interval = 5000

flower/views/dashboard.py

@@ -1,12 +1,10 @@
 import logging
-import time
 
 from tornado import web
 from tornado import gen
 
+from ..events import EventsState
 from ..views import BaseHandler
-from ..options import options
-from ..api.workers import ListWorkers
 
 
 logger = logging.getLogger(__name__)
@@ -19,64 +17,19 @@ def get(self):
         refresh = self.get_argument('refresh', default=False, type=bool)
         json = self.get_argument('json', default=False, type=bool)
 
-        events = self.application.events.state
+        events_state: EventsState = self.application.events.state
 
         if refresh:
             try:
                 self.application.update_workers()
             except Exception as e:
                 logger.exception('Failed to update workers: %s', e)
 
-        workers = {}
-        for name, values in events.counter.items():
-            if name not in events.workers:
-                continue
-            worker = events.workers[name]
-            info = dict(values)
-            info.update(self._as_dict(worker))
-            info.update(status=worker.alive)
-            workers[name] = info
-
-        if options.purge_offline_workers is not None:
-            timestamp = int(time.time())
-            offline_workers = []
-            for name, info in workers.items():
-                if info.get('status', True):
-                    continue
-
-                heartbeats = info.get('heartbeats', [])
-                last_heartbeat = int(max(heartbeats)) if heartbeats else None
-                if not last_heartbeat or timestamp - last_heartbeat > options.purge_offline_workers:
-                    offline_workers.append(name)
-
-            for name in offline_workers:
-                workers.pop(name)
-
+        workers = events_state.get_online_workers()
         if json:
             self.write(dict(data=list(workers.values())))
         else:
             self.render("dashboard.html",
                         workers=workers,
                         broker=self.application.capp.connection().as_uri(),
                         autorefresh=1 if self.application.options.auto_refresh else 0)
-
-    @classmethod
-    def _as_dict(cls, worker):
-        if hasattr(worker, '_fields'):
-            return dict((k, worker.__getattribute__(k)) for k in worker._fields)
-        else:
-            return cls._info(worker)
-
-    @classmethod
-    def _info(cls, worker):
-        _fields = ('hostname', 'pid', 'freq', 'heartbeats', 'clock',
-                   'active', 'processed', 'loadavg', 'sw_ident',
-                   'sw_ver', 'sw_sys')
-
-        def _keys():
-            for key in _fields:
-                value = getattr(worker, key, None)
-                if value is not None:
-                    yield key, value
-
-        return dict(_keys())

flower/views/monitor.py

@@ -2,12 +2,16 @@
 
 from tornado import gen
 
+from ..events import EventsState
 from ..views import BaseHandler
 
 
 class Metrics(BaseHandler):
     @gen.coroutine
     def get(self):
+        events_state: EventsState = self.application.events.state
+        events_state.remove_metrics_for_offline_workers()
+
         self.write(prometheus_client.generate_latest())
         self.set_header("Content-Type", "text/plain")
 

tests/unit/views/test_dashboard.py

@@ -73,7 +73,7 @@ def test_purge_offline_workers(self):
                           local_received=time.time()))
         self.app.events.state = state
 
-        with patch('flower.views.dashboard.options') as mock_options:
+        with patch('flower.events.options') as mock_options:
             mock_options.purge_offline_workers = 0
             r = self.get('/dashboard')