[core] add support for integration span hooks (#679)

* [core] add Tracer.on and Tracer.emit for span hooks

* [core] rename Tracer.emit to Tracer._emit

* [core] fix mispelling of Tracer._hooks

* [core] do not raise an exception or error log

* [falcon] add span hook documentation for Falcon

* [core] add tests for Tracer.on/Tracer._emit

* [core] add tracer hook argument tests

* [core] register span hooks on config object instead

* [core] fix flake8 issues

* Update ddtrace/settings.py

* Update ddtrace/settings.py

* Update ddtrace/settings.py

* [core] remove Hooks.__getattr__, add Hooks.on alias for Hooks.register

Author: brettlangdon

ddtrace/contrib/falcon/__init__.py

@@ -19,6 +19,28 @@
 
 To enable distributed tracing when using autopatching, set the
 ``DATADOG_FALCON_DISTRIBUTED_TRACING`` environment variable to ``True``.
+
+**Supported span hooks**
+
+The following is a list of available tracer hooks that can be used to intercept
+and modify spans created by this integration.
+
+- ``request``
+    - Called before the response has been finished
+    - ``def on_falcon_request(span, request, response)``
+
+
+Example::
+
+    import falcon
+    from ddtrace import config, patch_all
+    patch_all()
+
+    app = falcon.API()
+
+    @config.falcon.hooks.on('request')
+    def on_falcon_request(span, request, response):
+        span.set_tag('my.custom', 'tag')
 """
 from ...utils.importlib import require_modules
 

ddtrace/contrib/falcon/middleware.py

@@ -4,6 +4,7 @@
 from ddtrace.propagation.http import HTTPPropagator
 from ...compat import iteritems
 from ...ext import AppTypes
+from ...settings import config
 
 
 class TraceMiddleware(object):
@@ -76,6 +77,12 @@ def process_response(self, req, resp, resource, req_succeeded=None):
                 status = _detect_and_set_status_error(err_type, span)
 
         span.set_tag(httpx.STATUS_CODE, status)
+
+        # Emit span hook for this response
+        # DEV: Emit before closing so they can overwrite `span.resource` if they want
+        config.falcon.hooks._emit('request', span, req, resp)
+
+        # Close the span
         span.finish()
 
 

ddtrace/settings.py

@@ -1,7 +1,9 @@
+import collections
 import logging
 
 from copy import deepcopy
 
+from .span import Span
 from .pin import Pin
 from .utils.merge import deepmerge
 
@@ -102,11 +104,12 @@ def __init__(self, global_config, *args, **kwargs):
         :param kwargs:
         """
         super(IntegrationConfig, self).__init__(*args, **kwargs)
-
         self.global_config = global_config
+        self.hooks = Hooks()
 
     def __deepcopy__(self, memodict=None):
         new = IntegrationConfig(self.global_config, deepcopy(dict(self)))
+        new.hooks = deepcopy(self.hooks)
         return new
 
     def __repr__(self):
@@ -115,5 +118,120 @@ def __repr__(self):
         return '{}.{}({})'.format(cls.__module__, cls.__name__, keys)
 
 
+class Hooks(object):
+    """
+    Hooks configuration object is used for registering and calling hook functions
+
+    Example::
+
+        @config.falcon.hooks.on('request')
+        def on_request(span, request, response):
+            pass
+    """
+    __slots__ = ['_hooks']
+
+    def __init__(self):
+        self._hooks = collections.defaultdict(set)
+
+    def __deepcopy__(self, memodict=None):
+        hooks = Hooks()
+        hooks._hooks = deepcopy(self._hooks)
+        return hooks
+
+    def register(self, hook, func=None):
+        """
+        Function used to register a hook for the provided name.
+
+        Example::
+
+            def on_request(span, request, response):
+                pass
+
+            config.falcon.hooks.register('request', on_request)
+
+
+        If no function is provided then a decorator is returned::
+
+            @config.falcon.hooks.register('request')
+            def on_request(span, request, response):
+                pass
+
+        :param hook: The name of the hook to register the function for
+        :type hook: str
+        :param func: The function to register, or ``None`` if a decorator should be returned
+        :type func: function, None
+        :returns: Either a function decorator if ``func is None``, otherwise ``None``
+        :rtype: function, None
+        """
+        # If they didn't provide a function, then return a decorator
+        if not func:
+            def wrapper(func):
+                self.register(hook, func)
+                return func
+            return wrapper
+        self._hooks[hook].add(func)
+
+    # Provide shorthand `on` method for `register`
+    # >>> @config.falcon.hooks.on('request')
+    #     def on_request(span, request, response):
+    #        pass
+    on = register
+
+    def deregister(self, func):
+        """
+        Function to deregister a function from all hooks it was registered under
+
+        Example::
+
+            @config.falcon.hooks.on('request')
+            def on_request(span, request, response):
+                pass
+
+            config.falcon.hooks.deregister(on_request)
+
+
+        :param func: Function hook to register
+        :type func: function
+        """
+        for funcs in self._hooks.values():
+            if func in funcs:
+                funcs.remove(func)
+
+    def _emit(self, hook, span, *args, **kwargs):
+        """
+        Function used to call registered hook functions.
+
+        :param hook: The hook to call functions for
+        :type hook: str
+        :param span: The span to call the hook with
+        :type span: :class:`ddtrace.span.Span`
+        :param *args: Positional arguments to pass to the hook functions
+        :type args: list
+        :param **kwargs: Keyword arguments to pass to the hook functions
+        :type kwargs: dict
+        """
+        # Return early if no hooks are registered
+        if hook not in self._hooks:
+            return
+
+        # Return early if we don't have a Span
+        if not isinstance(span, Span):
+            return
+
+        # Call registered hooks
+        for func in self._hooks[hook]:
+            try:
+                func(span, *args, **kwargs)
+            except Exception as e:
+                # DEV: Use log.debug instead of log.error until we have a throttled logger
+                log.debug('Failed to run hook {} function {}: {}'.format(hook, func, e))
+
+    def __repr__(self):
+        """Return string representation of this class instance"""
+        cls = self.__class__
+        hooks = ','.join(self._hooks.keys())
+        return '{}.{}({})'.format(cls.__module__, cls.__name__, hooks)
+
+
 # Configure our global configuration object
 config = Config()

tests/contrib/falcon/test_suite.py

@@ -1,5 +1,6 @@
 from nose.tools import eq_, ok_
 
+from ddtrace import config
 from ddtrace.ext import errors as errx, http as httpx
 
 from tests.opentracer.utils import init_tracer
@@ -158,3 +159,20 @@ def test_200_ot(self):
         eq_(dd_span.resource, 'GET tests.contrib.falcon.app.resources.Resource200')
         eq_(dd_span.get_tag(httpx.STATUS_CODE), '200')
         eq_(dd_span.get_tag(httpx.URL), 'http://falconframework.org/200')
+
+    def test_falcon_request_hook(self):
+        @config.falcon.hooks.on('request')
+        def on_falcon_request(span, request, response):
+            span.set_tag('my.custom', 'tag')
+
+        out = self.simulate_get('/200')
+        eq_(out.status_code, 200)
+        eq_(out.content.decode('utf-8'), 'Success')
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(len(traces), 1)
+        eq_(len(traces[0]), 1)
+        span = traces[0][0]
+        eq_(span.name, 'falcon.request')
+
+        eq_(span.get_tag('my.custom'), 'tag')

tests/test_global_config.py

@@ -1,15 +1,19 @@
+import mock
 from unittest import TestCase
 
 from nose.tools import eq_, ok_, assert_raises
 
 from ddtrace import config as global_config
 from ddtrace.settings import Config, ConfigException
 
+from .test_tracer import get_dummy_tracer
+
 
 class GlobalConfigTestCase(TestCase):
     """Test the `Configuration` class that stores integration settings"""
     def setUp(self):
         self.config = Config()
+        self.tracer = get_dummy_tracer()
 
     def test_registration(self):
         # ensure an integration can register a new list of settings
@@ -90,3 +94,151 @@ def test_settings_merge_deep(self):
         ))
         eq_(self.config.requests['a']['b']['c'], True)
         eq_(self.config.requests['a']['b']['d'], True)
+
+    def test_settings_hook(self):
+        """
+        When calling `Hooks._emit()`
+            When there is a hook registered
+                we call the hook as expected
+        """
+        # Setup our hook
+        @self.config.web.hooks.on('request')
+        def on_web_request(span):
+            span.set_tag('web.request', '/')
+
+        # Create our span
+        span = self.tracer.start_span('web.request')
+        ok_('web.request' not in span.meta)
+
+        # Emit the span
+        self.config.web.hooks._emit('request', span)
+
+        # Assert we updated the span as expected
+        eq_(span.get_tag('web.request'), '/')
+
+    def test_settings_hook_args(self):
+        """
+        When calling `Hooks._emit()` with arguments
+            When there is a hook registered
+                we call the hook as expected
+        """
+        # Setup our hook
+        @self.config.web.hooks.on('request')
+        def on_web_request(span, request, response):
+            span.set_tag('web.request', request)
+            span.set_tag('web.response', response)
+
+        # Create our span
+        span = self.tracer.start_span('web.request')
+        ok_('web.request' not in span.meta)
+
+        # Emit the span
+        # DEV: The actual values don't matter, we just want to test args + kwargs usage
+        self.config.web.hooks._emit('request', span, 'request', response='response')
+
+        # Assert we updated the span as expected
+        eq_(span.get_tag('web.request'), 'request')
+        eq_(span.get_tag('web.response'), 'response')
+
+    def test_settings_hook_args_failure(self):
+        """
+        When calling `Hooks._emit()` with arguments
+            When there is a hook registered that is missing parameters
+                we do not raise an exception
+        """
+        # Setup our hook
+        # DEV: We are missing the required "response" argument
+        @self.config.web.hooks.on('request')
+        def on_web_request(span, request):
+            span.set_tag('web.request', request)
+
+        # Create our span
+        span = self.tracer.start_span('web.request')
+        ok_('web.request' not in span.meta)
+
+        # Emit the span
+        # DEV: This also asserts that no exception was raised
+        self.config.web.hooks._emit('request', span, 'request', response='response')
+
+        # Assert we did not update the span
+        ok_('web.request' not in span.meta)
+
+    def test_settings_multiple_hooks(self):
+        """
+        When calling `Hooks._emit()`
+            When there are multiple hooks registered
+                we do not raise an exception
+        """
+        # Setup our hooks
+        @self.config.web.hooks.on('request')
+        def on_web_request(span):
+            span.set_tag('web.request', '/')
+
+        @self.config.web.hooks.on('request')
+        def on_web_request2(span):
+            span.set_tag('web.status', 200)
+
+        @self.config.web.hooks.on('request')
+        def on_web_request3(span):
+            span.set_tag('web.method', 'GET')
+
+        # Create our span
+        span = self.tracer.start_span('web.request')
+        ok_('web.request' not in span.meta)
+        ok_('web.status' not in span.meta)
+        ok_('web.method' not in span.meta)
+
+        # Emit the span
+        self.config.web.hooks._emit('request', span)
+
+        # Assert we updated the span as expected
+        eq_(span.get_tag('web.request'), '/')
+        eq_(span.get_tag('web.status'), '200')
+        eq_(span.get_tag('web.method'), 'GET')
+
+    def test_settings_hook_failure(self):
+        """
+        When calling `Hooks._emit()`
+            When the hook raises an exception
+                we do not raise an exception
+        """
+        # Setup our failing hook
+        on_web_request = mock.Mock(side_effect=Exception)
+        self.config.web.hooks.register('request')(on_web_request)
+
+        # Create our span
+        span = self.tracer.start_span('web.request')
+
+        # Emit the span
+        # DEV: This is the test, to ensure no exceptions are raised
+        self.config.web.hooks._emit('request', span)
+        on_web_request.assert_called()
+
+    def test_settings_no_hook(self):
+        """
+        When calling `Hooks._emit()`
+            When no hook is registered
+                we do not raise an exception
+        """
+        # Create our span
+        span = self.tracer.start_span('web.request')
+
+        # Emit the span
+        # DEV: This is the test, to ensure no exceptions are raised
+        self.config.web.hooks._emit('request', span)
+        on_web_request.assert_called()
+
+    def test_settings_no_hook(self):
+        """
+        When calling `Hooks._emit()`
+            When no span is provided
+                we do not raise an exception
+        """
+        # Setup our hooks
+        @self.config.web.hooks.on('request')
+        def on_web_request(span):
+            span.set_tag('web.request', '/')
+
+        # Emit the span
+        # DEV: This is the test, to ensure no exceptions are raised
+        self.config.web.hooks._emit('request', None)