diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py index f52bc5cf5..27cf060ed 100644 --- a/datadog/dogstatsd/base.py +++ b/datadog/dogstatsd/base.py @@ -16,7 +16,10 @@ from threading import Lock # datadog -from datadog.dogstatsd.context import TimedContextManagerDecorator +from datadog.dogstatsd.context import ( + TimedContextManagerDecorator, + DistributedContextManagerDecorator, +) from datadog.dogstatsd.route import get_default_route from datadog.util.compat import text from datadog.util.format import normalize_tags @@ -375,6 +378,34 @@ def get_user(user_id): """ return TimedContextManagerDecorator(self, metric, tags, sample_rate, use_ms) + def distributed(self, metric=None, tags=None, sample_rate=None, use_ms=None): + """ + A decorator or context manager that will measure the distribution of a + function's/context's run time using custom metric distribution. + Optionally specify a list of tags or a sample rate. If the metric is not + defined as a decorator, the module name and function name will be used. + The metric is required as a context manager. + :: + + @statsd.distributed('user.query.time', sample_rate=0.5) + def get_user(user_id): + # Do what you need to ... + pass + + # Is equivalent to ... + with statsd.distributed('user.query.time', sample_rate=0.5): + # Do what you need to ... + pass + + # Is equivalent to ... + start = time.time() + try: + get_user(user_id) + finally: + statsd.distribution('user.query.time', time.time() - start) + """ + return DistributedContextManagerDecorator(self, metric, tags, sample_rate, use_ms) + def set(self, metric, value, tags=None, sample_rate=None): """ Sample a set value. diff --git a/datadog/dogstatsd/context.py b/datadog/dogstatsd/context.py index 9de4051dc..49442b4fc 100644 --- a/datadog/dogstatsd/context.py +++ b/datadog/dogstatsd/context.py @@ -17,6 +17,7 @@ class TimedContextManagerDecorator(object): """ def __init__(self, statsd, metric=None, tags=None, sample_rate=1, use_ms=None): self.statsd = statsd + self.timing_func = statsd.timing self.metric = metric self.tags = tags self.sample_rate = sample_rate @@ -60,7 +61,7 @@ def _send(self, start): elapsed = time() - start use_ms = self.use_ms if self.use_ms is not None else self.statsd.use_ms elapsed = int(round(1000 * elapsed)) if use_ms else elapsed - self.statsd.timing(self.metric, elapsed, self.tags, self.sample_rate) + self.timing_func(self.metric, elapsed, self.tags, self.sample_rate) self.elapsed = elapsed def start(self): @@ -68,3 +69,13 @@ def start(self): def stop(self): self.__exit__(None, None, None) + + +class DistributedContextManagerDecorator(TimedContextManagerDecorator): + """ + A context manager and a decorator which will report the elapsed time in + the context OR in a function call using the custom distribution metric. + """ + def __init__(self, statsd, metric=None, tags=None, sample_rate=1, use_ms=None): + super(DistributedContextManagerDecorator, self).__init__(statsd, metric, tags, sample_rate, use_ms) + self.timing_func = statsd.distribution diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py index afc1c9be7..884ca3ac2 100644 --- a/tests/unit/dogstatsd/test_statsd.py +++ b/tests/unit/dogstatsd/test_statsd.py @@ -379,6 +379,50 @@ def test_socket_overflown(self): c = [call("Socket send would block: Socket error, dropping the packet")] mock_log.warning.assert_has_calls(c * 2) + def test_distributed(self): + """ + Measure the distribution of a function's run time using distribution custom metric. + """ + # In seconds + @self.statsd.distributed('distributed.test') + def func(a, b, c=1, d=1): + """docstring""" + time.sleep(0.5) + return (a, b, c, d) + + assert_equal('func', func.__name__) + assert_equal('docstring', func.__doc__) + + result = func(1, 2, d=3) + # Assert it handles args and kwargs correctly. + assert_equal(result, (1, 2, 1, 3)) + + packet = self.recv(2).split("\n")[0] # ignore telemetry packet + name_value, type_ = packet.split('|') + name, value = name_value.split(':') + + assert_equal('d', type_) + assert_equal('distributed.test', name) + self.assert_almost_equal(0.5, float(value), 0.1) + + # Repeat, force timer value in milliseconds + @self.statsd.distributed('distributed.test', use_ms=True) + def func(a, b, c=1, d=1): + """docstring""" + time.sleep(0.5) + return (a, b, c, d) + + func(1, 2, d=3) + + packet = self.recv(2).split("\n")[0] # ignore telemetry packet + name_value, type_ = packet.split('|') + name, value = name_value.split(':') + + assert_equal('d', type_) + assert_equal('distributed.test', name) + self.assert_almost_equal(500, float(value), 100) + + def test_timed(self): """ Measure the distribution of a function's run time.