refactor tracing context extraction

Author: alexgallotta

bottlecap/src/lifecycle/invocation/processor.rs

@@ -12,6 +12,7 @@ use serde_json::{json, Value};
 use tokio::sync::{mpsc::Sender, watch};
 use tracing::debug;
 
+use crate::traces::propagation::extract_composite;
 use crate::{
     config::{self, AwsConfig},
     lifecycle::invocation::{
@@ -28,16 +29,10 @@ use crate::{
     telemetry::events::{ReportMetrics, RuntimeDoneMetrics, Status},
     traces::{
         context::SpanContext,
-        propagation::{
-            text_map_propagator::{
-                DatadogHeaderPropagator, DATADOG_PARENT_ID_KEY, DATADOG_SAMPLING_PRIORITY_KEY,
-                DATADOG_SPAN_ID_KEY, DATADOG_TRACE_ID_KEY,
-            },
-            DatadogCompositePropagator, Propagator,
-        },
         trace_processor,
     },
 };
+use crate::traces::propagation::datadog_propagation::{extract_tags_datadog_context, DATADOG_PARENT_ID_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_SPAN_ID_KEY, DATADOG_TRACE_ID_KEY};
 
 pub const MS_TO_NS: f64 = 1_000_000.0;
 pub const S_TO_NS: f64 = 1_000_000_000.0;
@@ -61,7 +56,6 @@ pub struct Processor {
     // Extracted span context from inferred span, headers, or payload
     pub extracted_span_context: Option<SpanContext>,
     // Used to extract the trace context from inferred span, headers, or payload
-    propagator: DatadogCompositePropagator,
     // Helper to send enhanced metrics
     enhanced_metrics: EnhancedMetrics,
     // AWS configuration from the Lambda environment
@@ -85,15 +79,12 @@ impl Processor {
             .get_canonical_resource_name()
             .unwrap_or(String::from("aws.lambda"));
 
-        let propagator = DatadogCompositePropagator::new(Arc::clone(&config));
-
         Processor {
             context_buffer: ContextBuffer::default(),
             inferrer: SpanInferrer::new(config.service_mapping.clone()),
             span: create_empty_span(String::from("aws.lambda"), resource, service),
             cold_start_span: None,
             extracted_span_context: None,
-            propagator,
             enhanced_metrics: EnhancedMetrics::new(metrics_aggregator, Arc::clone(&config)),
             aws_config: aws_config.clone(),
             tracer_detected: false,
@@ -463,18 +454,18 @@ impl Processor {
         headers: &HashMap<String, String>,
         payload_value: &Value,
     ) -> Option<SpanContext> {
-        if let Some(sc) = self.inferrer.get_span_context(&self.propagator) {
+        if let Some(sc) = self.inferrer.get_span_context(&self.config) {
             return Some(sc);
         }
 
         if let Some(payload_headers) = payload_value.get("headers") {
-            if let Some(sc) = self.propagator.extract(payload_headers) {
+            if let Some(sc) = extract_composite(&self.config, payload_headers) {
                 debug!("Extracted trace context from event headers");
                 return Some(sc);
             }
         }
 
-        if let Some(sc) = self.propagator.extract(headers) {
+        if let Some(sc) = extract_composite(&self.config, headers) {
             debug!("Extracted trace context from headers");
             return Some(sc);
         }
@@ -564,7 +555,7 @@ impl Processor {
 
             // Extract tags from headers
             // Used for 128 bit trace ids
-            tags = DatadogHeaderPropagator::extract_tags(headers);
+            tags = extract_tags_datadog_context(headers);
         }
 
         // We should always use the generated trace id from the tracer
@@ -648,6 +639,7 @@ mod tests {
     use base64::{engine::general_purpose::STANDARD, Engine};
     use dogstatsd::aggregator::Aggregator;
     use dogstatsd::metric::EMPTY_TAGS;
+    use crate::traces::propagation::datadog_propagation::DATADOG_TRACE_ID_KEY;
 
     fn setup() -> Processor {
         let aws_config = AwsConfig {

bottlecap/src/lifecycle/invocation/span_inferrer.rs

@@ -1,10 +1,10 @@
-use std::collections::HashMap;
-
 use datadog_trace_protobuf::pb::Span;
 use serde_json::Value;
+use std::collections::HashMap;
+use std::sync::Arc;
 use tracing::debug;
 
-use crate::config::AwsConfig;
+use crate::config::{AwsConfig, Config};
 
 use crate::lifecycle::invocation::{
     generate_span_id,
@@ -22,8 +22,9 @@ use crate::lifecycle::invocation::{
         Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG,
     },
 };
+use crate::traces::context::SpanContext;
+use crate::traces::propagation::extract_composite;
 use crate::traces::span_pointers::SpanPointer;
-use crate::traces::{context::SpanContext, propagation::Propagator};
 
 #[derive(Default)]
 pub struct SpanInferrer {
@@ -305,13 +306,17 @@ impl SpanInferrer {
     /// If the carrier is set, it will try to extract the span context,
     /// otherwise it will
     ///
-    pub fn get_span_context(&self, propagator: &impl Propagator) -> Option<SpanContext> {
+    pub fn get_span_context(&self, config: &Arc<Config>) -> Option<SpanContext> {
         // Step Functions `SpanContext` is deterministically generated
         if self.generated_span_context.is_some() {
             return self.generated_span_context.clone();
         }
 
-        if let Some(sc) = self.carrier.as_ref().and_then(|c| propagator.extract(c)) {
+        if let Some(sc) = self
+            .carrier
+            .as_ref()
+            .and_then(|c| extract_composite(config, c))
+        {
             debug!("Extracted trace context from inferred span");
             return Some(sc);
         }

bottlecap/src/lifecycle/invocation/triggers/step_function_event.rs

@@ -4,19 +4,16 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use sha2::{Digest, Sha256};
 
+use super::DATADOG_CARRIER_KEY;
 use crate::{
     lifecycle::invocation::triggers::{
         ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
     },
     traces::{
         context::{Sampling, SpanContext},
-        propagation::text_map_propagator::{
-            DatadogHeaderPropagator, DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_TAGS_KEY,
-        },
     },
 };
-
-use super::DATADOG_CARRIER_KEY;
+use crate::traces::propagation::datadog_propagation::{extract_tags_datadog_context, DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_TAGS_KEY};
 
 pub const DATADOG_LEGACY_LAMBDA_PAYLOAD: &str = "Payload";
 
@@ -153,7 +150,7 @@ impl StepFunctionEvent {
                     .parse()
                     .unwrap_or(Self::generate_trace_id(self.execution.id.clone()).0);
 
-                let tags = DatadogHeaderPropagator::extract_tags(&HashMap::from([(
+                let tags = extract_tags_datadog_context(&HashMap::from([(
                     DATADOG_TAGS_KEY.to_string(),
                     trace_tags.to_string(),
                 )]));
@@ -261,7 +258,7 @@ impl ServiceNameResolver for StepFunctionEvent {
 mod tests {
     use super::*;
     use crate::lifecycle::invocation::triggers::test_utils::read_json_file;
-    use crate::traces::propagation::text_map_propagator::DATADOG_SAMPLING_DECISION_KEY;
+    use crate::traces::propagation::datadog_propagation::DATADOG_SAMPLING_DECISION_KEY;
 
     #[test]
     fn test_new_event() {

bottlecap/src/lifecycle/listener.rs

@@ -14,10 +14,7 @@ use tokio::sync::Mutex;
 use tracing::{debug, error, warn};
 
 use crate::lifecycle::invocation::processor::Processor as InvocationProcessor;
-use crate::traces::propagation::text_map_propagator::{
-    DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_TAGS_KEY,
-    DATADOG_TRACE_ID_KEY,
-};
+use crate::traces::propagation::datadog_propagation::{DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_TAGS_KEY, DATADOG_TRACE_ID_KEY};
 
 const HELLO_PATH: &str = "/lambda/hello";
 const START_INVOCATION_PATH: &str = "/lambda/start-invocation";