TSH-18197 - Propagate trace context on websocket upgrade request (copy #8739)

.changesets/fix_websocket_trace_propagation.md

@@ -0,0 +1,10 @@
+### Fix: Propagate trace context on websocket upgrade request ([PR #8739](https://github.com/apollographql/router/pull/8739))
+
+GraphQL subscriptions in Apollo Router quite often leverage WebSocket connections to subgraphs. Currently, distributed trace context is not propagated when establishing these WebSocket connections, breaking trace continuity between the router and subgraph services. This results in disconnected traces in observability platforms like Datadog.
+
+The Solution: Inject trace propagation headers into the initial HTTP upgrade request that establishes the WebSocket connection.
+
+WebSocket trace propagation happens exclusively during the HTTP handshake. Once the WebSocket connection is established, headers cannot be added to individual messages—the WebSocket protocol does not support per-message headers.
+
+
+By [@theJC](https://github.com/theJC) in https://github.com/apollographql/router/pull/8739

apollo-router/src/plugins/subscription/subgraph.rs

@@ -35,6 +35,8 @@ use crate::plugins::subscription::WebSocketConfiguration;
 use crate::plugins::telemetry::config_new::events::log_event;
 use crate::plugins::telemetry::config_new::subgraph::events::SubgraphEventRequest;
 use crate::plugins::telemetry::consts::SUBGRAPH_REQUEST_SPAN_NAME;
+use crate::plugins::telemetry::otel::span_ext::OpenTelemetrySpanExt;
+use crate::plugins::telemetry::reload::otel::prepare_context;
 use crate::protocols::websocket::GraphqlWebSocket;
 use crate::protocols::websocket::convert_websocket_stream;
 use crate::services::OperationKind;
@@ -445,6 +447,14 @@ fn get_websocket_request(
     parts.headers.extend(request.headers_mut().drain());
     *request.headers_mut() = parts.headers;
 
+    // Inject trace propagation headers into the WebSocket upgrade request
+    opentelemetry::global::get_text_map_propagator(|propagator| {
+        propagator.inject_context(
+            &prepare_context(tracing::Span::current().context()),
+            &mut crate::otel_compat::HeaderInjector(request.headers_mut()),
+        );
+    });
+
     Ok(request)
 }
 

apollo-router/tests/integration/subscriptions/mod.rs

@@ -28,6 +28,7 @@ use wiremock::ResponseTemplate;
 use wiremock::matchers::method;
 
 pub mod callback;
+pub mod trace_propagation;
 pub mod ws_passthrough;
 
 #[derive(Clone)]