Add support for server timing response propagator (#536)

* Add support for server timing response propagator

* Fix lint

* Fix oteltest class name

* Remove intermediate variable

* Revert oteltest name for now

* Hardcode flags in propagator

pyproject.toml

@@ -23,16 +23,16 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
-    "cryptography>=2.0,<=42.0.8",
-    "protobuf>=4.23,<5.0",
-    "opentelemetry-api==1.27.0",
-    "opentelemetry-sdk==1.27.0",
-    "opentelemetry-instrumentation==0.48b0",
-    "opentelemetry-instrumentation-system-metrics==0.48b0",
-    "opentelemetry-semantic-conventions==0.48b0",
-    "opentelemetry-propagator-b3==1.27.0",
-    "opentelemetry-exporter-otlp-proto-grpc==1.27.0",
-    "opentelemetry-exporter-otlp-proto-http==1.27.0"
+  "cryptography>=2.0,<=42.0.8",
+  "protobuf>=4.23,<5.0",
+  "opentelemetry-api==1.27.0",
+  "opentelemetry-sdk==1.27.0",
+  "opentelemetry-propagator-b3==1.27.0",
+  "opentelemetry-exporter-otlp-proto-grpc==1.27.0",
+  "opentelemetry-exporter-otlp-proto-http==1.27.0",
+  "opentelemetry-instrumentation==0.48b0",
+  "opentelemetry-instrumentation-system-metrics==0.48b0",
+  "opentelemetry-semantic-conventions==0.48b0",
 ]
 
 [project.urls]

src/splunk_otel/distro.py

@@ -15,10 +15,18 @@
 import logging
 
 from opentelemetry.instrumentation.distro import BaseDistro
+from opentelemetry.instrumentation.propagators import set_global_response_propagator
 from opentelemetry.instrumentation.system_metrics import SystemMetricsInstrumentor
 from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_HEADERS
 
-from splunk_otel.env import DEFAULTS, OTEL_METRICS_ENABLED, SPLUNK_ACCESS_TOKEN, Env
+from splunk_otel.env import (
+    DEFAULTS,
+    OTEL_METRICS_ENABLED,
+    SPLUNK_ACCESS_TOKEN,
+    SPLUNK_TRACE_RESPONSE_HEADER_ENABLED,
+    Env,
+)
+from splunk_otel.propagator import ServerTimingResponsePropagator
 
 X_SF_TOKEN = "x-sf-token"  # noqa S105
 
@@ -36,6 +44,7 @@ def __init__(self):
     def _configure(self, **kwargs):  # noqa: ARG002
         self.set_env_defaults()
         self.configure_headers()
+        self.set_server_timing_propagator()
 
     def set_env_defaults(self):
         for key, value in DEFAULTS.items():
@@ -53,6 +62,10 @@ def load_instrumentor(self, entry_point, **kwargs):
         else:
             super().load_instrumentor(entry_point, **kwargs)
 
+    def set_server_timing_propagator(self):
+        if self.env.is_true(SPLUNK_TRACE_RESPONSE_HEADER_ENABLED, "true"):
+            set_global_response_propagator(ServerTimingResponsePropagator())
+
 
 def is_system_metrics_instrumentor(entry_point):
     if entry_point.name == "system_metrics":

src/splunk_otel/propagator.py

@@ -0,0 +1,60 @@
+# Copyright Splunk Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from opentelemetry import trace
+from opentelemetry.context.context import Context
+from opentelemetry.instrumentation.propagators import (
+    _HTTP_HEADER_ACCESS_CONTROL_EXPOSE_HEADERS,
+    ResponsePropagator,
+    default_setter,
+)
+from opentelemetry.propagators import textmap
+from opentelemetry.trace import format_span_id, format_trace_id
+
+
+class ServerTimingResponsePropagator(ResponsePropagator):
+    def inject(
+        self,
+        carrier: textmap.CarrierT,
+        context: typing.Optional[Context] = None,  # noqa: FA100
+        setter: textmap.Setter = default_setter,  # type: ignore
+    ) -> None:
+        """
+        Injects SpanContext into the HTTP response carrier.
+
+        e.g.:
+        Server-Timing: traceparent;desc="00-e899d68fca52b66d3facae0bdaf764db-159efb97d6b56568-01"
+        Access-Control-Expose-Headers: Server-Timing
+        """
+        span = trace.get_current_span(context)
+        span_context = span.get_span_context()
+        if span_context == trace.INVALID_SPAN_CONTEXT:
+            return
+
+        header_name = "Server-Timing"
+        trace_id = format_trace_id(span_context.trace_id)
+        span_id = format_span_id(span_context.span_id)
+
+        setter.set(
+            carrier,
+            header_name,
+            f'traceparent;desc="00-{trace_id}-{span_id}-01"',
+        )
+        setter.set(
+            carrier,
+            _HTTP_HEADER_ACCESS_CONTROL_EXPOSE_HEADERS,
+            header_name,
+        )

tests/ott_propagator.py

@@ -0,0 +1,79 @@
+from typing import Mapping, Optional, Sequence
+
+from oteltest import OtelTest, Telemetry
+from ott_lib import project_path
+
+PORT = 8888
+
+HOST = "127.0.0.1"
+
+
+def main():
+    from flask import Flask
+
+    app = Flask(__name__)
+
+    @app.route("/")
+    def home():
+        return "hello"
+
+    app.run(host=HOST, port=PORT)
+
+
+if __name__ == "__main__":
+    main()
+
+
+class OTT(OtelTest):
+    def environment_variables(self) -> Mapping[str, str]:
+        return {
+            "OTEL_SERVICE_NAME": "my-svc",
+        }
+
+    def requirements(self) -> Sequence[str]:
+        return [
+            project_path(),
+            "oteltest",
+            "flask",
+            "opentelemetry-instrumentation-flask==0.48b0",
+            "opentelemetry-exporter-otlp-proto-grpc==1.27.0",
+        ]
+
+    def wrapper_command(self) -> str:
+        return "opentelemetry-instrument"
+
+    def on_start(self) -> Optional[float]:  # noqa: FA100
+        import http.client
+        import time
+
+        time.sleep(6)
+
+        conn = http.client.HTTPConnection(HOST, PORT)
+        conn.request("GET", "/")
+
+        response = conn.getresponse()
+        assert_server_timing_headers_found(response)
+
+        conn.close()
+
+        return 6
+
+    def on_stop(self, tel: Telemetry, stdout: str, stderr: str, returncode: int) -> None:
+        pass
+
+    def is_http(self) -> bool:
+        pass
+
+
+def assert_server_timing_headers_found(response):
+    # Server-Timing: traceparent;desc="00-e899d68fca52b66d3facae0bdaf764db-159efb97d6b56568-01"
+    # Access-Control-Expose-Headers: Server-Timing
+    server_timing_header_found = False
+    access_control_header_found = False
+    for header, _ in response.getheaders():
+        if header == "Server-Timing":
+            server_timing_header_found = True
+        elif header == "Access-Control-Expose-Headers":
+            access_control_header_found = True
+    assert server_timing_header_found
+    assert access_control_header_found

tests/test_distro.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+from opentelemetry.instrumentation.propagators import get_global_response_propagator, set_global_response_propagator
 from splunk_otel.distro import SplunkDistro
 from splunk_otel.env import Env
 
@@ -47,6 +47,27 @@ def test_access_token_whitespace():
     assert "OTEL_EXPORTER_OTLP_HEADERS" not in env_store
 
 
+def test_server_timing_resp_prop_default():
+    set_global_response_propagator(None)
+    env_store = {}
+    configure_distro(env_store)
+    assert get_global_response_propagator()
+
+
+def test_server_timing_resp_prop_true():
+    set_global_response_propagator(None)
+    env_store = {"SPLUNK_TRACE_RESPONSE_HEADER_ENABLED": "true"}
+    configure_distro(env_store)
+    assert get_global_response_propagator()
+
+
+def test_server_timing_resp_prop_false():
+    set_global_response_propagator(None)
+    env_store = {"SPLUNK_TRACE_RESPONSE_HEADER_ENABLED": "false"}
+    configure_distro(env_store)
+    assert get_global_response_propagator() is None
+
+
 def configure_distro(env_store):
     sd = SplunkDistro()
     sd.env = Env(env_store)

tests/test_propagator.py

@@ -0,0 +1,35 @@
+# Copyright Splunk Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from opentelemetry import trace
+from splunk_otel.propagator import ServerTimingResponsePropagator
+
+
+def test_inject():
+    span = trace.NonRecordingSpan(
+        trace.SpanContext(
+            trace_id=1,
+            span_id=2,
+            is_remote=False,
+            trace_flags=trace.TraceFlags(1),
+            trace_state=trace.DEFAULT_TRACE_STATE,
+        ),
+    )
+
+    ctx = trace.set_span_in_context(span)
+    prop = ServerTimingResponsePropagator()
+    carrier = {}
+    prop.inject(carrier, ctx)
+    assert carrier["Access-Control-Expose-Headers"] == "Server-Timing"
+    assert carrier["Server-Timing"] == 'traceparent;desc="00-00000000000000000000000000000001-0000000000000002-01"'