diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py index 1bdb47f5d..731edbf45 100644 --- a/datadog/dogstatsd/base.py +++ b/datadog/dogstatsd/base.py @@ -23,6 +23,7 @@ from datadog.dogstatsd.context import ( TimedContextManagerDecorator, DistributedContextManagerDecorator, + HistogrammedContextManagerDecorator, ) from datadog.dogstatsd.route import get_default_route from datadog.dogstatsd.container import ContainerID @@ -732,6 +733,34 @@ def get_user(user_id): """ return DistributedContextManagerDecorator(self, metric, tags, sample_rate, use_ms) + def histogrammed(self, metric=None, tags=None, sample_rate=None, use_ms=None): + """ + A decorator or context manager that will measure the histogram of a + function's/context's run time using custom metric histogram. + 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.histogrammed("user.query.time", sample_rate=0.5) + def get_user(user_id): + # Do what you need to ... + pass + + # Is equivalent to ... + with statsd.histogrammed("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.histogram("user.query.time", time.time() - start) + """ + return HistogrammedContextManagerDecorator(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 90e9ce90e..0cbf842fc 100644 --- a/datadog/dogstatsd/context.py +++ b/datadog/dogstatsd/context.py @@ -86,3 +86,14 @@ class DistributedContextManagerDecorator(TimedContextManagerDecorator): 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 + + +class HistogrammedContextManagerDecorator(TimedContextManagerDecorator): + """ + A context manager and a decorator which will report the elapsed time in + the context OR in a function call using the custom histogram metric. + """ + + def __init__(self, statsd, metric=None, tags=None, sample_rate=1, use_ms=None): + super(HistogrammedContextManagerDecorator, self).__init__(statsd, metric, tags, sample_rate, use_ms) + self.timing_func = statsd.histogram diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py index 658804a19..993be1db1 100644 --- a/tests/unit/dogstatsd/test_statsd.py +++ b/tests/unit/dogstatsd/test_statsd.py @@ -697,6 +697,50 @@ def test_udp_socket_ensures_min_receive_buffer(self, mock_socket_create): MIN_SEND_BUFFER_SIZE, ) + def test_histogrammed(self): + """ + Measure the histogram of a function's run time using histogram custom metric. + """ + # In seconds + @self.statsd.histogrammed('histogram.test') + def func(arg1, arg2, kwarg1=1, kwarg2=1): + """docstring""" + time.sleep(0.1) + return (arg1, arg2, kwarg1, kwarg2) + + self.assertEqual('func', func.__name__) + self.assertEqual('docstring', func.__doc__) + + result = func(1, 2, kwarg2=3) + # Assert it handles args and kwargs correctly. + self.assertEqual(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(':') + + self.assertEqual('h', type_) + self.assertEqual('histogram.test', name) + self.assert_almost_equal(0.1, float(value), 0.09) + + # Repeat, force timer value in milliseconds + @self.statsd.histogrammed('histogrammed.test', use_ms=True) + def func(arg1, arg2, kwarg1=1, kwarg2=1): + """docstring""" + time.sleep(0.5) + return (arg1, arg2, kwarg1, kwarg2) + + func(1, 2, kwarg2=3) + + # Ignore telemetry packet + packet = self.recv(2, reset_wait=True).split("\n")[0] + name_value, type_ = packet.split('|') + name, value = name_value.split(':') + + self.assertEqual('h', type_) + self.assertEqual('histogrammed.test', name) + self.assert_almost_equal(500, float(value), 100) + def test_distributed(self): """ Measure the distribution of a function's run time using distribution custom metric.