Merge pull request #172 from jsocol/fix/119/async-decorator

Fix timing decorator for async functions

3 files changed, 43 insertions, 8 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c488438..71d3666 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,11 @@
 
 - Drops support for Python 2.
 
+### Fixed
+
+- Using a timing decorator on an async function should now properly measure the
+  execution time, instead of counting immediately. See #119.
+
 Version 3.3
 -----------
 
diff --git a/statsd/client/timer.py b/statsd/client/timer.py
index 2f2ab66..453197e 100644
--- a/statsd/client/timer.py
+++ b/statsd/client/timer.py
@@ -1,12 +1,6 @@
 import functools
-
-# Use timer that's not susceptible to time of day adjustments.
-try:
-    # perf_counter is only present on Py3.3+
-    from time import perf_counter as time_now
-except ImportError:
-    # fall back to using time
-    from time import time as time_now
+from inspect import iscoroutinefunction
+from time import perf_counter as time_now
 
 
 def safe_wraps(wrapper, *args, **kwargs):
@@ -29,6 +23,17 @@ class Timer(object):
 
     def __call__(self, f):
         """Thread-safe timing function decorator."""
+        if iscoroutinefunction(f):
+            @safe_wraps(f)
+            async def _async_wrapped(*args, **kwargs):
+                start_time = time_now()
+                try:
+                    return await f(*args, **kwargs)
+                finally:
+                    elapsed_time_ms = 1000.0 * (time_now() - start_time)
+                    self.client.timing(self.stat, elapsed_time_ms, self.rate)
+            return _async_wrapped
+
         @safe_wraps(f)
         def _wrapped(*args, **kwargs):
             start_time = time_now()
diff --git a/statsd/tests.py b/statsd/tests.py
index 69d4358..d1b2a96 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import random
 import re
@@ -651,6 +652,30 @@ def _test_timer_decorator_exceptions(cl, proto):
     _timer_check(cl._sock, 1, proto, 'foo', 'ms')
 
 
+def test_coroutine_timer_decorator():
+    event_loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(event_loop)
+
+    already = False
+    cl = _udp_client()
+
+    def _send(*_):
+        nonlocal already
+        assert already is True, '_send called before coroutine completed'
+
+    cl._send = _send
+
+    @cl.timer('bar')
+    async def inner():
+        nonlocal already
+        await asyncio.sleep(0)
+        already = True
+        return None
+
+    event_loop.run_until_complete(inner())
+    event_loop.close()
+
+
 def test_timer_decorator_exceptions_udp():
     cl = _udp_client()
     _test_timer_decorator_exceptions(cl, 'udp')