Async Timer, ExceptionTracker, and InprogressTracker implementations.

Signed-off-by: Cameron Lee <[email protected]>

Author: camlee

prometheus_client/context_managers.py

@@ -2,6 +2,7 @@
 from timeit import default_timer
 from types import TracebackType
 from typing import Any, Callable, Optional, Type, TYPE_CHECKING, TypeVar
+from inspect import iscoroutinefunction
 
 if sys.version_info >= (3, 8, 0):
     from typing import Literal
@@ -26,10 +27,16 @@ def __exit__(self, typ: Optional[Type[BaseException]], value: Optional[BaseExcep
             self._counter.inc()
         return False
 
+
     def __call__(self, f: "F") -> "F":
-        def wrapped(func, *args, **kwargs):
-            with self:
-                return func(*args, **kwargs)
+        if iscoroutinefunction(f):
+            async def wrapped(func, *args, **kwargs):
+                with self:
+                    return await func(*args, **kwargs)
+        else:
+            def wrapped(func, *args, **kwargs):
+                with self:
+                    return func(*args, **kwargs)
 
         return decorate(f, wrapped)
 
@@ -45,9 +52,14 @@ def __exit__(self, typ, value, traceback):
         self._gauge.dec()
 
     def __call__(self, f):
-        def wrapped(func, *args, **kwargs):
-            with self:
-                return func(*args, **kwargs)
+        if iscoroutinefunction(f):
+            async def wrapped(func, *args, **kwargs):
+                with self:
+                    return await func(*args, **kwargs)
+        else:
+            def wrapped(func, *args, **kwargs):
+                with self:
+                    return func(*args, **kwargs)
 
         return decorate(f, wrapped)
 
@@ -74,10 +86,17 @@ def labels(self, *args, **kw):
         self._metric = self._metric.labels(*args, **kw)
 
     def __call__(self, 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)
+        if iscoroutinefunction(f):
+            # If an async function was decorated, we support that by
+            # producing an async function and awaiting the original.
+            async def wrapped(func, *args, **kwargs):
+                with self._new_timer():
+                    return await func(*args, **kwargs)
+        else:
+            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)

tests/test_core.py

@@ -29,6 +29,10 @@ def assert_not_observable(fn, *args, **kwargs):
     assert False, "Did not raise a 'missing label values' exception"
 
 
+def assert_between(lower, value, upper, msg=""):
+    assert lower <= value <= upper, "%s is not between %s and %s%s" % (value, lower, upper, " : %s" % msg if msg else "")
+
+
 class TestCounter(aiounittest.AsyncTestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
@@ -215,7 +219,7 @@ def f():
         self.assertEqual(([], None, None, None), getargspec(f))
 
         f()
-        self.assertTrue(0.05 <= self.registry.get_sample_value('g') <= 0.1)
+        assert_between(0.05, self.registry.get_sample_value('g'), 0.1)
 
     async def test_time_async_function_decorator(self):
         self.assertEqual(0, self.registry.get_sample_value('g'))
@@ -227,7 +231,7 @@ async def f():
         self.assertEqual(([], None, None, None), getargspec(f))
 
         await f()
-        self.assertTrue(0.05 <= self.registry.get_sample_value('g') <= 0.1)
+        assert_between(0.05, self.registry.get_sample_value('g'), 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('g'))
@@ -307,7 +311,7 @@ def f():
 
         f()
         self.assertEqual(1, self.registry.get_sample_value('s_count'))
-        self.assertTrue(.05 < self.registry.get_sample_value('s_sum') < 0.1)
+        assert_between(.05, self.registry.get_sample_value('s_sum'), 0.1)
 
     async def test_async_function_decorator(self):
         self.assertEqual(0, self.registry.get_sample_value('s_count'))
@@ -320,7 +324,7 @@ async def f():
 
         await f()
         self.assertEqual(1, self.registry.get_sample_value('s_count'))
-        self.assertTrue(.05 < self.registry.get_sample_value('s_sum') < 0.1)
+        assert_between(.05, self.registry.get_sample_value('s_sum'), 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('s_count'))
@@ -475,7 +479,7 @@ def f():
         f()
         self.assertEqual(1, self.registry.get_sample_value('h_count'))
         self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
-        self.assertTrue(.05 < self.registry.get_sample_value('h_sum') < 0.1)
+        assert_between(.05, self.registry.get_sample_value('h_sum'), 0.1)
 
     async def test_async_function_decorator(self):
         self.assertEqual(0, self.registry.get_sample_value('h_count'))
@@ -490,7 +494,7 @@ async def f():
         await f()
         self.assertEqual(1, self.registry.get_sample_value('h_count'))
         self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
-        self.assertTrue(.05 < self.registry.get_sample_value('h_sum') < 0.1)
+        assert_between(.05, self.registry.get_sample_value('h_sum'), 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('h_count'))