From 40b85325401bd1c489630311b13a2ef5210b40ca Mon Sep 17 00:00:00 2001 From: Siddham Jain Date: Fri, 23 Oct 2026 00:38:42 +0530 Subject: [PATCH] Support async functions in metric decorators Signed-off-by: Siddham Jain --- docs/content/instrumenting/counter.md | 2 +- docs/content/instrumenting/gauge.md | 4 +- docs/content/instrumenting/histogram.md | 2 +- docs/content/instrumenting/summary.md | 2 +- prometheus_client/context_managers.py | 37 ++++++----- prometheus_client/decorator.py | 11 +++- tests/test_core.py | 82 +++++++++++++++++++++++++ 7 files changed, 115 insertions(+), 25 deletions(-) diff --git a/docs/content/instrumenting/counter.md b/docs/content/instrumenting/counter.md index 4876b612..6f7256c8 100644 --- a/docs/content/instrumenting/counter.md +++ b/docs/content/instrumenting/counter.md @@ -75,7 +75,7 @@ c.reset() ### `count_exceptions(exception=Exception)` Count exceptions raised in a block of code or function. Can be used as a -decorator or context manager. Increments the counter each time an exception +decorator (also on `async def` functions) or context manager. Increments the counter each time an exception of the given type is raised. ```python diff --git a/docs/content/instrumenting/gauge.md b/docs/content/instrumenting/gauge.md index 62294944..3d95ed6a 100644 --- a/docs/content/instrumenting/gauge.md +++ b/docs/content/instrumenting/gauge.md @@ -83,7 +83,7 @@ g.set_to_current_time() ### `track_inprogress()` Increment the gauge when a block of code or function is entered, and decrement -it when exited. Can be used as a decorator or context manager. +it when exited. Can be used as a decorator (also on `async def` functions) or context manager. ```python @g.track_inprogress() @@ -99,7 +99,7 @@ with g.track_inprogress(): Set the gauge to the duration in seconds of the most recent execution of a block of code or function. Unlike `Histogram.time()` and `Summary.time()`, which accumulate all observations, this overwrites the gauge with the latest -duration each time. Can be used as a decorator or context manager. +duration each time. Can be used as a decorator (also on `async def` functions) or context manager. ```python @g.time() diff --git a/docs/content/instrumenting/histogram.md b/docs/content/instrumenting/histogram.md index fa0ffe1a..24bccc87 100644 --- a/docs/content/instrumenting/histogram.md +++ b/docs/content/instrumenting/histogram.md @@ -77,7 +77,7 @@ h.observe(0.43, exemplar={'trace_id': 'abc123'}) Observe the duration in seconds of a block of code or function and add it to the histogram. Every call accumulates — unlike `Gauge.time()`, which only keeps the -most recent duration. Can be used as a decorator or context manager. +most recent duration. Can be used as a decorator (also on `async def` functions) or context manager. ```python @h.time() diff --git a/docs/content/instrumenting/summary.md b/docs/content/instrumenting/summary.md index 714dfd2f..53a21d58 100644 --- a/docs/content/instrumenting/summary.md +++ b/docs/content/instrumenting/summary.md @@ -59,7 +59,7 @@ s.observe(1024) # observe 1024 bytes Observe the duration in seconds of a block of code or function and add it to the summary. Every call accumulates — unlike `Gauge.time()`, which only keeps the -most recent duration. Can be used as a decorator or context manager. +most recent duration. Can be used as a decorator (also on `async def` functions) or context manager. ```python @s.time() diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py index 3e8d7ced..48563cd4 100644 --- a/prometheus_client/context_managers.py +++ b/prometheus_client/context_managers.py @@ -1,3 +1,4 @@ +import inspect from timeit import default_timer from types import TracebackType from typing import ( @@ -12,6 +13,20 @@ F = TypeVar("F", bound=Callable[..., Any]) +def _decorate(f: "F", context: Callable[[], Any]) -> "F": + """Run each call of f inside `with context():`, awaiting coroutine functions.""" + if inspect.iscoroutinefunction(f): + async def wrapped(func, *args, **kwargs): + with context(): + return await func(*args, **kwargs) + else: + def wrapped(func, *args, **kwargs): + with context(): + return func(*args, **kwargs) + + return decorate(f, wrapped) + + class ExceptionCounter: def __init__(self, counter: "Counter", exception: Union[Type[BaseException], Tuple[Type[BaseException], ...]]) -> None: self._counter = counter @@ -26,11 +41,7 @@ def __exit__(self, typ: Optional[Type[BaseException]], value: Optional[BaseExcep return False def __call__(self, f: "F") -> "F": - def wrapped(func, *args, **kwargs): - with self: - return func(*args, **kwargs) - - return decorate(f, wrapped) + return _decorate(f, lambda: self) class InprogressTracker: @@ -44,11 +55,7 @@ def __exit__(self, typ, value, traceback): self._gauge.dec() def __call__(self, f: "F") -> "F": - def wrapped(func, *args, **kwargs): - with self: - return func(*args, **kwargs) - - return decorate(f, wrapped) + return _decorate(f, lambda: self) class Timer: @@ -74,10 +81,6 @@ def labels(self, *args, **kw): self._metric = self._metric.labels(*args, **kw) def __call__(self, f: "F") -> "F": - def wrapped(func, *args, **kwargs): - # Obtaining new instance of timer every time - # ensures thread safety and reentrancy. - with self._new_timer(): - return func(*args, **kwargs) - - return decorate(f, wrapped) + # Obtaining new instance of timer every time + # ensures thread safety and reentrancy. + return _decorate(f, self._new_timer) diff --git a/prometheus_client/decorator.py b/prometheus_client/decorator.py index 1ad2c977..11be76f1 100644 --- a/prometheus_client/decorator.py +++ b/prometheus_client/decorator.py @@ -173,7 +173,7 @@ def make(self, src_templ, evaldict=None, addsource=False, **attrs): "Make a new function from a given template and update the signature" src = src_templ % vars(self) # expand name and signature evaldict = evaldict or {} - mo = DEF.match(src) + mo = DEF.search(src) if mo is None: raise SyntaxError('not a valid function template\n%s' % src) name = mo.group(1) # extract the function name @@ -222,8 +222,13 @@ def create(cls, obj, body, evaldict, defaults=None, func = obj self = cls(func, name, signature, defaults, doc, module) ibody = '\n'.join(' ' + line for line in body.splitlines()) - return self.make('def %(name)s(%(signature)s):\n' + ibody, - evaldict, addsource, **attrs) + caller = evaldict.get('_call_') # when called from `decorate` + if caller and inspect.iscoroutinefunction(caller): + body = ('async def %(name)s(%(signature)s):\n' + ibody).replace( + 'return', 'return await') + else: + body = 'def %(name)s(%(signature)s):\n' + ibody + return self.make(body, evaldict, addsource, **attrs) def decorate(func, caller): diff --git a/tests/test_core.py b/tests/test_core.py index 2c29f79a..c41fdb03 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,4 +1,6 @@ +import asyncio from concurrent.futures import ThreadPoolExecutor +import inspect import os from threading import Barrier, Event import time @@ -105,6 +107,26 @@ def test_block_decorator(self): self.assertTrue(raised) self.assertEqual(1, self.registry.get_sample_value('c_total')) + def test_async_function_decorator(self): + @self.counter.count_exceptions(ValueError) + async def f(r): + await asyncio.sleep(0) + if r: + raise ValueError + else: + raise TypeError + + self.assertTrue(inspect.iscoroutinefunction(f)) + self.assertEqual((["r"], None, None, None), getargspec(f)) + + with self.assertRaises(TypeError): + asyncio.run(f(False)) + self.assertEqual(0, self.registry.get_sample_value('c_total')) + + with self.assertRaises(ValueError): + asyncio.run(f(True)) + self.assertEqual(1, self.registry.get_sample_value('c_total')) + def test_count_exceptions_not_observable(self): counter = Counter('counter', 'help', labelnames=('label',), registry=self.registry) assert_not_observable(counter.count_exceptions) @@ -212,6 +234,20 @@ def f(): f() self.assertEqual(0, self.registry.get_sample_value('g')) + def test_inprogress_async_function_decorator(self): + self.assertEqual(0, self.registry.get_sample_value('g')) + + @self.gauge.track_inprogress() + async def f(): + await asyncio.sleep(0) + self.assertEqual(1, self.registry.get_sample_value('g')) + + self.assertTrue(inspect.iscoroutinefunction(f)) + self.assertEqual(([], None, None, None), getargspec(f)) + + asyncio.run(f()) + self.assertEqual(0, self.registry.get_sample_value('g')) + def test_inprogress_block_decorator(self): self.assertEqual(0, self.registry.get_sample_value('g')) with self.gauge.track_inprogress(): @@ -244,6 +280,19 @@ def f(): f() self.assertNotEqual(0, self.registry.get_sample_value('g')) + def test_time_async_function_decorator(self): + self.assertEqual(0, self.registry.get_sample_value('g')) + + @self.gauge.time() + async def f(): + await asyncio.sleep(.05) + + self.assertTrue(inspect.iscoroutinefunction(f)) + self.assertEqual(([], None, None, None), getargspec(f)) + + asyncio.run(f()) + self.assertLess(.04, self.registry.get_sample_value('g')) + def test_function_decorator_multithread(self): self.assertEqual(0, self.registry.get_sample_value('g')) workers = 2 @@ -372,6 +421,24 @@ def f(i=1): total_expected_duration *= rounding_coefficient self.assertLess(total_expected_duration, self.registry.get_sample_value('s_sum')) + def test_async_function_decorator(self): + self.assertEqual(0, self.registry.get_sample_value('s_count')) + + @self.summary.time() + async def f(duration): + await asyncio.sleep(duration) + + self.assertTrue(inspect.iscoroutinefunction(f)) + self.assertEqual((["duration"], None, None, None), getargspec(f)) + + async def main(): + # Overlapping calls must each be timed by their own timer. + await asyncio.gather(f(.05), f(.1)) + + asyncio.run(main()) + self.assertEqual(2, self.registry.get_sample_value('s_count')) + self.assertLess(.15 * .9, self.registry.get_sample_value('s_sum')) + def test_block_decorator(self): self.assertEqual(0, self.registry.get_sample_value('s_count')) with self.summary.time(): @@ -485,6 +552,21 @@ def f(): self.assertEqual(1, self.registry.get_sample_value('h_count')) self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'})) + def test_async_function_decorator(self): + self.assertEqual(0, self.registry.get_sample_value('h_count')) + + @self.histogram.time() + async def f(): + await asyncio.sleep(.05) + + self.assertTrue(inspect.iscoroutinefunction(f)) + self.assertEqual(([], None, None, None), getargspec(f)) + + asyncio.run(f()) + self.assertEqual(1, self.registry.get_sample_value('h_count')) + self.assertEqual(0, self.registry.get_sample_value('h_bucket', {'le': '0.025'})) + self.assertLess(.04, self.registry.get_sample_value('h_sum')) + def test_function_decorator_multithread(self): self.assertEqual(0, self.registry.get_sample_value('h_count')) workers = 3