1818import re
1919from typing import Iterable , Optional , overload
2020
21+ import attr
2122from prometheus_client import REGISTRY , Metric
2223from typing_extensions import Literal
2324
2728logger = logging .getLogger (__name__ )
2829
2930
30- def _setup_jemalloc_stats () -> None :
31- """Checks to see if jemalloc is loaded, and hooks up a collector to record
32- statistics exposed by jemalloc.
33- """
34-
35- # Try to find the loaded jemalloc shared library, if any. We need to
36- # introspect into what is loaded, rather than loading whatever is on the
37- # path, as if we load a *different* jemalloc version things will seg fault.
38-
39- # We look in `/proc/self/maps`, which only exists on linux.
40- if not os .path .exists ("/proc/self/maps" ):
41- logger .debug ("Not looking for jemalloc as no /proc/self/maps exist" )
42- return
43-
44- # We're looking for a path at the end of the line that includes
45- # "libjemalloc".
46- regex = re .compile (r"/\S+/libjemalloc.*$" )
47-
48- jemalloc_path = None
49- with open ("/proc/self/maps" ) as f :
50- for line in f :
51- match = regex .search (line .strip ())
52- if match :
53- jemalloc_path = match .group ()
54-
55- if not jemalloc_path :
56- # No loaded jemalloc was found.
57- logger .debug ("jemalloc not found" )
58- return
59-
60- logger .debug ("Found jemalloc at %s" , jemalloc_path )
61-
62- jemalloc = ctypes .CDLL (jemalloc_path )
31+ @attr .s (slots = True , frozen = True , auto_attribs = True )
32+ class JemallocStats :
33+ jemalloc : ctypes .CDLL
6334
6435 @overload
6536 def _mallctl (
66- name : str , read : Literal [True ] = True , write : Optional [int ] = None
37+ self , name : str , read : Literal [True ] = True , write : Optional [int ] = None
6738 ) -> int :
6839 ...
6940
7041 @overload
71- def _mallctl (name : str , read : Literal [False ], write : Optional [int ] = None ) -> None :
42+ def _mallctl (
43+ self , name : str , read : Literal [False ], write : Optional [int ] = None
44+ ) -> None :
7245 ...
7346
7447 def _mallctl (
75- name : str , read : bool = True , write : Optional [int ] = None
48+ self , name : str , read : bool = True , write : Optional [int ] = None
7649 ) -> Optional [int ]:
7750 """Wrapper around `mallctl` for reading and writing integers to
7851 jemalloc.
@@ -120,7 +93,7 @@ def _mallctl(
12093 # Where oldp/oldlenp is a buffer where the old value will be written to
12194 # (if not null), and newp/newlen is the buffer with the new value to set
12295 # (if not null). Note that they're all references *except* newlen.
123- result = jemalloc .mallctl (
96+ result = self . jemalloc .mallctl (
12497 name .encode ("ascii" ),
12598 input_var_ref ,
12699 input_len_ref ,
@@ -136,21 +109,80 @@ def _mallctl(
136109
137110 return input_var .value
138111
139- def _jemalloc_refresh_stats ( ) -> None :
112+ def refresh_stats ( self ) -> None :
140113 """Request that jemalloc updates its internal statistics. This needs to
141114 be called before querying for stats, otherwise it will return stale
142115 values.
143116 """
144117 try :
145- _mallctl ("epoch" , read = False , write = 1 )
118+ self . _mallctl ("epoch" , read = False , write = 1 )
146119 except Exception as e :
147120 logger .warning ("Failed to reload jemalloc stats: %s" , e )
148121
122+ def get_stat (self , name : str ) -> int :
123+ """Request the stat of the given name at the time of the last
124+ `refresh_stats` call. This may throw if we fail to read
125+ the stat.
126+ """
127+ return self ._mallctl (f"stats.{ name } " )
128+
129+
130+ _JEMALLOC_STATS : Optional [JemallocStats ] = None
131+
132+
133+ def get_jemalloc_stats () -> Optional [JemallocStats ]:
134+ """Returns an interface to jemalloc, if it is being used.
135+
136+ Note that this will always return None until `setup_jemalloc_stats` has been
137+ called.
138+ """
139+ return _JEMALLOC_STATS
140+
141+
142+ def _setup_jemalloc_stats () -> None :
143+ """Checks to see if jemalloc is loaded, and hooks up a collector to record
144+ statistics exposed by jemalloc.
145+ """
146+
147+ global _JEMALLOC_STATS
148+
149+ # Try to find the loaded jemalloc shared library, if any. We need to
150+ # introspect into what is loaded, rather than loading whatever is on the
151+ # path, as if we load a *different* jemalloc version things will seg fault.
152+
153+ # We look in `/proc/self/maps`, which only exists on linux.
154+ if not os .path .exists ("/proc/self/maps" ):
155+ logger .debug ("Not looking for jemalloc as no /proc/self/maps exist" )
156+ return
157+
158+ # We're looking for a path at the end of the line that includes
159+ # "libjemalloc".
160+ regex = re .compile (r"/\S+/libjemalloc.*$" )
161+
162+ jemalloc_path = None
163+ with open ("/proc/self/maps" ) as f :
164+ for line in f :
165+ match = regex .search (line .strip ())
166+ if match :
167+ jemalloc_path = match .group ()
168+
169+ if not jemalloc_path :
170+ # No loaded jemalloc was found.
171+ logger .debug ("jemalloc not found" )
172+ return
173+
174+ logger .debug ("Found jemalloc at %s" , jemalloc_path )
175+
176+ jemalloc_dll = ctypes .CDLL (jemalloc_path )
177+
178+ stats = JemallocStats (jemalloc_dll )
179+ _JEMALLOC_STATS = stats
180+
149181 class JemallocCollector (Collector ):
150182 """Metrics for internal jemalloc stats."""
151183
152184 def collect (self ) -> Iterable [Metric ]:
153- _jemalloc_refresh_stats ()
185+ stats . refresh_stats ()
154186
155187 g = GaugeMetricFamily (
156188 "jemalloc_stats_app_memory_bytes" ,
@@ -184,7 +216,7 @@ def collect(self) -> Iterable[Metric]:
184216 "metadata" ,
185217 ):
186218 try :
187- value = _mallctl ( f" stats.{ t } " )
219+ value = stats .get_stat ( t )
188220 except Exception as e :
189221 # There was an error fetching the value, skip.
190222 logger .warning ("Failed to read jemalloc stats.%s: %s" , t , e )
0 commit comments