add configuration for trace_id

NEXT_CHANGELOG.md

@@ -27,6 +27,24 @@ By [@USERNAME](https://github.com/USERNAME) in https://github.com/apollographql/
 ## ❗ BREAKING ❗
 ## 🚀 Features
 
+### Add configuration for trace ID ([Issue #2080](https://github.com/apollographql/router/issues/2080))
+
+If you want to expose in response headers the generated trace ID or the one you provided using propagation headers you can use this configuration:
+
+```yaml title="router.yaml"
+telemetry:
+  tracing:
+    experimental_expose_trace_id:
+      enabled: true # default: false
+      header_name: "my-trace-id" # default: "apollo-trace-id"
+    propagation:
+      custom_header: "x-request-id" # Specify your own trace_id with a custom header in request headers
+```
+
+Using this configuration you will have a response header called `my-trace-id` containing the trace ID. It could help you to debug a specific query if you want to grep your log with this trace id to have more context.
+
+By [@bnjjj](https://github.com/bnjjj) in https://github.com/apollographql/router/pull/2131
+
 ### Add a supergraph configmap option to the helm chart ([PR #2119](https://github.com/apollographql/router/pull/2119))
 
 Adds the capability to create a configmap containing your supergraph schema. Here's an example of how you could make use of this from your values.yaml and with the `helm` install command.

apollo-router/src/axum_factory/utils.rs

@@ -238,7 +238,9 @@ impl<B> MakeSpan<B> for PropagatingMakeSpan {
 
         // If there was no span from the request then it will default to the NOOP span.
         // Attaching the NOOP span has the effect of preventing further tracing.
-        if context.span().span_context().is_valid() {
+        if context.span().span_context().is_valid()
+            || context.span().span_context().trace_id() != opentelemetry::trace::TraceId::INVALID
+        {
             // We have a valid remote span, attach it to the current thread before creating the root span.
             let _context_guard = context.attach();
             tracing::span!(

apollo-router/src/configuration/mod.rs

@@ -207,6 +207,29 @@ impl Configuration {
         self.apollo_plugins
             .plugins
             .insert("include_subgraph_errors".to_string(), json!({"all": true}));
+        // Enable experimental_expose_trace_id
+        self.apollo_plugins
+            .plugins
+            .get_mut("telemetry")
+            .expect("telemetry plugin must be initialized at this point")
+            .as_object_mut()
+            .expect("configuration for telemetry must be an object")
+            .entry("tracing")
+            .and_modify(|e| {
+                e.as_object_mut()
+                    .expect("configuration for telemetry.tracing must be an object")
+                    .entry("experimental_expose_trace_id")
+                    .and_modify(|e| *e = json!({"enabled": true, "header_name": null}))
+                    .or_insert_with(|| json!({"enabled": true, "header_name": null}));
+            })
+            .or_insert_with(|| {
+                json!({
+                    "experimental_expose_trace_id": {
+                        "enabled": true,
+                        "header_name": null
+                    }
+                })
+            });
         self.supergraph.introspection = true;
         self.sandbox.enabled = true;
         self.homepage.enabled = false;

apollo-router/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap

@@ -1702,6 +1702,25 @@ expression: "&schema"
               "additionalProperties": false,
               "nullable": true
             },
+            "experimental_expose_trace_id": {
+              "description": "A way to expose trace id in response headers",
+              "type": "object",
+              "required": [
+                "enabled"
+              ],
+              "properties": {
+                "enabled": {
+                  "description": "Expose the trace_id in response headers",
+                  "type": "boolean"
+                },
+                "header_name": {
+                  "description": "Choose the header name to expose trace_id (default: apollo-trace-id)",
+                  "type": "string",
+                  "nullable": true
+                }
+              },
+              "additionalProperties": false
+            },
             "jaeger": {
               "type": "object",
               "oneOf": [
@@ -1843,6 +1862,11 @@ expression: "&schema"
                   "type": "boolean",
                   "nullable": true
                 },
+                "custom_header": {
+                  "description": "Select a custom header to set your own trace_id (header value must be convertible from hexadecimal to set a correct trace_id)",
+                  "type": "string",
+                  "nullable": true
+                },
                 "datadog": {
                   "type": "boolean",
                   "nullable": true

apollo-router/src/plugins/telemetry/apollo.rs

@@ -12,6 +12,7 @@ use serde::Deserialize;
 use serde::Serialize;
 use url::Url;
 
+use super::config::ExposeTraceId;
 use super::metrics::apollo::studio::ContextualizedStats;
 use super::metrics::apollo::studio::SingleStats;
 use super::metrics::apollo::studio::SingleStatsReport;
@@ -79,6 +80,11 @@ pub(crate) struct Config {
     // The purpose is to allow is to pass this in to the plugin.
     #[schemars(skip)]
     pub(crate) schema_id: String,
+
+    // Skipped because only useful at runtime, it's a copy of the configuration in tracing config
+    #[schemars(skip)]
+    #[serde(skip)]
+    pub(crate) expose_trace_id: ExposeTraceId,
 }
 
 fn apollo_key() -> Option<String> {
@@ -122,6 +128,7 @@ impl Default for Config {
             field_level_instrumentation_sampler: Some(SamplerOption::TraceIdRatioBased(0.01)),
             send_headers: ForwardHeaders::None,
             send_variable_values: ForwardValues::None,
+            expose_trace_id: ExposeTraceId::default(),
         }
     }
 }

apollo-router/src/plugins/telemetry/config.rs

@@ -2,6 +2,7 @@
 use std::borrow::Cow;
 use std::collections::BTreeMap;
 
+use axum::headers::HeaderName;
 use opentelemetry::sdk::Resource;
 use opentelemetry::Array;
 use opentelemetry::KeyValue;
@@ -11,6 +12,7 @@ use serde::Deserialize;
 
 use super::metrics::MetricsAttributesConf;
 use super::*;
+use crate::plugin::serde::deserialize_option_header_name;
 use crate::plugins::telemetry::metrics;
 
 #[derive(thiserror::Error, Debug)]
@@ -78,6 +80,11 @@ pub(crate) struct MetricsCommon {
 #[derive(Clone, Default, Debug, Deserialize, JsonSchema)]
 #[serde(deny_unknown_fields, rename_all = "snake_case")]
 pub(crate) struct Tracing {
+    // TODO: when deleting the `experimental_` prefix, check the usage when enabling dev mode
+    // When deleting, put a #[serde(alias = "experimental_expose_trace_id")] if we don't want to break things
+    /// A way to expose trace id in response headers
+    #[serde(default, rename = "experimental_expose_trace_id")]
+    pub(crate) expose_trace_id: ExposeTraceId,
     pub(crate) propagation: Option<Propagation>,
     pub(crate) trace_config: Option<Trace>,
     pub(crate) otlp: Option<otlp::Config>,
@@ -86,9 +93,24 @@ pub(crate) struct Tracing {
     pub(crate) datadog: Option<tracing::datadog::Config>,
 }
 
+#[derive(Clone, Default, Debug, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields, rename_all = "snake_case")]
+pub(crate) struct ExposeTraceId {
+    /// Expose the trace_id in response headers
+    pub(crate) enabled: bool,
+    /// Choose the header name to expose trace_id (default: apollo-trace-id)
+    #[schemars(with = "Option<String>")]
+    #[serde(deserialize_with = "deserialize_option_header_name")]
+    pub(crate) header_name: Option<HeaderName>,
+}
+
 #[derive(Clone, Default, Debug, Deserialize, JsonSchema)]
 #[serde(deny_unknown_fields, rename_all = "snake_case")]
 pub(crate) struct Propagation {
+    /// Select a custom header to set your own trace_id (header value must be convertible from hexadecimal to set a correct trace_id)
+    #[schemars(with = "Option<String>")]
+    #[serde(deserialize_with = "deserialize_option_header_name")]
+    pub(crate) custom_header: Option<HeaderName>,
     pub(crate) baggage: Option<bool>,
     pub(crate) trace_context: Option<bool>,
     pub(crate) jaeger: Option<bool>,

apollo-router/src/plugins/telemetry/mod.rs

@@ -16,6 +16,7 @@ use ::tracing::info_span;
 use ::tracing::subscriber::set_global_default;
 use ::tracing::Span;
 use ::tracing::Subscriber;
+use axum::headers::HeaderName;
 use futures::future::ready;
 use futures::future::BoxFuture;
 use futures::stream::once;
@@ -27,13 +28,20 @@ use http::HeaderValue;
 use multimap::MultiMap;
 use once_cell::sync::OnceCell;
 use opentelemetry::global;
+use opentelemetry::propagation::text_map_propagator::FieldIter;
+use opentelemetry::propagation::Extractor;
+use opentelemetry::propagation::Injector;
 use opentelemetry::propagation::TextMapPropagator;
 use opentelemetry::sdk::propagation::BaggagePropagator;
 use opentelemetry::sdk::propagation::TextMapCompositePropagator;
 use opentelemetry::sdk::propagation::TraceContextPropagator;
 use opentelemetry::sdk::trace::Builder;
+use opentelemetry::trace::SpanContext;
+use opentelemetry::trace::SpanId;
 use opentelemetry::trace::SpanKind;
 use opentelemetry::trace::TraceContextExt;
+use opentelemetry::trace::TraceFlags;
+use opentelemetry::trace::TraceState;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry::KeyValue;
 use rand::Rng;
@@ -87,6 +95,7 @@ use crate::spaceport::server::ReportSpaceport;
 use crate::spaceport::StatsContext;
 use crate::subgraph::Request;
 use crate::subgraph::Response;
+use crate::tracer::TraceId;
 use crate::Context;
 use crate::ExecutionRequest;
 use crate::ListenAddr;
@@ -111,6 +120,8 @@ const SUBGRAPH_ATTRIBUTES: &str = "apollo_telemetry::subgraph_metrics_attributes
 pub(crate) const STUDIO_EXCLUDE: &str = "apollo_telemetry::studio::exclude";
 pub(crate) const FTV1_DO_NOT_SAMPLE: &str = "apollo_telemetry::studio::ftv1_do_not_sample";
 const DEFAULT_SERVICE_NAME: &str = "apollo-router";
+const GLOBAL_TRACER_NAME: &str = "apollo-router";
+const DEFAULT_EXPOSE_TRACE_ID_HEADER: &str = "apollo-trace-id";
 
 static TELEMETRY_LOADED: OnceCell<bool> = OnceCell::new();
 static TELEMETRY_REFCOUNT: AtomicU8 = AtomicU8::new(0);
@@ -185,13 +196,15 @@ impl Plugin for Telemetry {
         let metrics_sender = self.apollo_metrics_sender.clone();
         let metrics = BasicMetrics::new(&self.meter_provider);
         let config = Arc::new(self.config.clone());
+        let config_map_res_first = config.clone();
         let config_map_res = config.clone();
         ServiceBuilder::new()
             .instrument(Self::supergraph_service_span(
                 self.field_level_instrumentation_ratio,
                 config.apollo.clone().unwrap_or_default(),
             ))
-            .map_response(|resp: SupergraphResponse| {
+            .map_response(move |mut resp: SupergraphResponse| {
+                let config = config_map_res_first.clone();
                 if let Ok(Some(usage_reporting)) =
                     resp.context.get::<_, UsageReporting>(USAGE_REPORTING)
                 {
@@ -201,6 +214,20 @@ impl Plugin for Telemetry {
                         &usage_reporting.stats_report_key.as_str(),
                     );
                 }
+                // To expose trace_id or not
+                let expose_trace_id_header = config.tracing.as_ref().and_then(|t| {
+                    t.expose_trace_id.enabled.then(|| {
+                        t.expose_trace_id.header_name.clone().unwrap_or(
+                            HeaderName::from_static(DEFAULT_EXPOSE_TRACE_ID_HEADER)
+                        )
+                    })
+                });
+                if let (Some(header_name), Some(trace_id)) = (
+                    expose_trace_id_header,
+                    TraceId::maybe_new().and_then(|t| HeaderValue::from_str(&t.to_string()).ok()),
+                ) {
+                    resp.response.headers_mut().append(header_name, trace_id);
+                }
                 resp
             })
             .map_future_with_request_data(
@@ -213,6 +240,7 @@ impl Plugin for Telemetry {
                     let metrics = metrics.clone();
                     let sender = metrics_sender.clone();
                     let start = Instant::now();
+
                     async move {
                         let mut result: Result<SupergraphResponse, BoxError> = fut.await;
                         result = Self::update_otel_metrics(
@@ -360,6 +388,9 @@ impl Telemetry {
             .apollo
             .as_mut()
             .expect("telemetry apollo config must be present");
+        if let Some(tracing_conf) = &config.tracing {
+            apollo.expose_trace_id = tracing_conf.expose_trace_id.clone();
+        }
 
         // If we have key and graph ref but no endpoint we start embedded spaceport
         let spaceport = match apollo {
@@ -395,7 +426,7 @@ impl Telemetry {
             let tracer_provider = Self::create_tracer_provider(&config)?;
 
             let tracer = tracer_provider.versioned_tracer(
-                "apollo-router",
+                GLOBAL_TRACER_NAME,
                 Some(env!("CARGO_PKG_VERSION")),
                 None,
             );
@@ -535,6 +566,11 @@ impl Telemetry {
         if propagation.datadog.unwrap_or_default() || tracing.datadog.is_some() {
             propagators.push(Box::new(opentelemetry_datadog::DatadogPropagator::default()));
         }
+        if let Some(custom_header) = &propagation.custom_header {
+            propagators.push(Box::new(CustomTraceIdPropagator::new(
+                custom_header.to_string(),
+            )));
+        }
 
         TextMapCompositePropagator::new(propagators)
     }
@@ -1245,6 +1281,76 @@ impl ApolloFtv1Handler {
     }
 }
 
+/// CustomTraceIdPropagator to set custom trace_id for our tracing system
+/// coming from headers
+#[derive(Debug)]
+struct CustomTraceIdPropagator {
+    header_name: String,
+    fields: [String; 1],
+}
+
+impl CustomTraceIdPropagator {
+    fn new(header_name: String) -> Self {
+        Self {
+            fields: [header_name.clone()],
+            header_name,
+        }
+    }
+
+    fn extract_span_context(&self, extractor: &dyn Extractor) -> Option<SpanContext> {
+        let trace_id = extractor.get(&self.header_name)?;
+
+        opentelemetry::global::tracer_provider().versioned_tracer(
+            GLOBAL_TRACER_NAME,
+            Some(env!("CARGO_PKG_VERSION")),
+            None,
+        );
+        // extract trace id
+        let trace_id = match opentelemetry::trace::TraceId::from_hex(trace_id) {
+            Ok(trace_id) => trace_id,
+            Err(err) => {
+                ::tracing::error!("cannot generate custom trace_id: {err}");
+                return None;
+            }
+        };
+
+        SpanContext::new(
+            trace_id,
+            SpanId::INVALID,
+            TraceFlags::default().with_sampled(true),
+            true,
+            TraceState::default(),
+        )
+        .into()
+    }
+}
+
+impl TextMapPropagator for CustomTraceIdPropagator {
+    fn inject_context(&self, cx: &opentelemetry::Context, injector: &mut dyn Injector) {
+        let span = cx.span();
+        let span_context = span.span_context();
+        if span_context.is_valid() {
+            let header_value = format!("{}", span_context.trace_id());
+            injector.set(&self.header_name, header_value);
+        }
+    }
+
+    fn extract_with_context(
+        &self,
+        cx: &opentelemetry::Context,
+        extractor: &dyn Extractor,
+    ) -> opentelemetry::Context {
+        cx.with_remote_span_context(
+            self.extract_span_context(extractor)
+                .unwrap_or_else(SpanContext::empty_context),
+        )
+    }
+
+    fn fields(&self) -> FieldIter<'_> {
+        FieldIter::new(self.fields.as_ref())
+    }
+}
+
 //
 // Please ensure that any tests added to the tests module use the tokio multi-threaded test executor.
 //

apollo-router/src/plugins/telemetry/tracing/apollo.rs

@@ -21,11 +21,13 @@ impl TracingConfigurator for Config {
                 schema_id,
                 buffer_size,
                 field_level_instrumentation_sampler,
+                expose_trace_id,
                 ..
             } => {
                 tracing::debug!("configuring exporter to Studio");
 
                 let exporter = apollo_telemetry::Exporter::builder()
+                    .expose_trace_id_config(expose_trace_id.clone())
                     .trace_config(trace_config.clone())
                     .endpoint(endpoint.clone())
                     .apollo_key(key)