Request/Response hooks for Tornado server and client (#426)

Author: owais

CHANGELOG.md

@@ -32,6 +32,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#407](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/407))
 - `opentelemetry-instrumentation-falcon` FalconInstrumentor now supports request/response hooks.
   ([#415](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/415))
+- `opentelemetry-instrumentation-tornado` Add request/response hooks.
+  ([#426](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/426))
 
 ### Removed
 - Remove `http.status_text` from span attributes

docs-requirements.txt

@@ -33,4 +33,5 @@ PyMySQL~=0.9.3
 pyramid>=1.7
 redis>=2.6
 sqlalchemy>=1.0
+tornado>=6.0
 ddtrace>=0.34.0

docs/instrumentation/tornado/tornado.rst

@@ -0,0 +1,7 @@
+OpenTelemetry Tornado Instrumentation
+======================================
+
+.. automodule:: opentelemetry.instrumentation.tornado
+    :members:
+    :undoc-members:
+    :show-inheritance:

instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py

@@ -33,6 +33,43 @@ def get(self):
     app = tornado.web.Application([(r"/", Handler)])
     app.listen(8080)
     tornado.ioloop.IOLoop.current().start()
+
+Hooks
+*******
+
+Tornado instrumentation supports extending tracing behaviour with the help of hooks.
+It's ``instrument()`` method accepts three optional functions that get called back with the
+created span and some other contextual information. Example:
+
+.. code-block:: python
+
+    # will be called for each incoming request to Tornado
+    # web server. `handler` is an instance of
+    # `tornado.web.RequestHandler`.
+    def server_request_hook(span, handler):
+        pass
+
+    # will be called just before sending out a request with
+    # `tornado.httpclient.AsyncHTTPClient.fetch`.
+    # `request` is an instance of ``tornado.httpclient.HTTPRequest`.
+    def client_request_hook(span, request):
+        pass
+
+    # will be called after a outgoing request made with
+    # `tornado.httpclient.AsyncHTTPClient.fetch` finishes.
+    # `response`` is an instance of ``Future[tornado.httpclient.HTTPResponse]`.
+    def client_resposne_hook(span, future):
+        pass
+
+    # apply tornado instrumentation with hooks
+    TornadoInstrumentor().instrument(
+        server_request_hook=server_request_hook,
+        client_request_hook=client_request_hook,
+        client_response_hook=client_resposne_hook
+    )
+
+API
+---
 """
 
 
@@ -96,9 +133,13 @@ def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")
         tracer = trace.get_tracer(__name__, __version__, tracer_provider)
 
+        client_request_hook = kwargs.get("client_request_hook", None)
+        client_response_hook = kwargs.get("client_response_hook", None)
+        server_request_hook = kwargs.get("server_request_hook", None)
+
         def handler_init(init, handler, args, kwargs):
             cls = handler.__class__
-            if patch_handler_class(tracer, cls):
+            if patch_handler_class(tracer, cls, server_request_hook):
                 self.patched_handlers.append(cls)
             return init(*args, **kwargs)
 
@@ -108,7 +149,9 @@ def handler_init(init, handler, args, kwargs):
         wrap_function_wrapper(
             "tornado.httpclient",
             "AsyncHTTPClient.fetch",
-            partial(fetch_async, tracer),
+            partial(
+                fetch_async, tracer, client_request_hook, client_response_hook
+            ),
         )
 
     def _uninstrument(self, **kwargs):
@@ -119,12 +162,12 @@ def _uninstrument(self, **kwargs):
         self.patched_handlers = []
 
 
-def patch_handler_class(tracer, cls):
+def patch_handler_class(tracer, cls, request_hook=None):
     if getattr(cls, _OTEL_PATCHED_KEY, False):
         return False
 
     setattr(cls, _OTEL_PATCHED_KEY, True)
-    _wrap(cls, "prepare", partial(_prepare, tracer))
+    _wrap(cls, "prepare", partial(_prepare, tracer, request_hook))
     _wrap(cls, "on_finish", partial(_on_finish, tracer))
     _wrap(cls, "log_exception", partial(_log_exception, tracer))
     return True
@@ -146,12 +189,14 @@ def _wrap(cls, method_name, wrapper):
     wrapt.apply_patch(cls, method_name, wrapper)
 
 
-def _prepare(tracer, func, handler, args, kwargs):
+def _prepare(tracer, request_hook, func, handler, args, kwargs):
     start_time = _time_ns()
     request = handler.request
     if _excluded_urls.url_disabled(request.uri):
         return func(*args, **kwargs)
-    _start_span(tracer, handler, start_time)
+    ctx = _start_span(tracer, handler, start_time)
+    if request_hook:
+        request_hook(ctx.span, handler)
     return func(*args, **kwargs)
 
 

instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/client.py

@@ -39,7 +39,7 @@ def _normalize_request(args, kwargs):
     return (new_args, new_kwargs)
 
 
-def fetch_async(tracer, func, _, args, kwargs):
+def fetch_async(tracer, request_hook, response_hook, func, _, args, kwargs):
     start_time = _time_ns()
 
     # Return immediately if no args were provided (error)
@@ -55,6 +55,8 @@ def fetch_async(tracer, func, _, args, kwargs):
     span = tracer.start_span(
         request.method, kind=trace.SpanKind.CLIENT, start_time=start_time,
     )
+    if request_hook:
+        request_hook(span, request)
 
     if span.is_recording():
         attributes = {
@@ -68,12 +70,16 @@ def fetch_async(tracer, func, _, args, kwargs):
         inject(request.headers)
         future = func(*args, **kwargs)
         future.add_done_callback(
-            functools.partial(_finish_tracing_callback, span=span)
+            functools.partial(
+                _finish_tracing_callback,
+                span=span,
+                response_hook=response_hook,
+            )
         )
         return future
 
 
-def _finish_tracing_callback(future, span):
+def _finish_tracing_callback(future, span, response_hook):
     status_code = None
     description = None
     exc = future.exception()
@@ -92,4 +98,6 @@ def _finish_tracing_callback(future, span):
                 description=description,
             )
         )
+    if response_hook:
+        response_hook(span, future)
     span.end()

instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py

@@ -42,7 +42,11 @@ def get_app(self):
         return app
 
     def setUp(self):
-        TornadoInstrumentor().instrument()
+        TornadoInstrumentor().instrument(
+            server_request_hook=getattr(self, "server_request_hook", None),
+            client_request_hook=getattr(self, "client_request_hook", None),
+            client_response_hook=getattr(self, "client_response_hook", None),
+        )
         super().setUp()
         # pylint: disable=protected-access
         self.env_patch = patch.dict(
@@ -367,6 +371,59 @@ def test_traced_attrs(self):
         self.memory_exporter.clear()
 
 
+class TornadoHookTest(TornadoTest):
+    _client_request_hook = None
+    _client_response_hook = None
+    _server_request_hook = None
+
+    def client_request_hook(self, span, handler):
+        if self._client_request_hook is not None:
+            self._client_request_hook(span, handler)
+
+    def client_response_hook(self, span, handler):
+        if self._client_response_hook is not None:
+            self._client_response_hook(span, handler)
+
+    def server_request_hook(self, span, handler):
+        if self._server_request_hook is not None:
+            self._server_request_hook(span, handler)
+
+    def test_hooks(self):
+        def server_request_hook(span, handler):
+            span.update_name("name from server hook")
+            handler.set_header("hello", "world")
+
+        def client_request_hook(span, request):
+            span.update_name("name from client hook")
+
+        def client_response_hook(span, request):
+            span.set_attribute("attr-from-hook", "value")
+
+        self._server_request_hook = server_request_hook
+        self._client_request_hook = client_request_hook
+        self._client_response_hook = client_response_hook
+
+        response = self.fetch("/")
+        self.assertEqual(response.headers.get("hello"), "world")
+
+        spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
+        self.assertEqual(len(spans), 3)
+        server_span = spans[1]
+        self.assertEqual(server_span.kind, SpanKind.SERVER)
+        self.assertEqual(server_span.name, "name from server hook")
+        self.assert_span_has_attributes(server_span, {"uri": "/"})
+        self.memory_exporter.clear()
+
+        client_span = spans[2]
+        self.assertEqual(client_span.kind, SpanKind.CLIENT)
+        self.assertEqual(client_span.name, "name from client hook")
+        self.assert_span_has_attributes(
+            client_span, {"attr-from-hook": "value"}
+        )
+
+        self.memory_exporter.clear()
+
+
 class TestTornadoUninstrument(TornadoTest):
     def test_uninstrument(self):
         response = self.fetch("/")