format and style

Author: alexgallotta

bottlecap/src/lifecycle/invocation/processor.rs

@@ -12,6 +12,10 @@ use serde_json::{json, Value};
 use tokio::sync::{mpsc::Sender, watch};
 use tracing::debug;
 
+use crate::traces::propagation::datadog_extraction::{
+    extract_tags_datadog_context, DATADOG_PARENT_ID_KEY, DATADOG_SAMPLING_PRIORITY_KEY,
+    DATADOG_SPAN_ID_KEY, DATADOG_TRACE_ID_KEY,
+};
 use crate::traces::propagation::extract_composite;
 use crate::{
     config::{self, AwsConfig},
@@ -27,12 +31,8 @@ use crate::{
     },
     tags::{lambda::tags::resolve_runtime_from_proc, provider},
     telemetry::events::{ReportMetrics, RuntimeDoneMetrics, Status},
-    traces::{
-        context::SpanContext,
-        trace_processor,
-    },
+    traces::{context::SpanContext, 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;
@@ -635,11 +635,11 @@ impl Processor {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::traces::propagation::datadog_extraction::DATADOG_TRACE_ID_KEY;
     use crate::LAMBDA_RUNTIME_SLUG;
     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/triggers/step_function_event.rs

@@ -5,15 +5,15 @@ use serde_json::Value;
 use sha2::{Digest, Sha256};
 
 use super::DATADOG_CARRIER_KEY;
+use crate::traces::propagation::datadog_extraction::{
+    extract_tags_datadog_context, DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_TAGS_KEY,
+};
 use crate::{
     lifecycle::invocation::triggers::{
         ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
     },
-    traces::{
-        context::{Sampling, SpanContext},
-    },
+    traces::context::{Sampling, SpanContext},
 };
-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";
 
@@ -258,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::datadog_propagation::DATADOG_SAMPLING_DECISION_KEY;
+    use crate::traces::propagation::datadog_extraction::DATADOG_SAMPLING_DECISION_KEY;
 
     #[test]
     fn test_new_event() {

bottlecap/src/lifecycle/listener.rs

@@ -14,7 +14,10 @@ use tokio::sync::Mutex;
 use tracing::{debug, error, warn};
 
 use crate::lifecycle::invocation::processor::Processor as InvocationProcessor;
-use crate::traces::propagation::datadog_propagation::{DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_TAGS_KEY, DATADOG_TRACE_ID_KEY};
+use crate::traces::propagation::datadog_extraction::{
+    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";

…races/propagation/datadog_propagation.rs …traces/propagation/datadog_extraction.rsbottlecap/src/traces/propagation/datadog_propagation.rs renamed to bottlecap/src/traces/propagation/datadog_extraction.rs bottlecap/src/traces/propagation/datadog_propagation.rs renamed to bottlecap/src/traces/propagation/datadog_extraction.rs

@@ -183,8 +183,8 @@ fn higher_order_bits_valid(trace_id_higher_order_bits: &str) -> bool {
 #[cfg(test)]
 #[allow(clippy::unwrap_used)]
 mod test {
-    use crate::traces::propagation::text_map_propagation::extract_context_standard_header;
     use super::*;
+    use crate::traces::propagation::text_map_extraction::extract_context_standard_header;
 
     #[test]
     fn test_extract_traceparent_propagator() {

bottlecap/src/traces/propagation/mod.rs

@@ -1,3 +1,6 @@
+use crate::traces::propagation::datadog_extraction::{
+    DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_LAST_PARENT_ID_KEY,
+};
 use crate::{
     config::{self, trace_propagation_style::TracePropagationStyle},
     traces::context::SpanContext,
@@ -6,15 +9,12 @@ use carrier::Extractor;
 use datadog_trace_protobuf::pb::SpanLink;
 use std::collections::VecDeque;
 use std::{collections::HashMap, sync::Arc};
-use text_map_propagation::{
-    BAGGAGE_PREFIX, TRACESTATE_KEY,
-};
-use crate::traces::propagation::datadog_propagation::{DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_LAST_PARENT_ID_KEY};
+use text_map_extraction::{BAGGAGE_PREFIX, TRACESTATE_KEY};
 
 pub mod carrier;
+pub mod datadog_extraction;
 pub mod error;
-pub mod text_map_propagation;
-pub mod datadog_propagation;
+pub mod text_map_extraction;
 
 pub fn extract_composite(
     config: &Arc<config::Config>,
@@ -33,15 +33,13 @@ pub fn extract_composite(
         }
 
     let contexts_found = extract_available_contexts(config, carrier);
-    if contexts_found.is_empty() {
-        return None;
-    }
-    let mut primary_context = resolve_primary_context_with_links(contexts_found);
-    if config.trace_propagation_http_baggage_enabled {
-        attach_baggage(&mut primary_context, carrier);
+    if let Some(mut primary_context) = resolve_primary_context_with_links(contexts_found) {
+        if config.trace_propagation_http_baggage_enabled {
+            attach_baggage(&mut primary_context, carrier);
+        }
+        return Some(primary_context);
     }
-
-    Some(primary_context)
+    None
 }
 
 fn extract_available_contexts(
@@ -50,26 +48,26 @@ fn extract_available_contexts(
 ) -> VecDeque<(TracePropagationStyle, SpanContext)> {
     let mut contexts_found: VecDeque<(TracePropagationStyle, SpanContext)> = VecDeque::new();
     for propagation_style in &config.trace_propagation_style_extract {
-            if let Some(context) = propagation_style.extract(carrier) {
-                contexts_found.push_back((*propagation_style, context));
-            }
+        if let Some(context) = propagation_style.extract(carrier) {
+            contexts_found.push_back((*propagation_style, context));
+        }
     }
     contexts_found
 }
 
 fn resolve_primary_context_with_links(
     mut all_contexts: VecDeque<(TracePropagationStyle, SpanContext)>,
-) -> SpanContext {
+) -> Option<SpanContext> {
     let dd_context: Option<SpanContext> = all_contexts
         .iter()
         .find(|(style, _)| *style == TracePropagationStyle::Datadog)
-        .and_then(|style_and_context| Some(style_and_context.1.clone()))
+        .map(|style_and_context| style_and_context.1.clone())
         .or(None);
 
-    let (_primary_style, mut primary_context) = all_contexts.pop_front().unwrap();
+    let (_primary_style, mut primary_context) = all_contexts.pop_front()?;
     let mut links = Vec::<SpanLink>::new();
 
-    for style_and_context in all_contexts.iter() {
+    for style_and_context in &all_contexts {
         let style = style_and_context.0;
         let context = style_and_context.1.clone();
 
@@ -129,7 +127,7 @@ fn resolve_primary_context_with_links(
 
     primary_context.links = links;
 
-    primary_context
+    Some(primary_context)
 }
 
 fn attach_baggage(context: &mut SpanContext, carrier: &dyn Extractor) {

…aces/propagation/text_map_propagation.rs …races/propagation/text_map_extraction.rsbottlecap/src/traces/propagation/text_map_propagation.rs renamed to bottlecap/src/traces/propagation/text_map_extraction.rs bottlecap/src/traces/propagation/text_map_propagation.rs renamed to bottlecap/src/traces/propagation/text_map_extraction.rs

@@ -2,11 +2,14 @@ use std::collections::HashMap;
 
 use crate::config::trace_propagation_style::TracePropagationStyle;
 use crate::traces::context::{Sampling, SpanContext};
+use crate::traces::propagation::datadog_extraction::{
+    extract_context_datadog_header, DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY,
+    DATADOG_LAST_PARENT_ID_KEY,
+};
 use crate::traces::propagation::{carrier::Extractor, error::Error};
 use lazy_static::lazy_static;
 use regex::Regex;
 use tracing::{debug, error, warn};
-use crate::traces::propagation::datadog_propagation::{extract_context_datadog_header, DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_LAST_PARENT_ID_KEY};
 
 // Traceparent Keys
 const TRACEPARENT_KEY: &str = "traceparent";
@@ -33,10 +36,8 @@ impl TracePropagationStyle {
     pub fn extract(&self, carrier: &dyn Extractor) -> Option<SpanContext> {
         match self {
             TracePropagationStyle::Datadog => extract_context_datadog_header(carrier),
-            TracePropagationStyle::B3Multi => None,
-            TracePropagationStyle::B3 => None,
             TracePropagationStyle::TraceContext => extract_context_standard_header(carrier),
-            TracePropagationStyle::None => None,
+            _ => None,
         }
     }
 }