Refactor more custom Collector to be stubbed to use hs.metrics_collector_registry

synapse/metrics/_gc.py

@@ -24,10 +24,11 @@ import gc
 import logging
 import platform
 import time
-from typing import Iterable
+from typing import Iterable, Optional
 
 from prometheus_client.core import (
     REGISTRY,
+    CollectorRegistry,
     CounterMetricFamily,
     Gauge,
     GaugeMetricFamily,
@@ -81,6 +82,10 @@ gc_time = Histogram(
 
 
 class GCCounts(Collector):
+    def __init__(self, registry: Optional[CollectorRegistry] = REGISTRY) -> None:
+        if registry is not None:
+            registry.register(self)
+
     def collect(self) -> Iterable[Metric]:
         cm = GaugeMetricFamily("python_gc_counts", "GC object counts", labels=["gen"])
         for n, m in enumerate(gc.get_count()):
@@ -101,7 +106,7 @@ def install_gc_manager() -> None:
     if running_on_pypy:
         return
 
-    REGISTRY.register(GCCounts())
+    GCCounts(registry=hs.metrics_collector_registry)
 
     gc.disable()
 
@@ -145,6 +150,10 @@ def install_gc_manager() -> None:
 
 
 class PyPyGCStats(Collector):
+    def __init__(self, registry: Optional[CollectorRegistry] = REGISTRY) -> None:
+        if registry is not None:
+            registry.register(self)
+
     def collect(self) -> Iterable[Metric]:
         # @stats is a pretty-printer object with __str__() returning a nice table,
         # plus some fields that contain data from that table.
@@ -208,4 +217,4 @@ class PyPyGCStats(Collector):
 
 
 if running_on_pypy:
-    REGISTRY.register(PyPyGCStats())
+    PyPyGCStats(registry=hs.metrics_collector_registry)

synapse/metrics/_reactor_metrics.py

@@ -22,10 +22,10 @@
 import logging
 import time
 from selectors import SelectSelector, _PollLikeSelector  # type: ignore[attr-defined]
-from typing import Any, Callable, Iterable
+from typing import Any, Callable, Iterable, Optional
 
 from prometheus_client import Histogram, Metric
-from prometheus_client.core import REGISTRY, GaugeMetricFamily
+from prometheus_client.core import REGISTRY, CollectorRegistry, GaugeMetricFamily
 
 from twisted.internet import reactor, selectreactor
 from twisted.internet.asyncioreactor import AsyncioSelectorReactor
@@ -110,9 +110,16 @@ class ObjWrapper:
 
 
 class ReactorLastSeenMetric(Collector):
-    def __init__(self, call_wrapper: CallWrapper):
+    def __init__(
+        self,
+        call_wrapper: CallWrapper,
+        registry: Optional[CollectorRegistry] = REGISTRY,
+    ):
         self._call_wrapper = call_wrapper
 
+        if registry is not None:
+            registry.register(self)
+
     def collect(self) -> Iterable[Metric]:
         cm = GaugeMetricFamily(
             "python_twisted_reactor_last_seen",
@@ -165,4 +172,4 @@ except Exception as e:
 
 
 if wrapper:
-    REGISTRY.register(ReactorLastSeenMetric(wrapper))
+    ReactorLastSeenMetric(wrapper, registry=hs.metrics_collector_registry)

synapse/metrics/background_process_metrics.py

@@ -38,7 +38,7 @@ from typing import (
 )
 
 from prometheus_client import Metric
-from prometheus_client.core import REGISTRY, Counter, Gauge
+from prometheus_client.core import REGISTRY, CollectorRegistry, Counter, Gauge
 from typing_extensions import ParamSpec
 
 from twisted.internet import defer
@@ -141,6 +141,10 @@ class _Collector(Collector):
     before they are returned.
     """
 
+    def __init__(self, registry: Optional[CollectorRegistry] = REGISTRY) -> None:
+        if registry is not None:
+            registry.register(self)
+
     def collect(self) -> Iterable[Metric]:
         global _background_processes_active_since_last_scrape
 
@@ -165,7 +169,7 @@ class _Collector(Collector):
             yield from m.collect()
 
 
-REGISTRY.register(_Collector())
+_Collector(registry=hs.metrics_collector_registry)
 
 
 class _BackgroundProcess:

synapse/metrics/jemalloc.py

@@ -26,7 +26,7 @@ import re
 from typing import Iterable, Literal, Optional, overload
 
 import attr
-from prometheus_client import REGISTRY, Metric
+from prometheus_client import REGISTRY, CollectorRegistry, Metric
 
 from synapse.metrics import GaugeMetricFamily
 from synapse.metrics._types import Collector
@@ -185,6 +185,10 @@ def _setup_jemalloc_stats() -> None:
     class JemallocCollector(Collector):
         """Metrics for internal jemalloc stats."""
 
+        def __init__(self, registry: Optional[CollectorRegistry] = REGISTRY) -> None:
+            if registry is not None:
+                registry.register(self)
+
         def collect(self) -> Iterable[Metric]:
             stats.refresh_stats()
 
@@ -230,7 +234,7 @@ def _setup_jemalloc_stats() -> None:
 
             yield g
 
-    REGISTRY.register(JemallocCollector())
+    JemallocCollector(registry=hs.metrics_collector_registry)
 
     logger.debug("Added jemalloc stats")
 

synapse/util/caches/__init__.py

@@ -27,7 +27,6 @@ from sys import intern
 from typing import Any, Callable, Dict, List, Optional, Sized, TypeVar
 
 import attr
-from prometheus_client import REGISTRY
 from prometheus_client.core import Gauge
 
 from synapse.config.cache import add_resizable_cache
@@ -94,7 +93,7 @@ response_cache_total = Gauge(
 
 
 # Register our custom cache metrics registry with the global registry
-REGISTRY.register(CACHE_METRIC_REGISTRY)
+hs.metrics_collector_registry.register(CACHE_METRIC_REGISTRY)
 
 
 class EvictionReason(Enum):

tests/storage/test_event_metrics.py

@@ -20,7 +20,6 @@
 #
 from prometheus_client import generate_latest
 
-from synapse.metrics import REGISTRY
 from synapse.types import UserID, create_requester
 
 from tests.unittest import HomeserverTestCase
@@ -61,7 +60,7 @@ class ExtremStatisticsTestCase(HomeserverTestCase):
         items = list(
             filter(
                 lambda x: b"synapse_forward_extremities_" in x and b"# HELP" not in x,
-                generate_latest(REGISTRY).split(b"\n"),
+                generate_latest(self.hs.metrics_collector_registry).split(b"\n"),
             )
         )