Support async functions in metric decorators

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

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()

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()

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()

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)

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):

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