Skip to content
Snippets Groups Projects
Unverified Commit cf65433d authored by reivilibre's avatar reivilibre Committed by GitHub
Browse files

Fix a memory leak when running the unit tests. (#13798)

parent eaed4e61
Branches
Tags
No related merge requests found
Fix a memory leak when running the unit tests.
\ No newline at end of file
......@@ -205,8 +205,9 @@ def register_cache(
add_resizable_cache(cache_name, resize_callback)
metric = CacheMetric(cache, cache_type, cache_name, collect_callback)
metric_name = "cache_%s_%s" % (cache_type, cache_name)
caches_by_name[cache_name] = cache
CACHE_METRIC_REGISTRY.register_hook(metric.collect)
CACHE_METRIC_REGISTRY.register_hook(metric_name, metric.collect)
return metric
......
......@@ -15,7 +15,7 @@
import logging
from functools import wraps
from types import TracebackType
from typing import Awaitable, Callable, Generator, List, Optional, Type, TypeVar
from typing import Awaitable, Callable, Dict, Generator, Optional, Type, TypeVar
from prometheus_client import CollectorRegistry, Counter, Metric
from typing_extensions import Concatenate, ParamSpec, Protocol
......@@ -220,21 +220,21 @@ class DynamicCollectorRegistry(CollectorRegistry):
def __init__(self) -> None:
super().__init__()
self._pre_update_hooks: List[Callable[[], None]] = []
self._pre_update_hooks: Dict[str, Callable[[], None]] = {}
def collect(self) -> Generator[Metric, None, None]:
"""
Collects metrics, calling pre-update hooks first.
"""
for pre_update_hook in self._pre_update_hooks:
for pre_update_hook in self._pre_update_hooks.values():
pre_update_hook()
yield from super().collect()
def register_hook(self, hook: Callable[[], None]) -> None:
def register_hook(self, metric_name: str, hook: Callable[[], None]) -> None:
"""
Registers a hook that is called before metric collection.
"""
self._pre_update_hooks.append(hook)
self._pre_update_hooks[metric_name] = hook
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment