Add cost information to protobuf traces

.changesets/maint_bryn_demand_control_studio_traces.md

@@ -0,0 +1,5 @@
+### Add cost information to protobuf traces ([PR #5430](https://github.com/apollographql/router/pull/5430))
+
+If `experimental_demand_control` is enabled, cost information for queries is now exported on Apollo protobuf traces for display in studio.
+
+By [@BrynCooke](https://github.com/BrynCooke) in https://github.com/apollographql/router/pull/5430

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

@@ -43,6 +43,7 @@ pub(crate) struct CostContext {
     pub(crate) estimated: f64,
     pub(crate) actual: f64,
     pub(crate) result: &'static str,
+    pub(crate) strategy: &'static str,
 }
 
 impl Default for CostContext {
@@ -51,6 +52,7 @@ impl Default for CostContext {
             estimated: 0.0,
             actual: 0.0,
             result: "COST_OK",
+            strategy: "COST_STRATEGY_UNKNOWN",
         }
     }
 }

apollo-router/src/plugins/demand_control/strategy/static_estimated.rs

@@ -22,6 +22,7 @@ impl StrategyImpl for StaticEstimated {
             .and_then(|cost| {
                 request.context.extensions().with_lock(|mut lock| {
                     let cost_result = lock.get_or_default_mut::<CostContext>();
+                    cost_result.strategy = "static_estimated";
                     cost_result.estimated = cost;
                     if cost > self.max {
                         Err(

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

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use opentelemetry::metrics::MeterProvider;
+use opentelemetry_api::Key;
 use opentelemetry_api::KeyValue;
 use parking_lot::Mutex;
 use schemars::JsonSchema;
@@ -10,6 +11,7 @@ use tower::BoxError;
 use super::instruments::Increment;
 use crate::metrics;
 use crate::plugins::demand_control::CostContext;
+use crate::plugins::telemetry::config::AttributeValue;
 use crate::plugins::telemetry::config_new::attributes::SupergraphAttributes;
 use crate::plugins::telemetry::config_new::conditions::Condition;
 use crate::plugins::telemetry::config_new::extendable::Extendable;
@@ -24,6 +26,15 @@ use crate::services::supergraph::Request;
 use crate::services::supergraph::Response;
 use crate::Context;
 
+pub(crate) const APOLLO_PRIVATE_COST_ESTIMATED: Key =
+    Key::from_static_str("apollo_private.cost.estimated");
+pub(crate) const APOLLO_PRIVATE_COST_ACTUAL: Key =
+    Key::from_static_str("apollo_private.cost.actual");
+pub(crate) const APOLLO_PRIVATE_COST_STRATEGY: Key =
+    Key::from_static_str("apollo_private.cost.strategy");
+pub(crate) const APOLLO_PRIVATE_COST_RESULT: Key =
+    Key::from_static_str("apollo_private.cost.result");
+
 static COST_ESTIMATED: &str = "cost.estimated";
 static COST_ACTUAL: &str = "cost.actual";
 static COST_DELTA: &str = "cost.delta";
@@ -271,6 +282,29 @@ pub(crate) enum CostValue {
     Result,
 }
 
+pub(crate) fn add_cost_attributes(context: &Context, custom_attributes: &mut Vec<KeyValue>) {
+    context.extensions().with_lock(|c| {
+        if let Some(cost) = c.get::<CostContext>().cloned() {
+            custom_attributes.push(KeyValue::new(
+                APOLLO_PRIVATE_COST_ESTIMATED.clone(),
+                AttributeValue::I64(cost.estimated as i64),
+            ));
+            custom_attributes.push(KeyValue::new(
+                APOLLO_PRIVATE_COST_ACTUAL.clone(),
+                AttributeValue::I64(cost.actual as i64),
+            ));
+            custom_attributes.push(KeyValue::new(
+                APOLLO_PRIVATE_COST_RESULT.clone(),
+                AttributeValue::String(cost.result.into()),
+            ));
+            custom_attributes.push(KeyValue::new(
+                APOLLO_PRIVATE_COST_STRATEGY.clone(),
+                AttributeValue::String(cost.strategy.into()),
+            ));
+        }
+    });
+}
+
 #[cfg(test)]
 mod test {
     use crate::context::OPERATION_NAME;

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

@@ -18,7 +18,7 @@ pub(crate) mod attributes;
 pub(crate) mod conditions;
 
 mod conditional;
-mod cost;
+pub(crate) mod cost;
 pub(crate) mod events;
 mod experimental_when_header;
 pub(crate) mod extendable;

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

@@ -92,6 +92,7 @@ use crate::plugins::telemetry::apollo_exporter::proto::reports::StatsContext;
 use crate::plugins::telemetry::config::AttributeValue;
 use crate::plugins::telemetry::config::MetricsCommon;
 use crate::plugins::telemetry::config::TracingCommon;
+use crate::plugins::telemetry::config_new::cost::add_cost_attributes;
 use crate::plugins::telemetry::config_new::graphql::GraphQLInstruments;
 use crate::plugins::telemetry::config_new::instruments::SupergraphInstruments;
 use crate::plugins::telemetry::dynamic_attribute::SpanDynAttribute;
@@ -588,15 +589,16 @@ impl Plugin for Telemetry {
 
                     (req.context.clone(), custom_instruments, custom_attributes, supergraph_events, custom_graphql_instruments)
                 },
-                move |(ctx, custom_instruments, custom_attributes, supergraph_events, custom_graphql_instruments): (Context, SupergraphInstruments, Vec<KeyValue>, SupergraphEvents, GraphQLInstruments), fut| {
+                move |(ctx, custom_instruments, mut custom_attributes, supergraph_events, custom_graphql_instruments): (Context, SupergraphInstruments, Vec<KeyValue>, SupergraphEvents, GraphQLInstruments), fut| {
                     let config = config_map_res.clone();
                     let sender = metrics_sender.clone();
                     let start = Instant::now();
 
                     async move {
                         let span = Span::current();
-                        span.set_span_dyn_attributes(custom_attributes);
                         let mut result: Result<SupergraphResponse, BoxError> = fut.await;
+                        add_cost_attributes(&ctx, &mut custom_attributes);
+                        span.set_span_dyn_attributes(custom_attributes);
                         match &result {
                             Ok(resp) => {
                                 span.set_span_dyn_attributes(config.instrumentation.spans.supergraph.attributes.on_response(resp));

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

@@ -55,6 +55,7 @@ use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::query_pla
 use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::query_plan_node::SequenceNode;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::Details;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::Http;
+use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::Limits;
 use crate::plugins::telemetry::apollo_exporter::proto::reports::trace::QueryPlanNode;
 use crate::plugins::telemetry::apollo_exporter::ApolloExporter;
 use crate::plugins::telemetry::apollo_otlp_exporter::ApolloOtlpExporter;
@@ -96,6 +97,14 @@ const APOLLO_PRIVATE_HTTP_RESPONSE_HEADERS: Key =
     Key::from_static_str("apollo_private.http.response_headers");
 pub(crate) const APOLLO_PRIVATE_OPERATION_SIGNATURE: Key =
     Key::from_static_str("apollo_private.operation_signature");
+pub(crate) const APOLLO_PRIVATE_COST_ESTIMATED: Key =
+    Key::from_static_str("apollo_private.cost.estimated");
+pub(crate) const APOLLO_PRIVATE_COST_ACTUAL: Key =
+    Key::from_static_str("apollo_private.cost.actual");
+pub(crate) const APOLLO_PRIVATE_COST_STRATEGY: Key =
+    Key::from_static_str("apollo_private.cost.strategy");
+pub(crate) const APOLLO_PRIVATE_COST_RESULT: Key =
+    Key::from_static_str("apollo_private.cost.result");
 pub(crate) const APOLLO_PRIVATE_FTV1: Key = Key::from_static_str("apollo_private.ftv1");
 const PATH: Key = Key::from_static_str("graphql.path");
 const SUBGRAPH_NAME: Key = Key::from_static_str("apollo.subgraph.name");
@@ -110,7 +119,7 @@ pub(crate) const OPERATION_SUBTYPE: Key = Key::from_static_str("apollo_private.o
 const EXT_TRACE_ID: Key = Key::from_static_str("trace_id");
 
 /// The set of attributes to include when sending to the Apollo Reports protocol.
-const REPORTS_INCLUDE_ATTRS: [Key; 18] = [
+const REPORTS_INCLUDE_ATTRS: [Key; 22] = [
     APOLLO_PRIVATE_REQUEST,
     APOLLO_PRIVATE_DURATION_NS_KEY,
     APOLLO_PRIVATE_SENT_TIME_OFFSET,
@@ -119,6 +128,10 @@ const REPORTS_INCLUDE_ATTRS: [Key; 18] = [
     APOLLO_PRIVATE_HTTP_RESPONSE_HEADERS,
     APOLLO_PRIVATE_OPERATION_SIGNATURE,
     APOLLO_PRIVATE_FTV1,
+    APOLLO_PRIVATE_COST_STRATEGY,
+    APOLLO_PRIVATE_COST_RESULT,
+    APOLLO_PRIVATE_COST_ESTIMATED,
+    APOLLO_PRIVATE_COST_ACTUAL,
     PATH,
     SUBGRAPH_NAME,
     CLIENT_NAME_KEY,
@@ -261,6 +274,7 @@ enum TreeData {
         operation_signature: String,
         operation_name: String,
         variables_json: HashMap<String, String>,
+        limits: Option<Limits>,
     },
     QueryPlanNode(QueryPlanNode),
     DeferPrimary(DeferNodePrimary),
@@ -395,13 +409,15 @@ impl Exporter {
                     operation_signature,
                     operation_name,
                     variables_json,
+                    limits,
                 } => {
                     root_trace.field_execution_weight = self.field_execution_weight;
                     root_trace.signature = operation_signature;
                     root_trace.details = Some(Details {
                         variables_json,
                         operation_name,
                     });
+                    root_trace.limits = limits;
                     results.push(root_trace.clone());
                 }
                 TreeData::Execution(operation_type) => {
@@ -572,6 +588,35 @@ impl Exporter {
                 )]
             }
             SUPERGRAPH_SPAN_NAME => {
+                let limits = span
+                    .attributes
+                    .get(&APOLLO_PRIVATE_COST_RESULT)
+                    .and_then(extract_string)
+                    .map(|result| {
+                        Limits {
+                            result,
+                            strategy: span
+                                .attributes
+                                .get(&APOLLO_PRIVATE_COST_STRATEGY)
+                                .and_then(extract_string)
+                                .unwrap_or_default(),
+                            cost_estimated: span
+                                .attributes
+                                .get(&APOLLO_PRIVATE_COST_ESTIMATED)
+                                .and_then(extract_u64)
+                                .unwrap_or_default(),
+                            cost_actual: span
+                                .attributes
+                                .get(&APOLLO_PRIVATE_COST_ACTUAL)
+                                .and_then(extract_u64)
+                                .unwrap_or_default(),
+                            // Not extracted yet
+                            depth: 0,
+                            height: 0,
+                            alias_count: 0,
+                            root_field_count: 0,
+                        }
+                    });
                 //Currently some data is in the supergraph span as we don't have the a request hook in plugin.
                 child_nodes.push(TreeData::Supergraph {
                     operation_signature: span
@@ -589,6 +634,7 @@ impl Exporter {
                         .get(&APOLLO_PRIVATE_GRAPHQL_VARIABLES)
                         .and_then(extract_json)
                         .unwrap_or_default(),
+                    limits,
                 });
                 child_nodes
             }
@@ -751,6 +797,7 @@ impl Exporter {
                         .and_then(extract_string)
                         .unwrap_or_default(),
                     variables_json: HashMap::new(),
+                    limits: None,
                 });
 
                 child_nodes.push(TreeData::Execution(
@@ -821,6 +868,18 @@ pub(crate) fn extract_i64(v: &Value) -> Option<i64> {
     }
 }
 
+pub(crate) fn extract_u64(v: &Value) -> Option<u64> {
+    if let Value::I64(v) = v {
+        if *v > 0 {
+            Some(*v as u64)
+        } else {
+            None
+        }
+    } else {
+        None
+    }
+}
+
 pub(crate) fn extract_ftv1_trace_with_error_count(
     v: &Value,
     error_config: &ErrorConfiguration,