Implement multiprocess support in the Redis mode

Signed-off-by: Stefano Rivera <stefano@rivera.za.net>

Author: stefanor

prometheus_client/redis.py

@@ -0,0 +1,90 @@
+import os
+from datetime import timedelta
+from threading import Event, Thread
+from typing import Any
+from urllib.parse import urlsplit
+
+from redis import Redis
+
+# For testing, a pool of otherwise anonymous FakeRedis instances are made
+# available by ID
+_fake_redis_pool: dict[int, Redis] = {}
+
+
+def redis_client() -> Redis:
+    """
+    Create a redis client for PROMETHEUS_REDIS_URL.
+
+    Configure the redis database via a URL in PROMETHEUS_REDIS_URL of the form
+    redis://localhost:6379/0
+    """
+    parsed_url = urlsplit(os.environ["PROMETHEUS_REDIS_URL"])
+    assert parsed_url.path.startswith("/")
+    assert parsed_url.path[1:].isdigit()
+    port = parsed_url.port or 6379
+    db = int(parsed_url.path[1:])
+
+    if parsed_url.scheme == "fakeredis":
+        from fakeredis import FakeRedis
+
+        if db not in _fake_redis_pool:
+            _fake_redis_pool[db] = FakeRedis()
+        return _fake_redis_pool[db]
+
+    assert parsed_url.scheme == "redis"
+    assert parsed_url.hostname
+    return Redis(host=parsed_url.hostname, port=port, db=db)
+
+
+# For each process identifier, a list of keys that should be kept from expiring
+_live_metrics: dict[str, set[str]] = {}
+
+
+def _key_expiry() -> timedelta:
+    """Return the configured expiry for multiprocess keys."""
+    return timedelta(seconds=int(os.environ.get("PROMETHEUS_REDIS_REFRESH_TTL", 20)))
+
+
+class KeepMetricsAliveThread(Thread):
+    """A daemon thread that keeps metrics from expiring as long as we live."""
+
+    stop: Event
+    identifier: str
+
+    def __init__(self, identifier: str, *args: Any, **kwargs: Any) -> None:
+        self.stop = Event()
+        self.identifier = identifier
+        super().__init__(*args, **kwargs)
+
+    def run(self) -> None:
+        delay = int(os.environ.get("PROMETHEUS_REDIS_REFRESH_FREQUENCY", 10))
+        expiry = _key_expiry()
+        client = redis_client()
+        while not self.stop.wait(delay):
+            for key in _live_metrics[self.identifier]:
+                client.expire(key, expiry)
+
+
+_daemon_threads: dict[str, KeepMetricsAliveThread] = {}
+
+
+def _keep_key_from_expiring(identifier: str, key: str) -> None:
+    """Stop key for process identifier from expiring as long as we are alive."""
+    _live_metrics.setdefault(identifier, set()).add(key)
+    if identifier not in _daemon_threads:
+        thread = KeepMetricsAliveThread(identifier=identifier, daemon=True)
+        thread.start()
+        _daemon_threads[identifier] = thread
+
+
+def mark_process_dead(identifier: str | int) -> None:
+    """Immediately expire all live* metrics for process identifier."""
+    thread = _daemon_threads.pop(str(identifier), None)
+    if thread is not None:
+        thread.stop.set()
+        thread.join()
+
+    keys = _live_metrics.pop(str(identifier), None)
+    if not keys:
+        return
+    redis_client().delete(*keys)

prometheus_client/redis_collector.py

@@ -1,39 +1,12 @@
-from collections.abc import Iterable
 import json
-import os
-from urllib.parse import urlsplit
+from collections.abc import Iterable
+from typing import cast
 
 from .metrics_core import Metric
+from .redis import redis_client
 from .registry import Collector, CollectorRegistry
 from .samples import Sample
-
-fake_redis_pool = {}
-
-
-def redis_client():
-    """
-    Create a redis client for PROMETHEUS_REDIS_URL.
-
-    Configure the redis database via a URL in PROMETHEUS_REDIS_URL of the form
-    redis://localhost:6379/0
-    """
-    from redis import Redis
-
-    parsed_url = urlsplit(os.environ["PROMETHEUS_REDIS_URL"])
-    assert parsed_url.path.startswith("/")
-    assert parsed_url.path[1:].isdigit()
-    port = parsed_url.port or 6379
-    db = int(parsed_url.path[1:])
-
-    if parsed_url.scheme == "fakeredis":
-        from fakeredis import FakeRedis
-
-        if db not in fake_redis_pool:
-            fake_redis_pool[db] = FakeRedis()
-        return fake_redis_pool[db]
-
-    assert parsed_url.scheme == "redis"
-    return Redis(host=parsed_url.hostname, port=port, db=db)
+from .values import MULTIPROCESS_MODE_T
 
 
 class RedisCollector(Collector):
@@ -56,30 +29,78 @@ def _iter_values(self) -> Iterable[tuple[bytes, str]]:
     def collect(self) -> Iterable[Metric]:
         metrics: dict[str, Metric] = {}
         histograms: set[str] = set()
+        multiprocess: dict[str, MULTIPROCESS_MODE_T] = {}
 
         for key, value_s in self._iter_values():
             # FIXME: Catch ValueError here, just in case?
-            prefix_b, typ_b, mmap_key = key.split(b":", 2)
+            prefix_b, typ_b, multiprocess_mode_b, mmap_key = key.split(b":", 3)
             assert prefix_b == b"value"
-            typ = typ_b.decode()
             value = float(value_s)
 
             metric_name, name, labels, help_text = json.loads(mmap_key)
 
             metric = metrics.get(metric_name)
             if metric is None:
+                typ = typ_b.decode()
                 metric = Metric(metric_name, help_text, typ)
                 metrics[metric_name] = metric
+
                 if typ in ("histogram", "gaugehistogram"):
                     histograms.add(metric_name)
 
+                multiprocess_mode = cast(
+                    MULTIPROCESS_MODE_T, multiprocess_mode_b.decode()
+                )
+                if typ in ("gauge", "gaugehistogram") and multiprocess_mode:
+                    multiprocess[metric_name] = multiprocess_mode
+
             metric.add_sample(name, labels, value)
 
+        for name, multiprocess_mode in multiprocess.items():
+            self._accumulate_multiprocess(metrics[name], multiprocess_mode)
+
         for name in histograms:
             self._fix_histogram(metrics[name])
 
         return metrics.values()
 
+    def _accumulate_multiprocess(
+        self, metric: Metric, multiprocess_mode: MULTIPROCESS_MODE_T
+    ) -> None:
+        """Merge metrics from multiple processes using multiprocess_mode."""
+        # We deal with live/dead with Redis expiry
+        if multiprocess_mode.startswith("live"):
+            multiprocess_mode = cast(
+                MULTIPROCESS_MODE_T, multiprocess_mode[len("live") :]
+            )
+        if multiprocess_mode == "all":
+            return
+
+        by_label: dict[tuple[tuple[str, ...], str], Sample] = {}
+
+        for sample in metric.samples:
+            labels = sample.labels
+            if "pid" in sample.labels:
+                labels = labels.copy()
+                labels.pop("pid")
+            key = (tuple(labels.values()), sample.name)
+            value = sample.value
+            if key in by_label:
+                current_value = by_label[key].value
+                if multiprocess_mode == "min" and value > current_value:
+                    continue
+                if multiprocess_mode == "max" and value < current_value:
+                    continue
+                if multiprocess_mode == "sum":
+                    value += current_value
+                if multiprocess_mode == "mostrecent":
+                    raise NotImplementedError(
+                        "The 'mostrecent' modes are not supported in RedisCollector"
+                    )
+            by_label[key] = Sample(sample.name, labels, value)
+
+        metric.samples = list(by_label.values())
+
     def _fix_histogram(self, metric: Metric) -> None:
         """
         Fix-up histogram samples.