Add support for custom attributes for spans and logs

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

@@ -0,0 +1,159 @@
+use std::collections::HashMap;
+
+use tracing_subscriber::layer::Context;
+use tracing_subscriber::registry::LookupSpan;
+use tracing_subscriber::Layer;
+use tracing_subscriber::Registry;
+
+use super::ROUTER_SPAN_NAME;
+use super::SUBGRAPH_SPAN_NAME;
+use super::SUPERGRAPH_SPAN_NAME;
+
+#[derive(Debug)]
+pub(crate) enum LogAttributes {
+    Router(HashMap<String, String>),
+    Supergraph(HashMap<String, String>),
+    Subgraph(HashMap<String, String>),
+}
+
+impl LogAttributes {
+    pub(crate) fn get_attributes(&self) -> &HashMap<String, String> {
+        match self {
+            LogAttributes::Router(attributes)
+            | LogAttributes::Subgraph(attributes)
+            | LogAttributes::Supergraph(attributes) => attributes,
+        }
+    }
+
+    fn insert(&mut self, span_name: &str, key: String, value: String) {
+        match span_name {
+            ROUTER_SPAN_NAME => {
+                if let Self::Router(attributes) = self {
+                    attributes.insert(key, value);
+                }
+            }
+            SUBGRAPH_SPAN_NAME => {
+                if let Self::Subgraph(attributes) = self {
+                    attributes.insert(key, value);
+                }
+            }
+            SUPERGRAPH_SPAN_NAME => {
+                if let Self::Supergraph(attributes) = self {
+                    attributes.insert(key, value);
+                }
+            }
+            _ => {
+                eprintln!("cannot add custom attributes to this span '{span_name}', it's only available on router/supergraph/subgraph spans");
+            }
+        }
+    }
+    fn extend(&mut self, span_name: &str, val: impl IntoIterator<Item = (String, String)>) {
+        match span_name {
+            ROUTER_SPAN_NAME => {
+                if let Self::Router(attributes) = self {
+                    attributes.extend(val);
+                }
+            }
+            SUBGRAPH_SPAN_NAME => {
+                if let Self::Subgraph(attributes) = self {
+                    attributes.extend(val);
+                }
+            }
+            SUPERGRAPH_SPAN_NAME => {
+                if let Self::Supergraph(attributes) = self {
+                    attributes.extend(val);
+                }
+            }
+            _ => {
+                eprintln!("cannot add custom attributes to this span '{span_name}', it's only available on router/supergraph/subgraph spans");
+            }
+        }
+    }
+}
+
+pub(crate) struct DynAttributeLayer;
+
+impl<S> Layer<S> for DynAttributeLayer
+where
+    S: tracing_core::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
+{
+    fn on_new_span(
+        &self,
+        _attrs: &tracing_core::span::Attributes<'_>,
+        id: &tracing_core::span::Id,
+        ctx: Context<'_, S>,
+    ) {
+        let span = ctx.span(id).expect("Span not found, this is a bug");
+        let custom_attributes = match span.name() {
+            ROUTER_SPAN_NAME => LogAttributes::Router(HashMap::new()),
+            SUBGRAPH_SPAN_NAME => LogAttributes::Subgraph(HashMap::new()),
+            SUPERGRAPH_SPAN_NAME => LogAttributes::Supergraph(HashMap::new()),
+            _ => {
+                return;
+            }
+        };
+        let mut extensions = span.extensions_mut();
+        if extensions.get_mut::<LogAttributes>().is_none() {
+            extensions.insert(custom_attributes);
+        }
+    }
+}
+
+impl DynAttributeLayer {
+    pub(crate) fn new() -> Self {
+        Self {}
+    }
+}
+
+pub(crate) trait DynAttribute {
+    fn set_dyn_attribute(&self, key: String, value: String);
+    fn set_dyn_attributes(&self, attributes: HashMap<String, String>);
+}
+
+impl DynAttribute for ::tracing::Span {
+    fn set_dyn_attribute(&self, key: String, value: String) {
+        self.with_subscriber(move |(id, dispatch)| {
+            if let Some(reg) = dispatch.downcast_ref::<Registry>() {
+                match reg.span(id) {
+                    None => eprintln!("no spanref, this is a bug"),
+                    Some(s) => {
+                        let mut extensions = s.extensions_mut();
+                        match extensions.get_mut::<LogAttributes>() {
+                            Some(attributes) => {
+                                attributes.insert(s.name(), key, value);
+                            }
+                            None => {
+                                eprintln!("no SubgraphRequestLogAttributes, this is a bug");
+                            }
+                        }
+                    }
+                };
+            } else {
+                eprintln!("no Registry, this is a bug");
+            }
+        });
+    }
+
+    fn set_dyn_attributes(&self, attributes: HashMap<String, String>) {
+        self.with_subscriber(move |(id, dispatch)| {
+            if let Some(reg) = dispatch.downcast_ref::<Registry>() {
+                match reg.span(id) {
+                    None => eprintln!("no spanref, this is a bug"),
+                    Some(s) => {
+                        let mut extensions = s.extensions_mut();
+                        match extensions.get_mut::<LogAttributes>() {
+                            Some(registered_attributes) => {
+                                registered_attributes.extend(s.name(), attributes);
+                            }
+                            None => {
+                                eprintln!("no LogAttributes, this is a bug");
+                            }
+                        }
+                    }
+                };
+            } else {
+                eprintln!("no Registry, this is a bug");
+            }
+        });
+    }
+}

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

@@ -17,6 +17,7 @@ use tracing_subscriber::fmt::time::SystemTime;
 use tracing_subscriber::fmt::FmtContext;
 use tracing_subscriber::registry::LookupSpan;
 
+use crate::plugins::telemetry::dynamic_attribute::LogAttributes;
 use crate::plugins::telemetry::reload::IsSampled;
 
 #[derive(Debug, Clone)]
@@ -153,6 +154,50 @@ impl TextFormatter {
         writer.write_char(' ')
     }
 
+    #[inline]
+    fn format_dyn_attributes<S, N>(
+        &self,
+        ctx: &FmtContext<'_, S, N>,
+        writer: &mut Writer<'_>,
+        event: &Event<'_>,
+    ) -> fmt::Result
+    where
+        S: Subscriber + for<'a> LookupSpan<'a>,
+        N: for<'a> FormatFields<'a> + 'static,
+    {
+        let span = event
+            .parent()
+            .and_then(|id| ctx.span(id))
+            .or_else(|| ctx.lookup_current());
+        if let Some(span) = span {
+            let ext = span.extensions();
+            match &ext.get::<LogAttributes>() {
+                Some(dyn_attributes) => {
+                    // TODO: put it in a method
+                    let attributes = dyn_attributes.get_attributes();
+                    // if writer.has_ansi_escapes() {
+                    //     let style = Style::new().dimmed();
+                    //     write!(writer, "{}", style.prefix())?;
+                    //     write!(writer, "[trace_id={trace_id}]")?;
+                    //     write!(writer, "{}", style.suffix())?;
+                    // } else {
+                    let attrs: Vec<String> = attributes
+                        .iter()
+                        .map(|(key, val)| format!("{key}={val}"))
+                        .collect();
+                    write!(writer, "[attributes=[{}]]", attrs.join(", "))?;
+                    // }
+                    writer.write_char(' ')?;
+                }
+                None => {
+                    eprintln!("Unable to find LogAttributes in extensions; this is a bug");
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     #[inline]
     fn format_request_id<S, N>(
         &self,
@@ -216,6 +261,7 @@ where
 
         self.format_level(meta.level(), &mut writer)?;
         self.format_request_id(ctx, &mut writer, event)?;
+        self.format_dyn_attributes(ctx, &mut writer, event)?;
         if self.display_target {
             self.format_target(meta.target(), &mut writer)?;
         }

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

@@ -48,6 +48,10 @@ use tokio::runtime::Handle;
 use tower::BoxError;
 use tower::ServiceBuilder;
 use tower::ServiceExt;
+use tracing_core::callsite::DefaultCallsite;
+use tracing_core::Callsite;
+use tracing_core::Interest;
+use tracing_core::Metadata;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use tracing_subscriber::fmt::format::JsonFields;
 use tracing_subscriber::Layer;
@@ -61,6 +65,7 @@ use self::apollo_exporter::Sender;
 use self::config::Conf;
 use self::config::Sampler;
 use self::config::SamplerOption;
+use self::dynamic_attribute::DynAttribute;
 use self::formatters::text::TextFormatter;
 use self::metrics::apollo::studio::SingleTypeStat;
 use self::metrics::AttributesForwardConf;
@@ -125,6 +130,7 @@ pub(crate) mod apollo;
 pub(crate) mod apollo_exporter;
 pub(crate) mod config;
 mod config_new;
+pub(crate) mod dynamic_attribute;
 mod endpoint;
 pub(crate) mod formatters;
 pub(crate) mod metrics;
@@ -310,8 +316,13 @@ impl Plugin for Telemetry {
                     "apollo_private.http.request_headers" = filter_headers(request.router_request.headers(), &apollo.send_headers).as_str(),
                     "apollo_private.http.response_headers" = field::Empty
                 );
+                // TODO add support of set_attribute for the span
+                // It doesn't need to be prefixed now.
+                span.set_dyn_attribute(String::from("custom_attribute_header"), headers.get("host").and_then(|h| h.to_str().ok()).map(|h| h.to_string()).unwrap_or_default());
+
                 span
             })
+            // TODO add map_future_with_request_data to log the request
             .map_future(move |fut| {
                 let start = Instant::now();
                 let config = config_later.clone();
@@ -321,6 +332,9 @@ impl Plugin for Telemetry {
 
                 async move {
                     let span = Span::current();
+                    ::tracing::info!("coucou");
+
+
                     let response: Result<router::Response, BoxError> = fut.await;
 
                     span.record(
@@ -546,6 +560,20 @@ impl Plugin for Telemetry {
     }
 }
 
+struct LateCallsite {
+    callsite: std::sync::OnceLock<DefaultCallsite>,
+}
+
+impl Callsite for LateCallsite {
+    fn set_interest(&self, interest: Interest) {
+        self.callsite.get().unwrap().set_interest(interest);
+    }
+
+    fn metadata(&self) -> &Metadata<'_> {
+        self.callsite.get().unwrap().metadata()
+    }
+}
+
 impl Telemetry {
     pub(crate) fn activate(&mut self) {
         // Only apply things if we were executing in the context of a vanilla the Apollo executable.
@@ -954,7 +982,7 @@ impl Telemetry {
         subgraph_name: Arc<String>,
         sub_request: &mut Request,
     ) -> Option<CacheAttributes> {
-        let body = dbg!(sub_request.subgraph_request.body_mut());
+        let body = sub_request.subgraph_request.body_mut();
         let hashed_query = hash_request(body);
         let representations = body
             .variables

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

@@ -27,6 +27,7 @@ use tracing_subscriber::EnvFilter;
 use tracing_subscriber::Registry;
 
 use super::config::SamplerOption;
+use super::dynamic_attribute::DynAttributeLayer;
 use super::metrics::span_metrics_exporter::SpanMetricsLayer;
 use crate::axum_factory::utils::REQUEST_SPAN_NAME;
 use crate::metrics::layer::MetricsLayer;
@@ -36,7 +37,7 @@ use crate::plugins::telemetry::formatters::text::TextFormatter;
 use crate::plugins::telemetry::formatters::FilteringFormatter;
 use crate::plugins::telemetry::tracing::reload::ReloadTracer;
 
-pub(crate) type LayeredRegistry = Layered<SpanMetricsLayer, Registry>;
+pub(crate) type LayeredRegistry = Layered<SpanMetricsLayer, Layered<DynAttributeLayer, Registry>>;
 
 pub(super) type LayeredTracer = Layered<
     Filtered<
@@ -118,6 +119,7 @@ pub(crate) fn init_telemetry(log_level: &str) -> Result<()> {
             // Env filter is separate because of https://github.com/tokio-rs/tracing/issues/1629
             // the tracing registry is only created once
             tracing_subscriber::registry()
+                .with(DynAttributeLayer::new())
                 .with(SpanMetricsLayer::default())
                 .with(opentelemetry_layer)
                 .with(fmt_layer)