feat(traces): [SVLS-8734] respect Datadog-Client-Computed-Stats header

bottlecap/src/lifecycle/invocation/context.rs

@@ -43,6 +43,11 @@ pub struct Context {
     /// tracing.
     ///
     pub extracted_span_context: Option<SpanContext>,
+    /// Whether the tracer has already computed stats for this invocation.
+    ///
+    /// Set from the `Datadog-Client-Computed-Stats` header when the tracer's
+    /// placeholder span (`dd-tracer-serverless-span`) is received.
+    pub client_computed_stats: bool,
 }
 
 /// Struct containing the information needed to reparent a span.
@@ -94,6 +99,7 @@ impl Default for Context {
             snapstart_restore_span: None,
             tracer_span: None,
             extracted_span_context: None,
+            client_computed_stats: false,
         }
     }
 }
@@ -508,7 +514,12 @@ impl ContextBuffer {
 
     /// Adds the tracer span to a `Context` in the buffer.
     ///
-    pub fn add_tracer_span(&mut self, request_id: &String, tracer_span: &Span) {
+    pub fn add_tracer_span(
+        &mut self,
+        request_id: &String,
+        tracer_span: &Span,
+        client_computed_stats: bool,
+    ) {
         if let Some(context) = self
             .buffer
             .iter_mut()
@@ -528,6 +539,7 @@ impl ContextBuffer {
                 .extend(tracer_span.metrics.clone());
 
             context.tracer_span = Some(tracer_span.clone());
+            context.client_computed_stats = client_computed_stats;
         } else {
             debug!("Could not add tracer span - context not found");
         }

bottlecap/src/lifecycle/invocation/processor.rs

@@ -610,9 +610,16 @@ impl Processor {
         trace_sender: &Arc<SendingTraceProcessor>,
         context: Context,
     ) {
+        let client_computed_stats = context.client_computed_stats;
         let (traces, body_size) = self.get_ctx_spans(context);
-        self.send_spans(traces, body_size, tags_provider, trace_sender)
-            .await;
+        self.send_spans(
+            traces,
+            body_size,
+            client_computed_stats,
+            tags_provider,
+            trace_sender,
+        )
+        .await;
     }
 
     fn get_ctx_spans(&mut self, context: Context) -> (Vec<Span>, usize) {
@@ -677,7 +684,7 @@ impl Processor {
             let traces = vec![cold_start_span.clone()];
             let body_size = size_of_val(cold_start_span);
 
-            self.send_spans(traces, body_size, tags_provider, trace_sender)
+            self.send_spans(traces, body_size, false, tags_provider, trace_sender)
                 .await;
         }
     }
@@ -689,6 +696,7 @@ impl Processor {
         &mut self,
         traces: Vec<Span>,
         body_size: usize,
+        client_computed_stats: bool,
         tags_provider: &Arc<provider::Provider>,
         trace_sender: &Arc<SendingTraceProcessor>,
     ) {
@@ -701,7 +709,7 @@ impl Processor {
             tracer_version: "",
             container_id: "",
             client_computed_top_level: false,
-            client_computed_stats: false,
+            client_computed_stats,
             dropped_p0_traces: 0,
             dropped_p0_spans: 0,
         };
@@ -1337,9 +1345,10 @@ impl Processor {
     ///
     /// This is used to enrich the invocation span with additional metadata from the tracers
     /// top level span, since we discard the tracer span when we create the invocation span.
-    pub fn add_tracer_span(&mut self, span: &Span) {
+    pub fn add_tracer_span(&mut self, span: &Span, client_computed_stats: bool) {
         if let Some(request_id) = span.meta.get("request_id") {
-            self.context_buffer.add_tracer_span(request_id, span);
+            self.context_buffer
+                .add_tracer_span(request_id, span, client_computed_stats);
         }
     }
 }
@@ -2110,4 +2119,92 @@ mod tests {
             "Should return None when no trace context found"
         );
     }
+
+    /// Verifies that `client_computed_stats` set on a context via `add_tracer_span` is
+    /// propagated all the way through `send_ctx_spans` to the `aws.lambda` payload sent
+    /// to the backend, so the extension does not generate duplicate stats.
+    #[tokio::test]
+    #[allow(clippy::unwrap_used)]
+    async fn test_client_computed_stats_propagated_to_aws_lambda_span() {
+        use crate::traces::stats_concentrator_service::StatsConcentratorService;
+        use crate::traces::stats_generator::StatsGenerator;
+        use libdd_trace_obfuscation::obfuscation_config::ObfuscationConfig;
+        use tokio::sync::mpsc;
+
+        let config = Arc::new(config::Config {
+            apm_dd_url: "https://trace.agent.datadoghq.com".to_string(),
+            ..config::Config::default()
+        });
+        let tags_provider = Arc::new(provider::Provider::new(
+            Arc::clone(&config),
+            LAMBDA_RUNTIME_SLUG.to_string(),
+            &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]),
+        ));
+        let aws_config = Arc::new(AwsConfig {
+            region: "us-east-1".into(),
+            aws_lwa_proxy_lambda_runtime_api: None,
+            function_name: "test-function".into(),
+            sandbox_init_time: Instant::now(),
+            runtime_api: "***".into(),
+            exec_wrapper: None,
+            initialization_type: "on-demand".into(),
+        });
+        let (aggregator_service, aggregator_handle) =
+            AggregatorService::new(EMPTY_TAGS, 1024).expect("failed to create aggregator service");
+        tokio::spawn(aggregator_service.run());
+        let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
+        let mut p = Processor::new(
+            Arc::clone(&tags_provider),
+            Arc::clone(&config),
+            aws_config,
+            aggregator_handle,
+            propagator,
+        );
+
+        let (trace_tx, mut trace_rx) = mpsc::channel(10);
+        let (stats_concentrator_service, stats_concentrator_handle) =
+            StatsConcentratorService::new(Arc::clone(&config));
+        tokio::spawn(stats_concentrator_service.run());
+        let trace_sender = Arc::new(SendingTraceProcessor {
+            appsec: None,
+            processor: Arc::new(trace_processor::ServerlessTraceProcessor {
+                obfuscation_config: Arc::new(
+                    ObfuscationConfig::new().expect("Failed to create ObfuscationConfig"),
+                ),
+            }),
+            trace_tx,
+            stats_generator: Arc::new(StatsGenerator::new(stats_concentrator_handle)),
+        });
+
+        let mut context = Context::from_request_id("req-1");
+        context.invocation_span.trace_id = 1;
+        context.invocation_span.span_id = 2;
+        context.client_computed_stats = true;
+
+        p.send_ctx_spans(&tags_provider, &trace_sender, context)
+            .await;
+
+        let payload = trace_rx.recv().await.expect("expected payload from trace_tx");
+        assert!(
+            payload.header_tags.client_computed_stats,
+            "client_computed_stats must be propagated to the aws.lambda span payload"
+        );
+
+        // Verify _dd.compute_stats is "0" in the built payload tags: client_computed_stats=true
+        // means the tracer has already computed stats, so neither extension nor backend should.
+        let send_data = payload.builder.build();
+        let libdd_trace_utils::tracer_payload::TracerPayloadCollection::V07(payloads) =
+            send_data.get_payloads()
+        else {
+            panic!("expected V07 payload");
+        };
+        for p in payloads {
+            assert_eq!(
+                p.tags
+                    .get(crate::tags::lambda::tags::COMPUTE_STATS_KEY),
+                Some(&"0".to_string()),
+                "_dd.compute_stats must be 0 when client_computed_stats is true"
+            );
+        }
+    }
 }

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -110,6 +110,7 @@ pub enum ProcessorCommand {
     },
     AddTracerSpan {
         span: Box<Span>,
+        client_computed_stats: bool,
     },
     OnOutOfMemoryError {
         timestamp: i64,
@@ -372,10 +373,12 @@ impl InvocationProcessorHandle {
     pub async fn add_tracer_span(
         &self,
         span: Span,
+        client_computed_stats: bool,
     ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
         self.sender
             .send(ProcessorCommand::AddTracerSpan {
                 span: Box::new(span),
+                client_computed_stats,
             })
             .await
     }
@@ -585,8 +588,11 @@ impl InvocationProcessorService {
                     let result = Ok(self.processor.set_cold_start_span_trace_id(trace_id));
                     let _ = response.send(result);
                 }
-                ProcessorCommand::AddTracerSpan { span } => {
-                    self.processor.add_tracer_span(&span);
+                ProcessorCommand::AddTracerSpan {
+                    span,
+                    client_computed_stats,
+                } => {
+                    self.processor.add_tracer_span(&span, client_computed_stats);
                 }
                 ProcessorCommand::OnOutOfMemoryError { timestamp } => {
                     self.processor.on_out_of_memory_error(timestamp);

bottlecap/src/tags/lambda/tags.rs

@@ -39,7 +39,7 @@ const VERSION_KEY: &str = "version";
 const SERVICE_KEY: &str = "service";
 
 // ComputeStatsKey is the tag key indicating whether trace stats should be computed
-const COMPUTE_STATS_KEY: &str = "_dd.compute_stats";
+pub(crate) const COMPUTE_STATS_KEY: &str = "_dd.compute_stats";
 // FunctionTagsKey is the tag key for a function's tags to be set on the top level tracepayload
 const FUNCTION_TAGS_KEY: &str = "_dd.tags.function";
 // TODO(astuyve) decide what to do with the version
@@ -127,12 +127,6 @@ fn tags_from_env(
         tags_map.extend(config.tags.clone());
     }
 
-    // The value of _dd.compute_stats is the opposite of config.compute_trace_stats_on_extension.
-    // "config.compute_trace_stats_on_extension == true" means computing stats on the extension side,
-    // so we set _dd.compute_stats to 0 so stats won't be computed on the backend side.
-    let compute_stats = i32::from(!config.compute_trace_stats_on_extension);
-    tags_map.insert(COMPUTE_STATS_KEY.to_string(), compute_stats.to_string());
-
     tags_map
 }
 
@@ -292,8 +286,7 @@ mod tests {
     fn test_new_from_config() {
         let metadata = HashMap::new();
         let tags = Lambda::new_from_config(Arc::new(Config::default()), &metadata);
-        assert_eq!(tags.tags_map.len(), 3);
-        assert_eq!(tags.tags_map.get(COMPUTE_STATS_KEY).unwrap(), "1");
+        assert_eq!(tags.tags_map.len(), 2);
         let arch = arch_to_platform();
         assert_eq!(
             tags.tags_map.get(ARCHITECTURE_KEY).unwrap(),
@@ -428,7 +421,7 @@ mod tests {
                 (parts[0].to_string(), parts[1].to_string())
             })
             .collect();
-        assert_eq!(fn_tags_map.len(), 14);
+        assert_eq!(fn_tags_map.len(), 13);
         assert_eq!(fn_tags_map.get("key1").unwrap(), "value1");
         assert_eq!(fn_tags_map.get("key2").unwrap(), "value2");
         assert_eq!(fn_tags_map.get(ACCOUNT_ID_KEY).unwrap(), "123456789012");
@@ -470,7 +463,7 @@ mod tests {
                 (parts[0].to_string(), parts[1].to_string())
             })
             .collect();
-        assert_eq!(fn_tags_map.len(), 14);
+        assert_eq!(fn_tags_map.len(), 13);
         assert_eq!(fn_tags_map.get("key1").unwrap(), "value1");
         assert_eq!(fn_tags_map.get("key2").unwrap(), "value2");
         assert_eq!(fn_tags_map.get(ACCOUNT_ID_KEY).unwrap(), "123456789012");

bottlecap/src/traces/trace_agent.rs

@@ -501,7 +501,8 @@ impl TraceAgent {
             );
         }
 
-        let tracer_header_tags = (&parts.headers).into();
+        let tracer_header_tags: libdd_trace_utils::tracer_header_tags::TracerHeaderTags<'_> =
+            (&parts.headers).into();
 
         let (body_size, mut traces): (usize, Vec<Vec<pb::Span>>) = match version {
             ApiVersion::V04 => {
@@ -531,7 +532,6 @@ impl TraceAgent {
                 }
             },
         };
-
         let mut reparenting_info = match invocation_processor_handle.get_reparenting_info().await {
             Ok(info) => info,
             Err(e) => {
@@ -586,7 +586,7 @@ impl TraceAgent {
 
                 if span.resource == INVOCATION_SPAN_RESOURCE
                     && let Err(e) = invocation_processor_handle
-                        .add_tracer_span(span.clone())
+                        .add_tracer_span(span.clone(), tracer_header_tags.client_computed_stats)
                         .await
                 {
                     error!("Failed to add tracer span to processor: {}", e);