From 788a1630b71e8a2bd9841ff3cf0babc9ae191b56 Mon Sep 17 00:00:00 2001
From: Owais Lone <owais@lone.pw>
Date: Tue, 13 Apr 2021 16:46:21 +0530
Subject: [PATCH] Added trace response headers to Tornado

---
 .github/workflows/test.yml                    |  2 +-
 CHANGELOG.md                                  |  2 ++
 .../instrumentation/tornado/__init__.py       | 14 +++++++++
 .../tests/test_instrumentation.py             | 31 +++++++++++++++++++
 4 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 9b5e1e5438..ae470547ca 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -6,7 +6,7 @@ on:
     - 'release/*'
   pull_request:
 env:
-  CORE_REPO_SHA: cad261e5dae1fe986c87e6965664b45cc9ab73c3
+  CORE_REPO_SHA: f6b04c483f6c416e1927f010c07e71a17a5d79d0
 
 jobs:
   build:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e398cd8d9d..17693e7452 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#299](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/299))
 - `opentelemetry-instrumenation-django` now supports request and response hooks.
   ([#407](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/407))
+- `opentelemetry-instrumenation-tornado` now supports trace response headers.
+  ([#433](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/433))
 
 ### Removed
 - Remove `http.status_text` from span attributes
diff --git a/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py b/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py
index 5b7654b593..be0ff1bb67 100644
--- a/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py
@@ -46,6 +46,10 @@ def get(self):
 
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    FuncSetter,
+    get_global_back_propagator,
+)
 from opentelemetry.instrumentation.tornado.version import __version__
 from opentelemetry.instrumentation.utils import (
     extract_attributes_from_object,
@@ -68,6 +72,8 @@ def get(self):
 _excluded_urls = get_excluded_urls("TORNADO")
 _traced_request_attrs = get_traced_request_attrs("TORNADO")
 
+back_propagation_setter = FuncSetter(tornado.web.RequestHandler.add_header)
+
 
 class TornadoInstrumentor(BaseInstrumentor):
     patched_handlers = []
@@ -211,6 +217,14 @@ def _start_span(tracer, handler, start_time) -> _TraceContext:
     activation.__enter__()  # pylint: disable=E1101
     ctx = _TraceContext(activation, span, token)
     setattr(handler, _HANDLER_CONTEXT_KEY, ctx)
+
+    # finish handler is called after the response is sent back to
+    # the client so it is too late to inject trace response headers
+    # there.
+    propagator = get_global_back_propagator()
+    if propagator:
+        propagator.inject(handler, setter=back_propagation_setter)
+
     return ctx
 
 
diff --git a/instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py
index d7e167b722..410c179368 100644
--- a/instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py
+++ b/instrumentation/opentelemetry-instrumentation-tornado/tests/test_instrumentation.py
@@ -18,6 +18,11 @@
 from tornado.testing import AsyncHTTPTestCase
 
 from opentelemetry import trace
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    get_global_back_propagator,
+    set_global_back_propagator,
+)
 from opentelemetry.instrumentation.tornado import (
     TornadoInstrumentor,
     patch_handler_class,
@@ -366,6 +371,32 @@ def test_traced_attrs(self):
         )
         self.memory_exporter.clear()
 
+    def test_response_headers(self):
+        orig = get_global_back_propagator()
+        set_global_back_propagator(TraceResponsePropagator())
+
+        response = self.fetch("/")
+        headers = response.headers
+
+        spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
+        self.assertEqual(len(spans), 3)
+        server_span = spans[1]
+
+        self.assertIn("traceresponse", headers)
+        self.assertEqual(
+            headers["access-control-expose-headers"], "traceresponse",
+        )
+        self.assertEqual(
+            headers["traceresponse"],
+            "00-{0}-{1}-01".format(
+                trace.format_trace_id(server_span.get_span_context().trace_id),
+                trace.format_span_id(server_span.get_span_context().span_id),
+            ),
+        )
+
+        self.memory_exporter.clear()
+        set_global_back_propagator(orig)
+
 
 class TestTornadoUninstrument(TornadoTest):
     def test_uninstrument(self):