Made Timer.__call__ thread-safe

3 files changed, 39 insertions, 23 deletions

diff --git a/docs/timing.rst b/docs/timing.rst
index ed85481..f8f4c32 100644
--- a/docs/timing.rst
+++ b/docs/timing.rst
@@ -52,9 +52,9 @@ block::
 Using a decorator
 =================
 
-The ``timer`` attribute can also be used as a function decorator. Every
-time the decorated function is called, the time it took to execute will
-be sent to the statsd server.
+The ``timer`` attribute decorates your methods in a thread-safe manner.
+Every time the decorated function is called, the time it took to execute
+will be sent to the statsd server.
 
 ::
 
@@ -70,10 +70,6 @@ be sent to the statsd server.
     myfunc(1, 2)
     myfunc(3, 7)
 
-.. warning::
-   Decorators are not thread-safe and may cause errors when decorated
-   functions are called concurrently. Use context managers or raw timers
-   instead.
 
 
 Using a Timer object directly
diff --git a/statsd/client.py b/statsd/client.py
index 8b4f8e0..4ea0b57 100644
--- a/statsd/client.py
+++ b/statsd/client.py
@@ -20,11 +20,17 @@ class Timer(object):
         self._start_time = None
 
     def __call__(self, f):
+        """Thread-safe timing function decorator."""
         @wraps(f)
-        def wrapper(*args, **kw):
-            with self:
-                return f(*args, **kw)
-        return wrapper
+        def _wrapped(*args, **kwargs):
+            start_time = time.time()
+            try:
+                return_value = f(*args, **kwargs)
+            finally:
+                elapsed_time_ms = int(round(1000 * (time.time() - start_time)))
+                self.client.timing(self.stat, elapsed_time_ms, self.rate)
+            return return_value
+        return _wrapped
 
     def __enter__(self):
         return self.start()
diff --git a/statsd/tests.py b/statsd/tests.py
index da52182..751815c 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -236,20 +236,28 @@ def test_timer_manager():
 
 
 def test_timer_decorator():
-    """StatsClient.timer is a decorator."""
+    """StatsClient.timer is a thread-safe decorator."""
     sc = _client()
 
+    @sc.timer('foo')
+    def foo(a, b):
+        return [a, b]
+
     @sc.timer('bar')
-    def bar():
-        pass
+    def bar(a, b):
+        return [b, a]
 
-    bar()
-    _timer_check(sc, 1, 'bar', 'ms')
+    # make sure it works with more than one decorator, called multiple times,
+    # and that parameters are handled correctly
+    eq_([4, 2], foo(4, 2))
+    _timer_check(sc, 1, 'foo', 'ms')
 
-    # Make sure the decorator works more than once:
-    bar()
+    eq_([2, 4], bar(4, 2))
     _timer_check(sc, 2, 'bar', 'ms')
 
+    eq_([6, 5], bar(5, 6))
+    _timer_check(sc, 3, 'bar', 'ms')
+
 
 def test_timer_capture():
     """You can capture the output of StatsClient.timer."""
@@ -273,13 +281,19 @@ def test_timer_context_rate():
 def test_timer_decorator_rate():
     sc = _client()
 
-    @sc.timer('bar', rate=0.1)
-    def bar():
-        pass
+    @sc.timer('foo', rate=0.1)
+    def foo(a, b):
+        return [b, a]
+
+    @sc.timer('bar', rate=0.2)
+    def bar(a, b=2, c=3):
+        return [c, b, a]
 
-    bar()
+    eq_([2, 4], foo(4, 2))
+    _timer_check(sc, 1, 'foo', 'ms|@0.1')
 
-    _timer_check(sc, 1, 'bar', 'ms|@0.1')
+    eq_([3, 2, 5], bar(5))
+    _timer_check(sc, 2, 'bar', 'ms|@0.2')
 
 
 def test_timer_context_exceptions():