Move jemalloc metrics further along

synapse/app/_base.py

@@ -76,11 +76,11 @@ from synapse.logging.context import PreserveLoggingContext
 from synapse.logging.opentracing import init_tracer
 from synapse.metrics import CPUMetrics, install_gc_manager, register_threadpool
 from synapse.metrics._gc import GCCounts, PyPyGCStats, running_on_pypy
+from synapse.metrics._reactor_metrics import setup_reactor_metrics
 from synapse.metrics.background_process_metrics import (
     BackgroundProcessCollector,
     wrap_as_background_process,
 )
-from synapse.metrics.jemalloc import setup_jemalloc_stats
 from synapse.module_api.callbacks.spamchecker_callbacks import load_legacy_spam_checkers
 from synapse.module_api.callbacks.third_party_event_rules_callbacks import (
     load_legacy_third_party_event_rules,
@@ -695,7 +695,7 @@ def setup_global_metrics(hs: "HomeServer") -> None:
         PyPyGCStats(registry=hs.metrics_collector_registry)
         GCCounts(registry=hs.metrics_collector_registry)
     BackgroundProcessCollector(registry=hs.metrics_collector_registry)
-    setup_jemalloc_stats()
+    setup_reactor_metrics(hs)
 
 
 def setup_sdnotify(hs: "HomeServer") -> None:

synapse/metrics/_reactor_metrics.py

@@ -22,7 +22,7 @@
 import logging
 import time
 from selectors import SelectSelector, _PollLikeSelector  # type: ignore[attr-defined]
-from typing import Any, Callable, Iterable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional
 
 from prometheus_client import Histogram, Metric
 from prometheus_client.core import REGISTRY, CollectorRegistry, GaugeMetricFamily
@@ -32,6 +32,9 @@ from twisted.internet.asyncioreactor import AsyncioSelectorReactor
 
 from synapse.metrics._types import Collector
 
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
+
 try:
     from selectors import KqueueSelector
 except ImportError:
@@ -171,5 +174,6 @@ except Exception as e:
     logger.warning("Configuring ReactorLastSeenMetric failed: %r", e)
 
 
-if wrapper:
-    ReactorLastSeenMetric(wrapper, registry=hs.metrics_collector_registry)
+def setup_reactor_metrics(hs: "HomeServer") -> None:
+    if wrapper:
+        ReactorLastSeenMetric(wrapper, registry=hs.metrics_collector_registry)

synapse/metrics/jemalloc.py

@@ -23,14 +23,16 @@ import ctypes
 import logging
 import os
 import re
-from typing import Iterable, Literal, Optional, overload
+from typing import TYPE_CHECKING, Iterable, Literal, Optional, overload
 
 import attr
 from prometheus_client import REGISTRY, CollectorRegistry, Metric
 
 from synapse.metrics import GaugeMetricFamily
 from synapse.metrics._types import Collector
-from synapse.server import HomeServer
+
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
 
 logger = logging.getLogger(__name__)
 
@@ -132,26 +134,11 @@ class JemallocStats:
         return self._mallctl(f"stats.{name}")
 
 
-_JEMALLOC_STATS: Optional[JemallocStats] = None
-
-
-def get_jemalloc_stats() -> Optional[JemallocStats]:
-    """Returns an interface to jemalloc, if it is being used.
-
-    Note that this will always return None until `setup_jemalloc_stats` has been
-    called.
-    """
-    # TODO
-    return _JEMALLOC_STATS
-
-
-def _setup_jemalloc_stats(hs: HomeServer) -> None:
+def _setup_jemalloc_stats(hs: "HomeServer") -> Optional[JemallocStats]:
     """Checks to see if jemalloc is loaded, and hooks up a collector to record
     statistics exposed by jemalloc.
     """
 
-    global _JEMALLOC_STATS
-
     # Try to find the loaded jemalloc shared library, if any. We need to
     # introspect into what is loaded, rather than loading whatever is on the
     # path, as if we load a *different* jemalloc version things will seg fault.
@@ -182,7 +169,6 @@ def _setup_jemalloc_stats(hs: HomeServer) -> None:
     jemalloc_dll = ctypes.CDLL(jemalloc_path)
 
     stats = JemallocStats(jemalloc_dll)
-    _JEMALLOC_STATS = stats
 
     class JemallocCollector(Collector):
         """Metrics for internal jemalloc stats."""
@@ -239,13 +225,14 @@ def _setup_jemalloc_stats(hs: HomeServer) -> None:
     JemallocCollector(registry=hs.metrics_collector_registry)
 
     logger.debug("Added jemalloc stats")
+    return stats
 
 
-def setup_jemalloc_stats(hs: HomeServer) -> None:
+def setup_jemalloc_stats(hs: "HomeServer") -> Optional[JemallocStats]:
     """Try to setup jemalloc stats, if jemalloc is loaded."""
 
     try:
-        _setup_jemalloc_stats(hs)
+        return _setup_jemalloc_stats(hs)
     except Exception as e:
         # This should only happen if we find the loaded jemalloc library, but
         # fail to load it somehow (e.g. we somehow picked the wrong version).

synapse/server.py

@@ -128,6 +128,7 @@ from synapse.http.matrixfederationclient import MatrixFederationHttpClient
 from synapse.media.media_repository import MediaRepository
 from synapse.metrics import register_threadpool
 from synapse.metrics.common_usage_metrics import CommonUsageMetricsManager
+from synapse.metrics.jemalloc import JemallocStats, setup_jemalloc_stats
 from synapse.module_api import ModuleApi
 from synapse.module_api.callbacks import ModuleApiCallbacks
 from synapse.notifier import Notifier, ReplicationNotifier
@@ -310,6 +311,7 @@ class HomeServer(metaclass=abc.ABCMeta):
         self.tls_server_context_factory: Optional[IOpenSSLContextFactory] = None
 
         self.metrics_collector_registry = CollectorRegistry(auto_describe=True)
+        self.jemalloc_stats: Optional[JemallocStats] = None
 
     def register_module_web_resource(self, path: str, resource: Resource) -> None:
         """Allows a module to register a web resource to be served at the given path.
@@ -360,6 +362,7 @@ class HomeServer(metaclass=abc.ABCMeta):
         logger.info("Setting up.")
         self.start_time = int(self.get_clock().time())
         self.datastores = Databases(self.DATASTORE_CLASS, self)
+        self.jemalloc_stats = setup_jemalloc_stats(self)
         logger.info("Finished setting up.")
 
         # Register background tasks required by this server. This must be done

synapse/util/caches/lrucache.py

@@ -50,7 +50,6 @@ from twisted.internet.interfaces import IReactorTime
 
 from synapse.config import cache as cache_config
 from synapse.metrics.background_process_metrics import wrap_as_background_process
-from synapse.metrics.jemalloc import get_jemalloc_stats
 from synapse.util import Clock, caches
 from synapse.util.caches import CacheMetric, EvictionReason, register_cache
 from synapse.util.caches.treecache import (
@@ -141,7 +140,7 @@ async def _expire_old_entries(
     evicting_due_to_memory = False
 
     # determine if we're evicting due to memory
-    jemalloc_interface = get_jemalloc_stats()
+    jemalloc_interface = hs.jemalloc_stats
     if jemalloc_interface and autotune_config:
         try:
             jemalloc_interface.refresh_stats()

tests/util/test_lrucache.py

@@ -21,7 +21,7 @@
 
 
 from typing import List, Tuple
-from unittest.mock import Mock, patch
+from unittest.mock import Mock
 
 from synapse.metrics.jemalloc import JemallocStats
 from synapse.types import JsonDict
@@ -342,10 +342,9 @@ class MemoryEvictionTestCase(unittest.HomeserverTestCase):
             }
         }
     )
-    @patch("synapse.util.caches.lrucache.get_jemalloc_stats")
-    def test_evict_memory(self, jemalloc_interface: Mock) -> None:
+    def test_evict_memory(self) -> None:
         mock_jemalloc_class = Mock(spec=JemallocStats)
-        jemalloc_interface.return_value = mock_jemalloc_class
+        self.hs.jemalloc_stats = mock_jemalloc_class()
 
         # set the return value of get_stat() to be greater than max_cache_memory_usage
         mock_jemalloc_class.get_stat.return_value = 924288000