2828import mypy .types
2929from mypy .erasetype import remove_instance_last_known_values
3030from mypy .errorcodes import ErrorCode
31- from mypy .nodes import ARG_NAMED_OPT , TempNode , Var
32- from mypy .plugin import FunctionSigContext , MethodSigContext , Plugin
31+ from mypy .nodes import ARG_NAMED_OPT , ListExpr , NameExpr , TempNode , Var
32+ from mypy .plugin import (
33+ FunctionLike ,
34+ FunctionSigContext ,
35+ MethodSigContext ,
36+ Plugin ,
37+ )
3338from mypy .typeops import bind_self
3439from mypy .types import (
3540 AnyType ,
4348 UnionType ,
4449)
4550
51+ PROMETHEUS_METRIC_MISSING_SERVER_NAME_LABEL = ErrorCode (
52+ "missing-server-name-label" ,
53+ "`SERVER_NAME_LABEL` required in metric" ,
54+ category = "per-homeserver-tenant-metrics" ,
55+ )
56+
4657
4758class SynapsePlugin (Plugin ):
59+ def get_function_signature_hook (
60+ self , fullname : str
61+ ) -> Optional [Callable [[FunctionSigContext ], FunctionLike ]]:
62+ if fullname in (
63+ "prometheus_client.metrics.Counter" ,
64+ # TODO: Add other prometheus_client metrics that need checking as we
65+ # refactor, see https://github.com/element-hq/synapse/issues/18592
66+ ):
67+ return check_prometheus_metric_instantiation
68+
69+ return None
70+
4871 def get_method_signature_hook (
4972 self , fullname : str
5073 ) -> Optional [Callable [[MethodSigContext ], CallableType ]]:
@@ -65,6 +88,85 @@ def get_method_signature_hook(
6588 return None
6689
6790
91+ def check_prometheus_metric_instantiation (ctx : FunctionSigContext ) -> CallableType :
92+ """
93+ Ensure that the `prometheus_client` metrics include the `SERVER_NAME_LABEL` label
94+ when instantiated.
95+
96+ This is important because we support multiple Synapse instances running in the same
97+ process, where all metrics share a single global `REGISTRY`. The `server_name` label
98+ ensures metrics are correctly separated by homeserver.
99+
100+ There are also some metrics that apply at the process level, such as CPU usage,
101+ Python garbage collection, Twisted reactor tick time which shouldn't have the
102+ `SERVER_NAME_LABEL`. In those cases, use use a type ignore comment to disable the
103+ check, e.g. `# type: ignore[missing-server-name-label]`.
104+ """
105+ # The true signature, this isn't being modified so this is what will be returned.
106+ signature : CallableType = ctx .default_signature
107+
108+ # Sanity check the arguments are still as expected in this version of
109+ # `prometheus_client`. ex. `Counter(name, documentation, labelnames, ...)`
110+ #
111+ # `signature.arg_names` should be: ["name", "documentation", "labelnames", ...]
112+ if len (signature .arg_names ) < 3 or signature .arg_names [2 ] != "labelnames" :
113+ ctx .api .fail (
114+ f"Expected the 3rd argument of { signature .name } to be 'labelnames', but got "
115+ f"{ signature .arg_names [2 ]} " ,
116+ ctx .context ,
117+ )
118+ return signature
119+
120+ # Ensure mypy is passing the correct number of arguments because we are doing some
121+ # dirty indexing into `ctx.args` later on.
122+ assert len (ctx .args ) == len (signature .arg_names ), (
123+ f"Expected the list of arguments in the { signature .name } signature ({ len (signature .arg_names )} )"
124+ f"to match the number of arguments from the function signature context ({ len (ctx .args )} )"
125+ )
126+
127+ # Check if the `labelnames` argument includes `SERVER_NAME_LABEL`
128+ #
129+ # `ctx.args` should look like this:
130+ # ```
131+ # [
132+ # [StrExpr("name")],
133+ # [StrExpr("documentation")],
134+ # [ListExpr([StrExpr("label1"), StrExpr("label2")])]
135+ # ...
136+ # ]
137+ # ```
138+ labelnames_arg_expression = ctx .args [2 ][0 ] if len (ctx .args [2 ]) > 0 else None
139+ if isinstance (labelnames_arg_expression , ListExpr ):
140+ # Check if the `labelnames` argument includes the `server_name` label (`SERVER_NAME_LABEL`).
141+ for labelname_expression in labelnames_arg_expression .items :
142+ if (
143+ isinstance (labelname_expression , NameExpr )
144+ and labelname_expression .fullname == "synapse.metrics.SERVER_NAME_LABEL"
145+ ):
146+ # Found the `SERVER_NAME_LABEL`, all good!
147+ break
148+ else :
149+ ctx .api .fail (
150+ f"Expected { signature .name } to include `SERVER_NAME_LABEL` in the list of labels. "
151+ "If this is a process-level metric (vs homeserver-level), use a type ignore comment "
152+ "to disable this check." ,
153+ ctx .context ,
154+ code = PROMETHEUS_METRIC_MISSING_SERVER_NAME_LABEL ,
155+ )
156+ else :
157+ ctx .api .fail (
158+ f"Expected the `labelnames` argument of { signature .name } to be a list of label names "
159+ f"(including `SERVER_NAME_LABEL`), but got { labelnames_arg_expression } . "
160+ "If this is a process-level metric (vs homeserver-level), use a type ignore comment "
161+ "to disable this check." ,
162+ ctx .context ,
163+ code = PROMETHEUS_METRIC_MISSING_SERVER_NAME_LABEL ,
164+ )
165+ return signature
166+
167+ return signature
168+
169+
68170def _get_true_return_type (signature : CallableType ) -> mypy .types .Type :
69171 """
70172 Get the "final" return type of a callable which might return an Awaitable/Deferred.
0 commit comments