feat: add parent and trace id from proxy to all spans

Author: alexgallotta

bottlecap/src/bin/bottlecap/main.rs

@@ -23,6 +23,7 @@ use bottlecap::{
         agent::LogsAgent,
         flusher::{build_fqdn_logs, Flusher as LogsFlusher},
     },
+    lwa::proxy::start_lwa_proxy,
     metrics::enhanced::lambda::Lambda as enhanced_metrics,
     secrets::decrypt,
     tags::{lambda, provider::Provider as TagProvider},
@@ -70,7 +71,6 @@ use tokio::sync::Mutex as TokioMutex;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error};
 use tracing_subscriber::EnvFilter;
-use bottlecap::lwa::proxy::start_lwa_proxy;
 
 #[derive(Clone, Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -301,16 +301,16 @@ async fn extension_loop_active(
         Arc::clone(config),
         aws_config,
     )));
-    let _ = start_lwa_proxy(Arc::clone(&invocation_processor));
 
-
-    let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor {
+    let trace_processor = Arc::new(TokioMutex::new(trace_processor::ServerlessTraceProcessor {
         obfuscation_config: Arc::new(
             obfuscation_config::ObfuscationConfig::new()
                 .map_err(|e| Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?,
         ),
         resolved_api_key: resolved_api_key.clone(),
-    });
+        override_trace_id: None,
+        substitute_parent_span_id: None,
+    }));
 
     let stats_flusher = Arc::new(stats_flusher::ServerlessStatsFlusher {
         buffer: Arc::new(TokioMutex::new(Vec::new())),
@@ -333,6 +333,12 @@ async fn extension_loop_active(
         )
         .await,
     );
+
+    let _ = start_lwa_proxy(
+        Arc::clone(&invocation_processor),
+        Arc::clone(&trace_processor),
+    );
+
     let trace_agent_tx = trace_agent.get_sender_copy();
 
     tokio::spawn(async move {

bottlecap/src/lifecycle/invocation/processor.rs

@@ -9,16 +9,18 @@ use datadog_trace_protobuf::pb::Span;
 use datadog_trace_utils::{send_data::SendData, tracer_header_tags};
 use serde_json::{json, Value};
 use tokio::sync::mpsc::Sender;
+use tokio::sync::Mutex;
 use tracing::debug;
 
+use crate::traces::trace_processor::TraceProcessor;
 use crate::{
-    config::{self, AwsConfig},
+    config,
+    config::AwsConfig,
     lifecycle::invocation::{context::ContextBuffer, span_inferrer::SpanInferrer},
     tags::provider,
     traces::{
         context::SpanContext,
         propagation::{DatadogCompositePropagator, Propagator},
-        trace_processor,
     },
 };
 
@@ -97,7 +99,7 @@ impl Processor {
         duration_ms: f64,
         config: Arc<config::Config>,
         tags_provider: Arc<provider::Provider>,
-        trace_processor: Arc<dyn trace_processor::TraceProcessor + Send + Sync>,
+        trace_processor: Arc<Mutex<impl TraceProcessor>>,
         trace_agent_tx: Sender<SendData>,
     ) {
         self.context_buffer
@@ -143,12 +145,13 @@ impl Processor {
                 client_computed_stats: false,
             };
 
-            let send_data = trace_processor.process_traces(
+            let send_data = trace_processor.lock().await.process_traces(
                 config.clone(),
                 tags_provider.clone(),
                 header_tags,
                 vec![traces],
                 body_size,
+                true,
             );
 
             if let Err(e) = trace_agent_tx.send(send_data).await {