From 4bf4ec042213d1ea93043162b0f94904ce49d5bf Mon Sep 17 00:00:00 2001
From: Owais Lone <owais@lone.pw>
Date: Tue, 13 Apr 2021 17:24:49 +0530
Subject: [PATCH] Added opt-in support to return traceresponse headers for
 server instrumentations.

This allows users to configure their web apps to inject trace context
as headers in HTTP responses. This is useful when client side apps need
to connect their spans with the server side spans e.g, in RUM products.

Today the most practical way to do this is to use the Server-Timing
header but in near future we might use the traceresponse header as
described here: https://w3c.github.io/trace-context/#trace-context-http-response-headers-format

Added trace response propagation support for:

    Django
    Falcon
    Flask
    Pyramid
    Tornado
---
 .github/workflows/test.yml                    |  4 +--
 CHANGELOG.md                                  | 10 ++++++
 .../instrumentation/django/middleware.py      |  8 +++++
 .../tests/test_middleware.py                  | 35 ++++++++++++++++++-
 .../instrumentation/falcon/__init__.py        |  9 +++++
 .../tests/test_falcon.py                      | 29 ++++++++++++++-
 .../instrumentation/flask/__init__.py         | 10 ++++++
 .../tests/test_programmatic.py                | 30 ++++++++++++++++
 .../instrumentation/pyramid/callbacks.py      |  7 ++++
 .../tests/test_programmatic.py                | 27 ++++++++++++++
 .../instrumentation/tornado/__init__.py       | 14 ++++++++
 .../tests/test_instrumentation.py             | 31 ++++++++++++++++
 .../instrumentation/wsgi/__init__.py          |  9 +++++
 13 files changed, 219 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e9c5e3eeda..2c54b5d22f 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: 6bd163f6d670319eba6693b8465a068a1828f484
 
 jobs:
   build:
@@ -98,6 +98,6 @@ jobs:
         uses: actions/cache@v2
         with:
           path: .tox
-          key: tox-cache-${{ matrix.tox-environment }}-${{ hashFiles('tox.ini', 'dev-requirements.txt', 'docs-requirements.txt') }}
+          key: tox-cache-v2-${{ matrix.tox-environment }}-${{ hashFiles('tox.ini', 'dev-requirements.txt', 'docs-requirements.txt') }}
       - name: run tox
         run: tox -e ${{ matrix.tox-environment }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52e351539e..957fe94334 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,6 +34,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#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))
+- `opentelemetry-instrumenation-django` now supports trace response headers.
+  ([#436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/436))
+- `opentelemetry-instrumenation-tornado` now supports trace response headers.
+  ([#436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/436))
+- `opentelemetry-instrumenation-pyramid` now supports trace response headers.
+  ([#436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/436))
+- `opentelemetry-instrumenation-falcon` now supports trace response headers.
+  ([#436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/436))
+- `opentelemetry-instrumenation-flask` now supports trace response headers.
+  ([#436](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/436))
 
 ### Removed
 - Remove `http.status_text` from span attributes
diff --git a/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py b/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py
index aeeb42bcb9..9b526c521f 100644
--- a/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py
@@ -20,6 +20,9 @@
 
 from opentelemetry.context import attach, detach
 from opentelemetry.instrumentation.django.version import __version__
+from opentelemetry.instrumentation.propagators import (
+    get_global_response_propagator,
+)
 from opentelemetry.instrumentation.utils import extract_attributes_from_object
 from opentelemetry.instrumentation.wsgi import (
     add_response_attributes,
@@ -179,6 +182,11 @@ def process_response(self, request, response):
                 response,
             )
 
+            propagator = get_global_response_propagator()
+            if propagator:
+                propagator.inject(response)
+
+            # record any exceptions raised while processing the request
             exception = request.META.pop(self._environ_exception_key, None)
             if _DjangoMiddleware._otel_response_hook:
                 _DjangoMiddleware._otel_response_hook(  # pylint: disable=not-callable
diff --git a/instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py b/instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py
index fb12dd24f4..b6174e3000 100644
--- a/instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py
@@ -26,10 +26,19 @@
     DjangoInstrumentor,
     _DjangoMiddleware,
 )
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    set_global_response_propagator,
+)
 from opentelemetry.sdk.trace import Span
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
-from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace import (
+    SpanKind,
+    StatusCode,
+    format_span_id,
+    format_trace_id,
+)
 from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
 
 # pylint: disable=import-error
@@ -309,3 +318,27 @@ def response_hook(span, request, response):
         self.assertIsInstance(response_hook_args[1], HttpRequest)
         self.assertIsInstance(response_hook_args[2], HttpResponse)
         self.assertEqual(response_hook_args[2], response)
+
+    def test_trace_response_headers(self):
+        response = Client().get("/span_name/1234/")
+
+        self.assertNotIn("Server-Timing", response.headers)
+        self.memory_exporter.clear()
+
+        set_global_response_propagator(TraceResponsePropagator())
+
+        response = Client().get("/span_name/1234/")
+        span = self.memory_exporter.get_finished_spans()[0]
+
+        self.assertIn("traceresponse", response.headers)
+        self.assertEqual(
+            response.headers["Access-Control-Expose-Headers"], "traceresponse",
+        )
+        self.assertEqual(
+            response.headers["traceresponse"],
+            "00-{0}-{1}-01".format(
+                format_trace_id(span.get_span_context().trace_id),
+                format_span_id(span.get_span_context().span_id),
+            ),
+        )
+        self.memory_exporter.clear()
diff --git a/instrumentation/opentelemetry-instrumentation-falcon/src/opentelemetry/instrumentation/falcon/__init__.py b/instrumentation/opentelemetry-instrumentation-falcon/src/opentelemetry/instrumentation/falcon/__init__.py
index 9aa3c14ea0..e5b6bae7d2 100644
--- a/instrumentation/opentelemetry-instrumentation-falcon/src/opentelemetry/instrumentation/falcon/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-falcon/src/opentelemetry/instrumentation/falcon/__init__.py
@@ -99,6 +99,10 @@ def response_hook(span, req, resp):
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.falcon.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    FuncSetter,
+    get_global_response_propagator,
+)
 from opentelemetry.instrumentation.utils import (
     extract_attributes_from_object,
     http_status_to_status_code,
@@ -119,6 +123,7 @@ def response_hook(span, req, resp):
 
 _excluded_urls = get_excluded_urls("FALCON")
 _traced_request_attrs = get_traced_request_attrs("FALCON")
+_response_propagation_setter = FuncSetter(falcon.api.Response.append_header)
 
 
 class FalconInstrumentor(BaseInstrumentor):
@@ -273,5 +278,9 @@ def process_response(
                 )
             )
 
+        propagator = get_global_response_propagator()
+        if propagator:
+            propagator.inject(resp, setter=_response_propagation_setter)
+
         if self._response_hook:
             self._response_hook(span, req, resp)
diff --git a/instrumentation/opentelemetry-instrumentation-falcon/tests/test_falcon.py b/instrumentation/opentelemetry-instrumentation-falcon/tests/test_falcon.py
index 61c752afcd..016529a424 100644
--- a/instrumentation/opentelemetry-instrumentation-falcon/tests/test_falcon.py
+++ b/instrumentation/opentelemetry-instrumentation-falcon/tests/test_falcon.py
@@ -17,8 +17,13 @@
 from falcon import testing
 
 from opentelemetry.instrumentation.falcon import FalconInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    get_global_response_propagator,
+    set_global_response_propagator,
+)
 from opentelemetry.test.test_base import TestBase
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import StatusCode, format_span_id, format_trace_id
 from opentelemetry.util.http import get_excluded_urls, get_traced_request_attrs
 
 from .app import make_app
@@ -197,6 +202,28 @@ def test_traced_request_attributes(self):
         self.assertEqual(span.attributes["query_string"], "q=abc")
         self.assertNotIn("not_available_attr", span.attributes)
 
+    def test_trace_response(self):
+        orig = get_global_response_propagator()
+        set_global_response_propagator(TraceResponsePropagator())
+
+        response = self.client().simulate_get(path="/hello?q=abc")
+        headers = response.headers
+        span = self.memory_exporter.get_finished_spans()[0]
+
+        self.assertIn("traceresponse", headers)
+        self.assertEqual(
+            headers["access-control-expose-headers"], "traceresponse",
+        )
+        self.assertEqual(
+            headers["traceresponse"],
+            "00-{0}-{1}-01".format(
+                format_trace_id(span.get_span_context().trace_id),
+                format_span_id(span.get_span_context().span_id),
+            ),
+        )
+
+        set_global_response_propagator(orig)
+
     def test_traced_not_recording(self):
         mock_tracer = Mock()
         mock_span = Mock()
diff --git a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py
index f657a1f3bd..2b114f6821 100644
--- a/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py
@@ -55,6 +55,9 @@ def hello():
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.flask.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    get_global_response_propagator,
+)
 from opentelemetry.propagate import extract
 from opentelemetry.util._time import _time_ns
 from opentelemetry.util.http import get_excluded_urls
@@ -91,6 +94,13 @@ def _start_response(status, response_headers, *args, **kwargs):
             if not _excluded_urls.url_disabled(flask.request.url):
                 span = flask.request.environ.get(_ENVIRON_SPAN_KEY)
 
+                propagator = get_global_response_propagator()
+                if propagator:
+                    propagator.inject(
+                        response_headers,
+                        setter=otel_wsgi.default_response_propagation_setter,
+                    )
+
                 if span:
                     otel_wsgi.add_response_attributes(
                         span, status, response_headers
diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py
index 3c62dd751d..aead569de8 100644
--- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py
+++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py
@@ -18,6 +18,11 @@
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.flask import FlaskInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    get_global_response_propagator,
+    set_global_response_propagator,
+)
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.util.http import get_excluded_urls
@@ -119,6 +124,31 @@ def test_simple(self):
         self.assertEqual(span_list[0].kind, trace.SpanKind.SERVER)
         self.assertEqual(span_list[0].attributes, expected_attrs)
 
+    def test_trace_response(self):
+        orig = get_global_response_propagator()
+
+        set_global_response_propagator(TraceResponsePropagator())
+        response = self.client.get("/hello/123")
+        headers = response.headers
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+        span = span_list[0]
+
+        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(span.get_span_context().trace_id),
+                trace.format_span_id(span.get_span_context().span_id),
+            ),
+        )
+
+        set_global_response_propagator(orig)
+
     def test_not_recording(self):
         mock_tracer = Mock()
         mock_span = Mock()
diff --git a/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py b/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py
index 45f65793ad..9e777623eb 100644
--- a/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py
+++ b/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py
@@ -21,6 +21,9 @@
 
 import opentelemetry.instrumentation.wsgi as otel_wsgi
 from opentelemetry import context, trace
+from opentelemetry.instrumentation.propagators import (
+    get_global_response_propagator,
+)
 from opentelemetry.instrumentation.pyramid.version import __version__
 from opentelemetry.propagate import extract
 from opentelemetry.util._time import _time_ns
@@ -157,6 +160,10 @@ def trace_tween(request):
                     response_or_exception.headers,
                 )
 
+                propagator = get_global_response_propagator()
+                if propagator:
+                    propagator.inject(response.headers)
+
                 activation = request.environ.get(_ENVIRON_ACTIVATION_KEY)
 
                 if isinstance(response_or_exception, HTTPException):
diff --git a/instrumentation/opentelemetry-instrumentation-pyramid/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-pyramid/tests/test_programmatic.py
index ea3fd266a1..ce4151ad95 100644
--- a/instrumentation/opentelemetry-instrumentation-pyramid/tests/test_programmatic.py
+++ b/instrumentation/opentelemetry-instrumentation-pyramid/tests/test_programmatic.py
@@ -17,6 +17,11 @@
 from pyramid.config import Configurator
 
 from opentelemetry import trace
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    get_global_response_propagator,
+    set_global_response_propagator,
+)
 from opentelemetry.instrumentation.pyramid import PyramidInstrumentor
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
@@ -98,6 +103,28 @@ def test_simple(self):
         self.assertEqual(span_list[0].kind, trace.SpanKind.SERVER)
         self.assertEqual(span_list[0].attributes, expected_attrs)
 
+    def test_response_headers(self):
+        orig = get_global_response_propagator()
+        set_global_response_propagator(TraceResponsePropagator())
+
+        response = self.client.get("/hello/500")
+        headers = response.headers
+        span = self.memory_exporter.get_finished_spans()[0]
+
+        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(span.get_span_context().trace_id),
+                trace.format_span_id(span.get_span_context().span_id),
+            ),
+        )
+
+        set_global_response_propagator(orig)
+
     def test_not_recording(self):
         mock_tracer = Mock()
         mock_span = Mock()
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 2fd6c7ffae..d7cf74b93f 100644
--- a/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-tornado/src/opentelemetry/instrumentation/tornado/__init__.py
@@ -83,6 +83,10 @@ def client_resposne_hook(span, future):
 
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.propagators import (
+    FuncSetter,
+    get_global_response_propagator,
+)
 from opentelemetry.instrumentation.tornado.version import __version__
 from opentelemetry.instrumentation.utils import (
     extract_attributes_from_object,
@@ -105,6 +109,8 @@ 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)
+
 
 class TornadoInstrumentor(BaseInstrumentor):
     patched_handlers = []
@@ -256,6 +262,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_response_propagator()
+    if propagator:
+        propagator.inject(handler, setter=response_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 d16e69e933..003f544008 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_response_propagator,
+    set_global_response_propagator,
+)
 from opentelemetry.instrumentation.tornado import (
     TornadoInstrumentor,
     patch_handler_class,
@@ -370,6 +375,32 @@ def test_traced_attrs(self):
         )
         self.memory_exporter.clear()
 
+    def test_response_headers(self):
+        orig = get_global_response_propagator()
+        set_global_response_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_response_propagator(orig)
+
 
 class TornadoHookTest(TornadoTest):
     _client_request_hook = None
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
index 2dffe0c5a1..406c275ea7 100644
--- a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
@@ -59,6 +59,7 @@ def hello():
 import wsgiref.util as wsgiref_util
 
 from opentelemetry import context, trace
+from opentelemetry.instrumentation.propagators import Setter as ResponseSetter
 from opentelemetry.instrumentation.utils import http_status_to_status_code
 from opentelemetry.instrumentation.wsgi.version import __version__
 from opentelemetry.propagate import extract
@@ -264,3 +265,11 @@ def _end_span_after_iterating(iterable, span, tracer, token):
             close()
         span.end()
         context.detach(token)
+
+
+class ResponsePropagationSetter(ResponseSetter):
+    def set(self, carrier, key, value):  # pylint: disable=no-self-use
+        carrier.append((key, value))
+
+
+default_response_propagation_setter = ResponsePropagationSetter()