keep tracer header tags on additional endpoints

Author: duncanista

bottlecap/src/traces/trace_aggregator.rs

@@ -1,4 +1,5 @@
 use libdd_trace_utils::send_data::SendDataBuilder;
+use libdd_trace_utils::trace_utils::TracerHeaderTags;
 use std::collections::VecDeque;
 
 /// Maximum content size per payload uncompressed in bytes,
@@ -7,16 +8,71 @@ use std::collections::VecDeque;
 /// <https://github.com/DataDog/datadog-agent/blob/9d57c10a9eeb3916e661d35dbd23c6e36395a99d/pkg/trace/writer/trace.go#L27-L31>
 pub const MAX_CONTENT_SIZE_BYTES: usize = 3_200_000;
 
+/// Owned version of `TracerHeaderTags<'a>` so it can be stored across async
+/// boundaries without lifetime issues.
+pub struct OwnedTracerHeaderTags {
+    pub lang: String,
+    pub lang_version: String,
+    pub lang_interpreter: String,
+    pub lang_vendor: String,
+    pub tracer_version: String,
+    pub container_id: String,
+    pub client_computed_top_level: bool,
+    pub client_computed_stats: bool,
+    pub dropped_p0_traces: usize,
+    pub dropped_p0_spans: usize,
+}
+
+impl From<TracerHeaderTags<'_>> for OwnedTracerHeaderTags {
+    fn from(tags: TracerHeaderTags<'_>) -> Self {
+        Self {
+            lang: tags.lang.to_string(),
+            lang_version: tags.lang_version.to_string(),
+            lang_interpreter: tags.lang_interpreter.to_string(),
+            lang_vendor: tags.lang_vendor.to_string(),
+            tracer_version: tags.tracer_version.to_string(),
+            container_id: tags.container_id.to_string(),
+            client_computed_top_level: tags.client_computed_top_level,
+            client_computed_stats: tags.client_computed_stats,
+            dropped_p0_traces: tags.dropped_p0_traces,
+            dropped_p0_spans: tags.dropped_p0_spans,
+        }
+    }
+}
+
+impl OwnedTracerHeaderTags {
+    #[must_use] 
+    pub fn to_tracer_header_tags(&self) -> TracerHeaderTags<'_> {
+        TracerHeaderTags {
+            lang: &self.lang,
+            lang_version: &self.lang_version,
+            lang_interpreter: &self.lang_interpreter,
+            lang_vendor: &self.lang_vendor,
+            tracer_version: &self.tracer_version,
+            container_id: &self.container_id,
+            client_computed_top_level: self.client_computed_top_level,
+            client_computed_stats: self.client_computed_stats,
+            dropped_p0_traces: self.dropped_p0_traces,
+            dropped_p0_spans: self.dropped_p0_spans,
+        }
+    }
+}
+
 // Bundle SendDataBuilder with payload size because SendDataBuilder doesn't
 // expose a getter for the size
 pub struct SendDataBuilderInfo {
     pub builder: SendDataBuilder,
     pub size: usize,
+    pub header_tags: OwnedTracerHeaderTags,
 }
 
 impl SendDataBuilderInfo {
-    pub fn new(builder: SendDataBuilder, size: usize) -> Self {
-        Self { builder, size }
+    pub fn new(builder: SendDataBuilder, size: usize, header_tags: OwnedTracerHeaderTags) -> Self {
+        Self {
+            builder,
+            size,
+            header_tags,
+        }
     }
 }
 
@@ -25,7 +81,7 @@ impl SendDataBuilderInfo {
 pub struct TraceAggregator {
     queue: VecDeque<SendDataBuilderInfo>,
     max_content_size_bytes: usize,
-    buffer: Vec<SendDataBuilder>,
+    buffer: Vec<SendDataBuilderInfo>,
 }
 
 impl Default for TraceAggregator {
@@ -55,7 +111,7 @@ impl TraceAggregator {
     }
 
     /// Returns a batch of trace payloads, subject to the max content size.
-    pub fn get_batch(&mut self) -> Vec<SendDataBuilder> {
+    pub fn get_batch(&mut self) -> Vec<SendDataBuilderInfo> {
         let mut batch_size = 0;
 
         // Fill the batch
@@ -70,7 +126,7 @@ impl TraceAggregator {
                     break;
                 }
                 batch_size += payload_size;
-                self.buffer.push(payload_info.builder);
+                self.buffer.push(payload_info);
             } else {
                 break;
             }
@@ -95,8 +151,8 @@ mod tests {
 
     use super::*;
 
-    fn make_builder(size: usize) -> SendDataBuilder {
-        let tracer_header_tags = TracerHeaderTags {
+    fn make_header_tags() -> TracerHeaderTags<'static> {
+        TracerHeaderTags {
             lang: "lang",
             lang_version: "lang_version",
             lang_interpreter: "lang_interpreter",
@@ -107,21 +163,26 @@ mod tests {
             client_computed_stats: true,
             dropped_p0_traces: 0,
             dropped_p0_spans: 0,
-        };
-        SendDataBuilder::new(
+        }
+    }
+
+    fn make_builder_info(size: usize) -> SendDataBuilderInfo {
+        let tracer_header_tags = make_header_tags();
+        let builder = SendDataBuilder::new(
             size,
             TracerPayloadCollection::V07(Vec::new()),
-            tracer_header_tags,
+            tracer_header_tags.clone(),
             &Endpoint::from_slice("localhost"),
-        )
+        );
+        SendDataBuilderInfo::new(builder, size, OwnedTracerHeaderTags::from(tracer_header_tags))
     }
 
     #[test]
     fn test_add() {
         let mut aggregator = TraceAggregator::default();
         let size = 1;
 
-        aggregator.add(SendDataBuilderInfo::new(make_builder(size), size));
+        aggregator.add(make_builder_info(size));
         assert_eq!(aggregator.queue.len(), 1);
     }
 
@@ -130,7 +191,7 @@ mod tests {
         let mut aggregator = TraceAggregator::default();
         let size = 1;
 
-        aggregator.add(SendDataBuilderInfo::new(make_builder(size), size));
+        aggregator.add(make_builder_info(size));
         assert_eq!(aggregator.queue.len(), 1);
         let batch = aggregator.get_batch();
         assert_eq!(batch.len(), 1);
@@ -142,9 +203,9 @@ mod tests {
         let size = 1;
 
         // Add 3 payloads
-        aggregator.add(SendDataBuilderInfo::new(make_builder(size), size));
-        aggregator.add(SendDataBuilderInfo::new(make_builder(size), size));
-        aggregator.add(SendDataBuilderInfo::new(make_builder(size), size));
+        aggregator.add(make_builder_info(size));
+        aggregator.add(make_builder_info(size));
+        aggregator.add(make_builder_info(size));
 
         // The batch should only contain the first 2 payloads
         let first_batch = aggregator.get_batch();

bottlecap/src/traces/trace_aggregator_service.rs

@@ -1,4 +1,3 @@
-use libdd_trace_utils::send_data::SendDataBuilder;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, error};
 
@@ -8,7 +7,7 @@ use crate::traces::trace_aggregator::{
 
 pub enum AggregatorCommand {
     InsertPayload(Box<SendDataBuilderInfo>),
-    GetBatches(oneshot::Sender<Vec<Vec<SendDataBuilder>>>),
+    GetBatches(oneshot::Sender<Vec<Vec<SendDataBuilderInfo>>>),
     Clear,
     Shutdown,
 }
@@ -27,7 +26,7 @@ impl AggregatorHandle {
             .send(AggregatorCommand::InsertPayload(Box::new(payload_info)))
     }
 
-    pub async fn get_batches(&self) -> Result<Vec<Vec<SendDataBuilder>>, String> {
+    pub async fn get_batches(&self) -> Result<Vec<Vec<SendDataBuilderInfo>>, String> {
         let (response_tx, response_rx) = oneshot::channel();
         self.tx
             .send(AggregatorCommand::GetBatches(response_tx))
@@ -105,9 +104,11 @@ impl AggregatorService {
 #[allow(clippy::unwrap_used)]
 mod tests {
     use super::*;
+    use crate::traces::trace_aggregator::OwnedTracerHeaderTags;
     use libdd_common::Endpoint;
     use libdd_trace_utils::{
-        trace_utils::TracerHeaderTags, tracer_payload::TracerPayloadCollection,
+        send_data::SendDataBuilder, trace_utils::TracerHeaderTags,
+        tracer_payload::TracerPayloadCollection,
     };
 
     #[tokio::test]
@@ -131,6 +132,7 @@ mod tests {
             dropped_p0_spans: 0,
         };
         let size = 1;
+        let owned_tags = OwnedTracerHeaderTags::from(tracer_header_tags.clone());
         let payload = SendDataBuilder::new(
             size,
             TracerPayloadCollection::V07(Vec::new()),
@@ -139,7 +141,7 @@ mod tests {
         );
 
         handle
-            .insert_payload(SendDataBuilderInfo::new(payload, size))
+            .insert_payload(SendDataBuilderInfo::new(payload, size, owned_tags))
             .unwrap();
 
         let batches = handle.get_batches().await.unwrap();

bottlecap/src/traces/trace_flusher.rs

@@ -5,9 +5,8 @@ use dogstatsd::api_key::ApiKeyFactory;
 use libdd_common::Endpoint;
 use libdd_trace_utils::{
     config_utils::trace_intake_url_prefixed,
-    send_data::{SendData, SendDataBuilder},
+    send_data::SendData,
     trace_utils::{self},
-    tracer_header_tags::TracerHeaderTags,
     tracer_payload::TracerPayloadCollection,
 };
 use std::str::FromStr;
@@ -118,25 +117,29 @@ impl TraceFlusher {
         let mut batch_tasks = JoinSet::new();
 
         for trace_builders in all_batches {
-            let traces: Vec<_> = trace_builders
+            let traces_with_tags: Vec<_> = trace_builders
                 .into_iter()
-                .map(|builder| builder.with_api_key(api_key.as_str()))
-                .map(SendDataBuilder::build)
+                .map(|info| {
+                    let trace = info.builder.with_api_key(api_key.as_str()).build();
+                    (trace, info.header_tags)
+                })
                 .collect();
 
             // Send to ADDITIONAL endpoints for dual-shipping.
             // Construct separate SendData objects per endpoint by cloning the inner
             // V07 payload data (TracerPayload is Clone, but SendData is not).
             for endpoint in self.additional_endpoints.clone() {
-                let additional_traces: Vec<_> = traces
+                let additional_traces: Vec<_> = traces_with_tags
                     .iter()
-                    .filter_map(|trace| match trace.get_payloads() {
+                    .filter_map(|(trace, tags)| match trace.get_payloads() {
                         TracerPayloadCollection::V07(payloads) => Some(SendData::new(
                             trace.len(),
                             TracerPayloadCollection::V07(payloads.clone()),
-                            TracerHeaderTags::default(),
+                            tags.to_tracer_header_tags(),
                             &endpoint,
                         )),
+                        // All payloads in the extension are V07 (produced by
+                        // collect_pb_trace_chunks), so this branch is unreachable.
                         _ => None,
                     })
                     .collect();
@@ -146,6 +149,7 @@ impl TraceFlusher {
             }
 
             // Send to PRIMARY endpoint (moves traces into the task).
+            let traces: Vec<_> = traces_with_tags.into_iter().map(|(t, _)| t).collect();
             let client_clone = http_client.clone();
             batch_tasks.spawn(async move { Self::send_traces(traces, client_clone).await });
         }

bottlecap/src/traces/trace_processor.rs

@@ -31,8 +31,8 @@ use tokio::sync::mpsc::Sender;
 use tokio::sync::mpsc::error::SendError;
 use tracing::{debug, error};
 
-use super::stats_generator::StatsGenerator;
-use super::trace_aggregator::SendDataBuilderInfo;
+use crate::traces::stats_generator::StatsGenerator;
+use crate::traces::trace_aggregator::{OwnedTracerHeaderTags, SendDataBuilderInfo};
 
 #[derive(Clone)]
 #[allow(clippy::module_name_repetitions)]
@@ -372,6 +372,8 @@ impl TraceProcessor for ServerlessTraceProcessor {
             }
         };
 
+        let owned_header_tags = OwnedTracerHeaderTags::from(header_tags.clone());
+
         // Move original payload into builder (no clone needed)
         let builder = SendDataBuilder::new(body_size, payload, header_tags, &endpoint)
             .with_compression(Compression::Zstd(config.apm_config_compression_level))
@@ -383,7 +385,7 @@ impl TraceProcessor for ServerlessTraceProcessor {
             ));
 
         (
-            SendDataBuilderInfo::new(builder, body_size),
+            SendDataBuilderInfo::new(builder, body_size, owned_header_tags),
             payloads_for_stats,
         )
     }