Tornado: Capture custom request/response headers as span attributes (#950)

Author: ashu658

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `opentelemetry-instrumentation-flask` Flask: Capture custom request/response headers in span attributes
   ([#952])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/952)
 
+- `opentelemetry-instrumentation-tornado` Tornado: Capture custom request/response headers in span attributes
+  ([#950])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/950)
+
 ### Added
 
 - `opentelemetry-instrumentation-aws-lambda` `SpanKind.SERVER` by default, add more cases for `SpanKind.CONSUMER` services. ([#926](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/926))

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

@@ -129,7 +129,15 @@ def client_resposne_hook(span, future):
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util._time import _time_ns
-from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_custom_headers,
+    get_excluded_urls,
+    get_traced_request_attrs,
+    normalise_request_header_name,
+    normalise_response_header_name,
+)
 
 from .client import fetch_async  # pylint: disable=E0401
 
@@ -141,7 +149,6 @@ def client_resposne_hook(span, future):
 
 _excluded_urls = get_excluded_urls("TORNADO")
 _traced_request_attrs = get_traced_request_attrs("TORNADO")
-
 response_propagation_setter = FuncSetter(tornado.web.RequestHandler.add_header)
 
 
@@ -257,6 +264,32 @@ def _log_exception(tracer, func, handler, args, kwargs):
     return func(*args, **kwargs)
 
 
+def _add_custom_request_headers(span, request_headers):
+    custom_request_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+    )
+    attributes = {}
+    for header_name in custom_request_headers_name:
+        header_values = request_headers.get(header_name)
+        if header_values:
+            key = normalise_request_header_name(header_name.lower())
+            attributes[key] = [header_values]
+    span.set_attributes(attributes)
+
+
+def _add_custom_response_headers(span, response_headers):
+    custom_response_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+    )
+    attributes = {}
+    for header_name in custom_response_headers_name:
+        header_values = response_headers.get(header_name)
+        if header_values:
+            key = normalise_response_header_name(header_name.lower())
+            attributes[key] = [header_values]
+    span.set_attributes(attributes)
+
+
 def _get_attributes_from_request(request):
     attrs = {
         SpanAttributes.HTTP_METHOD: request.method,
@@ -307,6 +340,8 @@ def _start_span(tracer, handler, start_time) -> _TraceContext:
         for key, value in attributes.items():
             span.set_attribute(key, value)
         span.set_attribute("tornado.handler", _get_full_handler_name(handler))
+        if span.kind == trace.SpanKind.SERVER:
+            _add_custom_request_headers(span, handler.request.headers)
 
     activation = trace.use_span(span, end_on_exit=True)
     activation.__enter__()  # pylint: disable=E1101
@@ -360,6 +395,8 @@ def _finish_span(tracer, handler, error=None):
                 description=otel_status_description,
             )
         )
+        if ctx.span.kind == trace.SpanKind.SERVER:
+            _add_custom_response_headers(ctx.span, handler._headers)
 
     ctx.activation.__exit__(*finish_args)  # pylint: disable=E1101
     if ctx.token:

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

@@ -32,7 +32,12 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import SpanKind
-from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_excluded_urls,
+    get_traced_request_attrs,
+)
 
 from .tornado_test_app import (
     AsyncHandler,
@@ -604,3 +609,122 @@ def test_mark_span_internal_in_presence_of_another_span(self):
         self.assertEqual(
             test_span.context.span_id, tornado_handler_span.parent.span_id
         )
+
+
+class TestTornadoCustomRequestResponseHeadersAddedWithServerSpan(TornadoTest):
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3"
+        },
+    )
+    def test_custom_request_headers_added_in_server_span(self):
+        headers = {
+            "Custom-Test-Header-1": "Test Value 1",
+            "Custom-Test-Header-2": "TestValue2,TestValue3",
+        }
+        response = self.fetch("/", headers=headers)
+        self.assertEqual(response.code, 201)
+        _, tornado_span, _ = self.sorted_spans(
+            self.memory_exporter.get_finished_spans()
+        )
+        expected = {
+            "http.request.header.custom_test_header_1": ("Test Value 1",),
+            "http.request.header.custom_test_header_2": (
+                "TestValue2,TestValue3",
+            ),
+        }
+        self.assertEqual(tornado_span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(tornado_span, expected)
+
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header"
+        },
+    )
+    def test_custom_response_headers_added_in_server_span(self):
+        response = self.fetch("/test_custom_response_headers")
+        self.assertEqual(response.code, 200)
+        tornado_span, _ = self.sorted_spans(
+            self.memory_exporter.get_finished_spans()
+        )
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("0",),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        self.assertEqual(tornado_span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(tornado_span, expected)
+
+
+class TestTornadoCustomRequestResponseHeadersNotAddedWithInternalSpan(
+    TornadoTest
+):
+    def get_app(self):
+        tracer = trace.get_tracer(__name__)
+        app = make_app(tracer)
+
+        def middleware(request):
+            """Wraps the request with a server span"""
+            with tracer.start_as_current_span(
+                "test", kind=trace.SpanKind.SERVER
+            ):
+                app(request)
+
+        return middleware
+
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3"
+        },
+    )
+    def test_custom_request_headers_not_added_in_internal_span(self):
+        headers = {
+            "Custom-Test-Header-1": "Test Value 1",
+            "Custom-Test-Header-2": "TestValue2,TestValue3",
+        }
+        response = self.fetch("/", headers=headers)
+        self.assertEqual(response.code, 201)
+        _, tornado_span, _, _ = self.sorted_spans(
+            self.memory_exporter.get_finished_spans()
+        )
+        not_expected = {
+            "http.request.header.custom_test_header_1": ("Test Value 1",),
+            "http.request.header.custom_test_header_2": (
+                "TestValue2,TestValue3",
+            ),
+        }
+        self.assertEqual(tornado_span.kind, trace.SpanKind.INTERNAL)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, tornado_span.attributes)
+
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header"
+        },
+    )
+    def test_custom_response_headers_not_added_in_internal_span(self):
+        response = self.fetch("/test_custom_response_headers")
+        self.assertEqual(response.code, 200)
+        tornado_span, _, _ = self.sorted_spans(
+            self.memory_exporter.get_finished_spans()
+        )
+        not_expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("0",),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        self.assertEqual(tornado_span.kind, trace.SpanKind.INTERNAL)
+        for key, _ in not_expected.items():
+            self.assertNotIn(key, tornado_span.attributes)

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

@@ -95,6 +95,16 @@ def get(self):
         self.set_status(200)
 
 
+class CustomResponseHeaderHandler(tornado.web.RequestHandler):
+    def get(self):
+        self.set_header("content-type", "text/plain; charset=utf-8")
+        self.set_header("content-length", "0")
+        self.set_header(
+            "my-custom-header", "my-custom-value-1,my-custom-header-2"
+        )
+        self.set_status(200)
+
+
 def make_app(tracer):
     app = tornado.web.Application(
         [
@@ -105,6 +115,7 @@ def make_app(tracer):
             (r"/on_finish", FinishedHandler),
             (r"/healthz", HealthCheckHandler),
             (r"/ping", HealthCheckHandler),
+            (r"/test_custom_response_headers", CustomResponseHeaderHandler),
         ]
     )
     app.tracer = tracer