Add traceresponse headers for asgi apps (FastAPI, Starlette)

This asgi version is modeled after the original wsgi version in open-telemetry#436.

CHANGELOG.md

@@ -18,6 +18,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `opentelemetry-instrumentation-aws-lambda` Adds support for configurable flush timeout  via `OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT` property. ([#825](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/825))
 
+### Added
+
+`opentelemetry-instrumenation-asgi` now returns a `traceresponse` response header.
+  ([#817](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/817))
+
 ### Fixed
 
 - `opentelemetry-exporter-richconsole` Fixed attribute error on parentless spans.
@@ -28,6 +33,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [1.7.1-0.26b1](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.7.0-0.26b0) - 2021-11-11
 
+### Added
+
 - `opentelemetry-instrumentation-aws-lambda` Add instrumentation for AWS Lambda Service - pkg metadata files (Part 1/2)
   ([#739](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/739))
 - Add support for Python 3.10

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

@@ -103,9 +103,12 @@ def client_response_hook(span: Span, message: dict):
 
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.asgi.version import __version__  # noqa
+from opentelemetry.instrumentation.propagators import (
+    get_global_response_propagator,
+)
 from opentelemetry.instrumentation.utils import http_status_to_status_code
 from opentelemetry.propagate import extract
-from opentelemetry.propagators.textmap import Getter
+from opentelemetry.propagators.textmap import Getter, Setter
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import Span
 from opentelemetry.trace.status import Status, StatusCode
@@ -152,6 +155,30 @@ def keys(self, carrier: dict) -> typing.List[str]:
 asgi_getter = ASGIGetter()
 
 
+class ASGISetter(Setter):
+    def set(
+        self, carrier: dict, key: str, value: str
+    ) -> None:  # pylint: disable=no-self-use
+        """Sets response header values on an ASGI scope according to `the spec <https://asgi.readthedocs.io/en/latest/specs/www.html#response-start-send-event>`_.
+
+        Args:
+            carrier: ASGI scope object
+            key: response header name to set
+            value: response header value
+        Returns:
+            None
+        """
+        headers = carrier.get("headers")
+        if not headers:
+            headers = []
+            carrier["headers"] = headers
+
+        headers.append([key.lower().encode(), value.encode()])
+
+
+asgi_setter = ASGISetter()
+
+
 def collect_request_attributes(scope):
     """Collects HTTP request attributes from the ASGI scope and returns a
     dictionary to be used as span creation attributes."""
@@ -341,6 +368,11 @@ async def wrapped_send(message):
                                 set_status_code(span, 200)
                                 set_status_code(send_span, 200)
                             send_span.set_attribute("type", message["type"])
+
+                        propagator = get_global_response_propagator()
+                        if propagator:
+                            propagator.inject(message, setter=asgi_setter)
+
                         await send(message)
 
                 await self.app(scope, wrapped_receive, wrapped_send)

instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_middleware.py

@@ -18,13 +18,19 @@
 
 import opentelemetry.instrumentation.asgi as otel_asgi
 from opentelemetry import trace as trace_api
+from opentelemetry.instrumentation.propagators import (
+    TraceResponsePropagator,
+    get_global_response_propagator,
+    set_global_response_propagator,
+)
 from opentelemetry.sdk import resources
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.asgitestutil import (
     AsgiTestBase,
     setup_testing_defaults,
 )
 from opentelemetry.test.test_base import TestBase
+from opentelemetry.trace import format_span_id, format_trace_id
 
 
 async def http_app(scope, receive, send):
@@ -287,6 +293,45 @@ def update_expected_user_agent(expected):
         outputs = self.get_all_output()
         self.validate_outputs(outputs, modifiers=[update_expected_user_agent])
 
+    def test_traceresponse_header(self):
+        """Test a traceresponse header is sent when a global propagator is set."""
+
+        orig = get_global_response_propagator()
+        set_global_response_propagator(TraceResponsePropagator())
+
+        app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
+        self.seed_app(app)
+        self.send_default_request()
+
+        # traceresponse header corresponds to http.response.start span
+        span = self.memory_exporter.get_finished_spans()[1]
+        self.assertDictEqual(
+            dict(span.attributes),
+            {
+                SpanAttributes.HTTP_STATUS_CODE: 200,
+                "type": "http.response.start",
+            },
+        )
+
+        response_start, response_body, *_ = self.get_all_output()
+        self.assertEqual(response_body["body"], b"*")
+        self.assertEqual(response_start["status"], 200)
+
+        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.assertListEqual(
+            response_start["headers"],
+            [
+                [b"Content-Type", b"text/plain"],
+                [b"traceresponse", f"{traceresponse}".encode()],
+                [b"access-control-expose-headers", b"traceresponse"],
+            ],
+        )
+
+        set_global_response_propagator(orig)
+
     def test_websocket(self):
         self.scope = {
             "type": "websocket",
@@ -359,6 +404,50 @@ def test_websocket(self):
             self.assertEqual(span.kind, expected["kind"])
             self.assertDictEqual(dict(span.attributes), expected["attributes"])
 
+    def test_websocket_traceresponse_header(self):
+        """Test a traceresponse header is set for websocket messages"""
+
+        orig = get_global_response_propagator()
+        set_global_response_propagator(TraceResponsePropagator())
+
+        self.scope = {
+            "type": "websocket",
+            "http_version": "1.1",
+            "scheme": "ws",
+            "path": "/",
+            "query_string": b"",
+            "headers": [],
+            "client": ("127.0.0.1", 32767),
+            "server": ("127.0.0.1", 80),
+        }
+        app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
+        self.seed_app(app)
+        self.send_input({"type": "websocket.connect"})
+        self.send_input({"type": "websocket.receive", "text": "ping"})
+        self.send_input({"type": "websocket.disconnect"})
+        _, socket_send, *_ = self.get_all_output()
+
+        # traceresponse header corresponds to the 2nd websocket.send span
+        span = self.memory_exporter.get_finished_spans()[3]
+        self.assertDictEqual(
+            dict(span.attributes),
+            {SpanAttributes.HTTP_STATUS_CODE: 200, "type": "websocket.send"},
+        )
+
+        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.assertListEqual(
+            socket_send["headers"],
+            [
+                [b"traceresponse", f"{traceresponse}".encode()],
+                [b"access-control-expose-headers", b"traceresponse"],
+            ],
+        )
+
+        set_global_response_propagator(orig)
+
     def test_lifespan(self):
         self.scope["type"] = "lifespan"
         app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)