chore: pass SendDataBuilderInfo instead of SendData until flush time (#745)

# Background
Right now `SendData` is passed around across channels.

# This PR

Instead of passing `SendData`, pass `SendDataBuilderInfo`, which bundles
`SendDataBuilder` and payload size. Just before flush, call
`SendDataBuilder.build()` to build `SendData`.

# Motivation
DataDog/libdatadog#1140 (comment)
It is suggested that the function `set_api_key()` shouldn't be added on
`SendData`, but should be added on `SendDataBuilder`. Because need to
call `set_api_key()` just before flush, we need to make sure the object
is `SendDataBuilder` instead of `SendData` until flush time.

And because we need payload size in Trace Aggregator, and
`SendDataBuilder` doesn't expose this field, we need to pass it
explicitly along with `SendDataBuilder`.

# Next steps
Update #717
#732 so that
`get_api_key()` is called just before flush.

# Dependency
DataDog/libdatadog#1140

Author: lym953

bottlecap/Cargo.lock

@@ -52,11 +52,11 @@ rustls-native-certs = { version = "0.8.1", optional = true }
 # be found in the clippy.toml file adjacent to this Cargo.toml.
 datadog-protos = { version = "0.1.0", default-features = false, git = "https://github.com/DataDog/saluki/", rev = "c89b58e5784b985819baf11f13f7d35876741222"}
 ddsketch-agent = { version = "0.1.0", default-features = false, git = "https://github.com/DataDog/saluki/" }
-ddcommon = { git = "https://github.com/DataDog/libdatadog", rev = "ca19adc5c782be316210b69510ebb6c8dff12d87" }
-datadog-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev = "ca19adc5c782be316210b69510ebb6c8dff12d87"  }
-datadog-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "ca19adc5c782be316210b69510ebb6c8dff12d87" , features = ["mini_agent"] }
-datadog-trace-normalization = { git = "https://github.com/DataDog/libdatadog",  rev = "ca19adc5c782be316210b69510ebb6c8dff12d87" }
-datadog-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "ca19adc5c782be316210b69510ebb6c8dff12d87"  }
+ddcommon = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2" }
+datadog-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2"  }
+datadog-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2" , features = ["mini_agent"] }
+datadog-trace-normalization = { git = "https://github.com/DataDog/libdatadog",  rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2" }
+datadog-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2"  }
 dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "c3d8ed4f90591c6958921145d485463860307f78", default-features = false }
 datadog-trace-agent = { git = "https://github.com/DataDog/serverless-components", rev = "c3d8ed4f90591c6958921145d485463860307f78" }
 datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "c3d8ed4f90591c6958921145d485463860307f78", default-features = false }

bottlecap/Cargo.toml

@@ -50,7 +50,8 @@ use bottlecap::{
         proxy_flusher::Flusher as ProxyFlusher,
         stats_aggregator::StatsAggregator,
         stats_flusher::{self, StatsFlusher},
-        stats_processor, trace_agent, trace_aggregator,
+        stats_processor, trace_agent,
+        trace_aggregator::{self, SendDataBuilderInfo},
         trace_flusher::{self, ServerlessTraceFlusher, TraceFlusher},
         trace_processor,
     },
@@ -802,7 +803,7 @@ async fn handle_event_bus_event(
     invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
     tags_provider: Arc<TagProvider>,
     trace_processor: Arc<trace_processor::ServerlessTraceProcessor>,
-    trace_agent_channel: Sender<datadog_trace_utils::send_data::SendData>,
+    trace_agent_channel: Sender<SendDataBuilderInfo>,
 ) -> Option<TelemetryEvent> {
     match event {
         Event::Metric(event) => {
@@ -1024,7 +1025,7 @@ fn start_trace_agent(
     invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
     trace_aggregator: Arc<TokioMutex<trace_aggregator::TraceAggregator>>,
 ) -> (
-    Sender<datadog_trace_utils::send_data::SendData>,
+    Sender<SendDataBuilderInfo>,
     Arc<trace_flusher::ServerlessTraceFlusher>,
     Arc<trace_processor::ServerlessTraceProcessor>,
     Arc<stats_flusher::ServerlessStatsFlusher>,
@@ -1146,7 +1147,7 @@ fn start_otlp_agent(
     config: &Arc<Config>,
     tags_provider: Arc<TagProvider>,
     trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
-    trace_tx: Sender<SendData>,
+    trace_tx: Sender<SendDataBuilderInfo>,
 ) -> Option<CancellationToken> {
     if !should_enable_otlp_agent(config) {
         return None;

bottlecap/src/bin/bottlecap/main.rs

@@ -6,7 +6,7 @@ use std::{
 
 use chrono::{DateTime, Utc};
 use datadog_trace_protobuf::pb::Span;
-use datadog_trace_utils::{send_data::SendData, tracer_header_tags};
+use datadog_trace_utils::tracer_header_tags;
 use dogstatsd::aggregator::Aggregator as MetricsAggregator;
 use serde_json::{json, Value};
 use tokio::sync::{mpsc::Sender, watch};
@@ -35,6 +35,7 @@ use crate::{
             },
             DatadogCompositePropagator, Propagator,
         },
+        trace_aggregator::SendDataBuilderInfo,
         trace_processor::{self, TraceProcessor},
     },
 };
@@ -284,7 +285,7 @@ impl Processor {
         status: Status,
         tags_provider: Arc<provider::Provider>,
         trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
-        trace_agent_tx: Sender<SendData>,
+        trace_agent_tx: Sender<SendDataBuilderInfo>,
         timestamp: i64,
     ) {
         // Set the runtime duration metric
@@ -328,7 +329,7 @@ impl Processor {
         status: Status,
         tags_provider: Arc<provider::Provider>,
         trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
-        trace_agent_tx: Sender<SendData>,
+        trace_agent_tx: Sender<SendDataBuilderInfo>,
     ) {
         let context = self.enrich_ctx_at_platform_done(request_id, status);
 
@@ -415,7 +416,7 @@ impl Processor {
         &mut self,
         tags_provider: &Arc<provider::Provider>,
         trace_processor: &Arc<dyn TraceProcessor + Send + Sync>,
-        trace_agent_tx: &Sender<SendData>,
+        trace_agent_tx: &Sender<SendDataBuilderInfo>,
         context: Context,
     ) {
         let mut body_size = std::mem::size_of_val(&context.invocation_span);
@@ -467,7 +468,7 @@ impl Processor {
         &mut self,
         tags_provider: &Arc<provider::Provider>,
         trace_processor: &Arc<dyn TraceProcessor + Send + Sync>,
-        trace_agent_tx: &Sender<SendData>,
+        trace_agent_tx: &Sender<SendDataBuilderInfo>,
     ) {
         if let Some(cold_start_context) = self.context_buffer.get_context_with_cold_start() {
             if let Some(cold_start_span) = &mut cold_start_context.cold_start_span {
@@ -500,7 +501,7 @@ impl Processor {
         body_size: usize,
         tags_provider: &Arc<provider::Provider>,
         trace_processor: &Arc<dyn TraceProcessor + Send + Sync>,
-        trace_agent_tx: &Sender<SendData>,
+        trace_agent_tx: &Sender<SendDataBuilderInfo>,
     ) {
         // todo: figure out what to do here
         let header_tags = tracer_header_tags::TracerHeaderTags {
@@ -516,7 +517,7 @@ impl Processor {
             dropped_p0_spans: 0,
         };
 
-        let send_data: SendData = trace_processor.process_traces(
+        let send_data_builder_info: SendDataBuilderInfo = trace_processor.process_traces(
             self.config.clone(),
             tags_provider.clone(),
             header_tags,
@@ -525,7 +526,7 @@ impl Processor {
             self.inferrer.span_pointers.clone(),
         );
 
-        if let Err(e) = trace_agent_tx.send(send_data).await {
+        if let Err(e) = trace_agent_tx.send(send_data_builder_info).await {
             debug!("Failed to send context spans to agent: {e}");
         }
     }

bottlecap/src/lifecycle/invocation/processor.rs

@@ -5,7 +5,6 @@ use axum::{
     routing::post,
     Router,
 };
-use datadog_trace_utils::send_data::SendData;
 use datadog_trace_utils::trace_utils::TracerHeaderTags as DatadogTracerHeaderTags;
 use serde_json::json;
 use std::net::SocketAddr;
@@ -19,7 +18,7 @@ use crate::{
     http::{extract_request_body, handler_not_found},
     otlp::processor::Processor as OtlpProcessor,
     tags::provider,
-    traces::trace_processor::TraceProcessor,
+    traces::{trace_aggregator::SendDataBuilderInfo, trace_processor::TraceProcessor},
 };
 
 const OTLP_AGENT_HTTP_PORT: u16 = 4318;
@@ -29,15 +28,15 @@ type AgentState = (
     Arc<provider::Provider>,
     OtlpProcessor,
     Arc<dyn TraceProcessor + Send + Sync>,
-    Sender<SendData>,
+    Sender<SendDataBuilderInfo>,
 );
 
 pub struct Agent {
     config: Arc<Config>,
     tags_provider: Arc<provider::Provider>,
     processor: OtlpProcessor,
     trace_processor: Arc<dyn TraceProcessor + Send + Sync>,
-    trace_tx: Sender<SendData>,
+    trace_tx: Sender<SendDataBuilderInfo>,
     port: u16,
     cancel_token: CancellationToken,
 }
@@ -47,7 +46,7 @@ impl Agent {
         config: Arc<Config>,
         tags_provider: Arc<provider::Provider>,
         trace_processor: Arc<dyn TraceProcessor + Send + Sync>,
-        trace_tx: Sender<SendData>,
+        trace_tx: Sender<SendDataBuilderInfo>,
     ) -> Self {
         let port = Self::parse_port(
             &config.otlp_config_receiver_protocols_http_endpoint,
@@ -157,16 +156,7 @@ impl Agent {
 
         let tracer_header_tags: DatadogTracerHeaderTags = (&parts.headers).into();
         let body_size = size_of_val(&traces);
-        let send_data = trace_processor.process_traces(
-            config,
-            tags_provider,
-            tracer_header_tags,
-            traces,
-            body_size,
-            None,
-        );
-
-        if send_data.is_empty() {
+        if body_size == 0 {
             return (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 json!({ "message": "Not sending traces, processor returned empty data" })
@@ -175,7 +165,16 @@ impl Agent {
                 .into_response();
         }
 
-        match trace_tx.send(send_data).await {
+        let send_data_builder = trace_processor.process_traces(
+            config,
+            tags_provider,
+            tracer_header_tags,
+            traces,
+            body_size,
+            None,
+        );
+
+        match trace_tx.send(send_data_builder).await {
             Ok(()) => {
                 debug!("OTLP | Successfully buffered traces to be flushed.");
                 (

bottlecap/src/otlp/agent.rs

@@ -29,12 +29,13 @@ use crate::{
     tags::provider,
     traces::{
         proxy_aggregator::{self, ProxyRequest},
-        stats_aggregator, stats_processor, trace_aggregator, trace_processor,
-        INVOCATION_SPAN_RESOURCE,
+        stats_aggregator, stats_processor,
+        trace_aggregator::{self, SendDataBuilderInfo},
+        trace_processor, INVOCATION_SPAN_RESOURCE,
     },
 };
 use datadog_trace_protobuf::pb;
-use datadog_trace_utils::trace_utils::{self, SendData};
+use datadog_trace_utils::trace_utils::{self};
 use ddcommon::hyper_migration;
 
 const TRACE_AGENT_PORT: usize = 8126;
@@ -69,7 +70,7 @@ const LAMBDA_LOAD_SPAN: &str = "aws.lambda.load";
 pub struct TraceState {
     pub config: Arc<config::Config>,
     pub trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
-    pub trace_tx: Sender<SendData>,
+    pub trace_tx: Sender<SendDataBuilderInfo>,
     pub invocation_processor: Arc<Mutex<InvocationProcessor>>,
     pub tags_provider: Arc<provider::Provider>,
 }
@@ -94,7 +95,7 @@ pub struct TraceAgent {
     pub proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
     pub tags_provider: Arc<provider::Provider>,
     invocation_processor: Arc<Mutex<InvocationProcessor>>,
-    tx: Sender<SendData>,
+    tx: Sender<SendDataBuilderInfo>,
     shutdown_token: CancellationToken,
 }
 
@@ -120,16 +121,16 @@ impl TraceAgent {
         // setup a channel to send processed traces to our flusher. tx is passed through each
         // endpoint_handler to the trace processor, which uses it to send de-serialized
         // processed trace payloads to our trace flusher.
-        let (trace_tx, mut trace_rx): (Sender<SendData>, Receiver<SendData>) =
+        let (trace_tx, mut trace_rx): (Sender<SendDataBuilderInfo>, Receiver<SendDataBuilderInfo>) =
             mpsc::channel(TRACER_PAYLOAD_CHANNEL_BUFFER_SIZE);
 
         // start our trace flusher. receives trace payloads and handles buffering + deciding when to
         // flush to backend.
 
         tokio::spawn(async move {
-            while let Some(tracer_payload) = trace_rx.recv().await {
+            while let Some(tracer_payload_info) = trace_rx.recv().await {
                 let mut aggregator = trace_aggregator.lock().await;
-                aggregator.add(tracer_payload);
+                aggregator.add(tracer_payload_info);
             }
         });
 
@@ -392,7 +393,7 @@ impl TraceAgent {
         config: Arc<config::Config>,
         request: Request,
         trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
-        trace_tx: Sender<SendData>,
+        trace_tx: Sender<SendDataBuilderInfo>,
         invocation_processor: Arc<Mutex<InvocationProcessor>>,
         tags_provider: Arc<provider::Provider>,
         version: ApiVersion,
@@ -537,7 +538,7 @@ impl TraceAgent {
     }
 
     #[must_use]
-    pub fn get_sender_copy(&self) -> Sender<SendData> {
+    pub fn get_sender_copy(&self) -> Sender<SendDataBuilderInfo> {
         self.tx.clone()
     }