Add histogram context manager and decorator

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.

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

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.