add configuration for trace_id

NEXT_CHANGELOG.md

@@ -71,6 +71,26 @@ There is now a separate subcommand for config related operations:
 
 ## 🚀 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_response_trace_id:
+      enabled: true # default: false
+      header_name: "my-trace-id" # default: "apollo-trace-id"
+    propagation:
+      # If you have your own way to generate a trace id and you want to pass it via a custom request header
+      request:
+        header_name: my-trace-id
+```
+
+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 configuration for logging and add more logs 
 
 By default some logs containing sensible data (like request body, response body, headers) are not displayed even if we set the right log level.

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

@@ -213,6 +213,29 @@ impl Configuration {
         self.apollo_plugins
             .plugins
             .insert("include_subgraph_errors".to_string(), json!({"all": true}));
+        // Enable experimental_response_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_response_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_response_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

@@ -1802,6 +1802,25 @@ expression: "&schema"
               "additionalProperties": false,
               "nullable": true
             },
+            "experimental_response_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": [
@@ -1951,6 +1970,21 @@ expression: "&schema"
                   "type": "boolean",
                   "nullable": true
                 },
+                "request": {
+                  "description": "Select a custom request header to set your own trace_id (header value must be convertible from hexadecimal to set a correct trace_id)",
+                  "type": "object",
+                  "required": [
+                    "header_name"
+                  ],
+                  "properties": {
+                    "header_name": {
+                      "description": "Choose the header name to expose trace_id (default: apollo-trace-id)",
+                      "type": "string"
+                    }
+                  },
+                  "additionalProperties": false,
+                  "nullable": true
+                },
                 "trace_context": {
                   "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;
@@ -13,6 +14,8 @@ use serde::Deserialize;
 use super::metrics::MetricsAttributesConf;
 use super::*;
 use crate::configuration::ConfigurationError;
+use crate::plugin::serde::deserialize_header_name;
+use crate::plugin::serde::deserialize_option_header_name;
 use crate::plugin::serde::deserialize_regex;
 use crate::plugins::telemetry::metrics;
 
@@ -82,6 +85,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_response_trace_id")] if we don't want to break things
+    /// A way to expose trace id in response headers
+    #[serde(default, rename = "experimental_response_trace_id")]
+    pub(crate) response_trace_id: ExposeTraceId,
     pub(crate) propagation: Option<Propagation>,
     pub(crate) trace_config: Option<Trace>,
     pub(crate) otlp: Option<otlp::Config>,
@@ -246,16 +254,38 @@ impl Default for LoggingFormat {
     }
 }
 
+#[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 request header to set your own trace_id (header value must be convertible from hexadecimal to set a correct trace_id)
+    pub(crate) request: Option<PropagationRequestTraceId>,
     pub(crate) baggage: Option<bool>,
     pub(crate) trace_context: Option<bool>,
     pub(crate) jaeger: Option<bool>,
     pub(crate) datadog: Option<bool>,
     pub(crate) zipkin: Option<bool>,
 }
 
+#[derive(Clone, Debug, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields, rename_all = "snake_case")]
+pub(crate) struct PropagationRequestTraceId {
+    /// Choose the header name to expose trace_id (default: apollo-trace-id)
+    #[schemars(with = "String")]
+    #[serde(deserialize_with = "deserialize_header_name")]
+    pub(crate) header_name: HeaderName,
+}
+
 #[derive(Default, Debug, Clone, Deserialize, JsonSchema)]
 #[serde(deny_unknown_fields)]
 #[non_exhaustive]