1 files changed, 16 insertions, 4 deletions
diff --git a/synapse/metrics/jemalloc.py b/synapse/metrics/jemalloc.py
index 98ed9c0829..6bc329f04a 100644
--- a/synapse/metrics/jemalloc.py
+++ b/synapse/metrics/jemalloc.py
@@ -16,11 +16,13 @@ import ctypes
 import logging
 import os
 import re
-from typing import Iterable, Optional
+from typing import Iterable, Optional, overload
 
-from prometheus_client import Metric
+from prometheus_client import REGISTRY, Metric
+from typing_extensions import Literal
 
-from synapse.metrics import REGISTRY, GaugeMetricFamily
+from synapse.metrics import GaugeMetricFamily
+from synapse.metrics._types import Collector
 
 logger = logging.getLogger(__name__)
 
@@ -59,6 +61,16 @@ def _setup_jemalloc_stats() -> None:
 
     jemalloc = ctypes.CDLL(jemalloc_path)
 
+    @overload
+    def _mallctl(
+        name: str, read: Literal[True] = True, write: Optional[int] = None
+    ) -> int:
+        ...
+
+    @overload
+    def _mallctl(name: str, read: Literal[False], write: Optional[int] = None) -> None:
+        ...
+
     def _mallctl(
         name: str, read: bool = True, write: Optional[int] = None
     ) -> Optional[int]:
@@ -134,7 +146,7 @@ def _setup_jemalloc_stats() -> None:
         except Exception as e:
             logger.warning("Failed to reload jemalloc stats: %s", e)
 
-    class JemallocCollector:
+    class JemallocCollector(Collector):
         """Metrics for internal jemalloc stats."""
 
         def collect(self) -> Iterable[Metric]: