From ecced22be8a100a8967a3455d1501561ac1b916a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 11:48:33 -0400 Subject: [PATCH 01/12] headers `HeaderMap` to `HashMap` --- bottlecap/src/lifecycle/listener.rs | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/bottlecap/src/lifecycle/listener.rs b/bottlecap/src/lifecycle/listener.rs index 4b11717f1..14dfbd0cd 100644 --- a/bottlecap/src/lifecycle/listener.rs +++ b/bottlecap/src/lifecycle/listener.rs @@ -1,6 +1,7 @@ // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::convert::Infallible; use std::net::SocketAddr; use std::sync::Arc; @@ -83,13 +84,25 @@ impl Listener { invocation_processor: Arc>, ) -> http::Result> { debug!("Received start invocation request"); - let (_, body) = req.into_parts(); + let (parts, body) = req.into_parts(); match hyper::body::to_bytes(body).await { Ok(b) => { let body = b.to_vec(); let mut processor = invocation_processor.lock().await; - processor.on_invocation_start(body); + // HeaderMap to HashMap + let headers: HashMap = parts + .headers + .iter() + .map(|(k, v)| { + ( + k.as_str().to_lowercase().to_string(), + v.to_str().unwrap_or_default().to_string(), + ) + }) + .collect(); + + processor.on_invocation_start(headers, body); let mut response = Response::builder().status(200); if processor.span.trace_id != 0 { From b1894ce53d2eab5eabc03e022da516fbf8619c93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 11:51:49 -0400 Subject: [PATCH 02/12] add `Send` to propagators traits --- bottlecap/src/traces/propagation/mod.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bottlecap/src/traces/propagation/mod.rs b/bottlecap/src/traces/propagation/mod.rs index e93d81329..e25e9d35d 100644 --- a/bottlecap/src/traces/propagation/mod.rs +++ b/bottlecap/src/traces/propagation/mod.rs @@ -21,7 +21,7 @@ pub trait Propagator { } pub struct DatadogCompositePropagator { - propagators: Vec>, + propagators: Vec>, config: Arc, } @@ -64,17 +64,17 @@ impl Propagator for DatadogCompositePropagator { impl DatadogCompositePropagator { #[must_use] pub fn new(config: Arc) -> Self { - let propagators: Vec> = config + let propagators: Vec> = config .trace_propagation_style_extract .iter() .filter_map(|style| match style { TracePropagationStyle::Datadog => { Some(Box::new(text_map_propagator::DatadogHeaderPropagator) - as Box) + as Box) } TracePropagationStyle::TraceContext => { Some(Box::new(text_map_propagator::TraceContextPropagator) - as Box) + as Box) } _ => None, }) From 8d827af85c2e523c2f218ad6a0265517a6dc7f33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 12:34:53 -0400 Subject: [PATCH 03/12] add `serde_json::Value` extractor + injector --- bottlecap/src/traces/propagation/carrier.rs | 55 +++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/bottlecap/src/traces/propagation/carrier.rs b/bottlecap/src/traces/propagation/carrier.rs index fc5ef9dbb..d0f2182fa 100644 --- a/bottlecap/src/traces/propagation/carrier.rs +++ b/bottlecap/src/traces/propagation/carrier.rs @@ -3,6 +3,8 @@ /// use std::collections::HashMap; +use serde_json::Value; + /// Injector provides an interface for a carrier to be used /// with a Propagator to inject a Context into the carrier. /// @@ -38,6 +40,35 @@ impl Extractor for HashMap { } } +impl Injector for Value { + /// Set a key and value in the `Value`. + fn set(&mut self, key: &str, value: String) { + if let Value::Object(map) = self { + map.insert(key.to_lowercase(), Value::String(value)); + } + } +} + +impl Extractor for Value { + /// Get a value for a key from the `Value`. + fn get(&self, key: &str) -> Option<&str> { + if let Value::Object(map) = self { + map.get(&key.to_lowercase()).and_then(|v| v.as_str()) + } else { + None + } + } + + /// Collect all the keys from the `Value`. + fn keys(&self) -> Vec<&str> { + if let Value::Object(map) = self { + map.keys().map(String::as_str).collect::>() + } else { + Vec::new() + } + } +} + #[cfg(test)] mod test { use super::*; @@ -65,4 +96,28 @@ mod test { assert!(got.contains(&"headername1")); assert!(got.contains(&"headername2")); } + + #[test] + fn serde_value_get() { + let mut carrier = Value::Object(serde_json::Map::new()); + carrier.set("headerName", "value".to_string()); + + assert_eq!( + Extractor::get(&carrier, "HEADERNAME"), + Some("value"), + "case insensitive extraction" + ); + } + + #[test] + fn serde_value_keys() { + let mut carrier = Value::Object(serde_json::Map::new()); + carrier.set("headerName1", "value1".to_string()); + carrier.set("headerName2", "value2".to_string()); + + let got = Extractor::keys(&carrier); + assert_eq!(got.len(), 2); + assert!(got.contains(&"headername1")); + assert!(got.contains(&"headername2")); + } } From 4b6e60a4b9aa6f22efe57514eef969d97e953bf6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 12:35:15 -0400 Subject: [PATCH 04/12] add `get_carrier` to `Trigger` trait --- bottlecap/src/lifecycle/invocation/triggers/mod.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/bottlecap/src/lifecycle/invocation/triggers/mod.rs b/bottlecap/src/lifecycle/invocation/triggers/mod.rs index 5eb32ec6b..f04db8a81 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/mod.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/mod.rs @@ -13,6 +13,7 @@ pub trait Trigger: Sized { fn enrich_span(&self, span: &mut Span); fn get_tags(&self) -> HashMap; fn get_arn(&self, region: &str) -> String; + fn get_carrier(&self) -> HashMap; fn is_async(&self) -> bool; } From bd17cc80375c32763f386179bb35f74d7e102ffe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 12:35:33 -0400 Subject: [PATCH 05/12] add `get_carrier` method to current inferred spans --- .../lifecycle/invocation/triggers/api_gateway_http_event.rs | 4 ++++ .../lifecycle/invocation/triggers/api_gateway_rest_event.rs | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs index effc3e3c8..932541a00 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs @@ -182,6 +182,10 @@ impl Trigger for APIGatewayHttpEvent { .get("x-amz-invocation-type") .is_some_and(|v| v == "Event") } + + fn get_carrier(&self) -> HashMap { + self.headers.clone() + } } #[cfg(test)] diff --git a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs index 7a737d576..2ae79c40a 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs +++ b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs @@ -171,6 +171,10 @@ impl Trigger for APIGatewayRestEvent { .get("x-amz-invocation-type") .is_some_and(|v| v == "Event") } + + fn get_carrier(&self) -> HashMap { + self.headers.clone() + } } #[cfg(test)] From ece8bde90e53f25e94df1f4fe02534ab716d6de6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 17:03:43 -0400 Subject: [PATCH 06/12] update `span_inferrer.rs` to use `get_carrier` methods for distributed tracing --- .../src/lifecycle/invocation/span_inferrer.rs | 111 ++++++++++-------- 1 file changed, 62 insertions(+), 49 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs index 7b2a0eefc..6141d8dd2 100644 --- a/bottlecap/src/lifecycle/invocation/span_inferrer.rs +++ b/bottlecap/src/lifecycle/invocation/span_inferrer.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; + use datadog_trace_protobuf::pb::Span; use rand::Rng; use serde_json::Value; @@ -16,6 +18,7 @@ const FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG: &str = "function_trigger.event_sour pub struct SpanInferrer { inferred_span: Option, is_async_span: bool, + carrier: Option>, } impl Default for SpanInferrer { @@ -30,6 +33,7 @@ impl SpanInferrer { Self { inferred_span: None, is_async_span: false, + carrier: None, } } @@ -37,58 +41,56 @@ impl SpanInferrer { /// and try matching it to a `Trigger` implementation, which will create /// an inferred span and set it to `self.inferred_span` /// - pub fn infer_span(&mut self, payload: &[u8], aws_config: &AwsConfig) { + pub fn infer_span(&mut self, payload_value: &Value, aws_config: &AwsConfig) { self.inferred_span = None; - if let Ok(payload_value) = serde_json::from_slice::(payload) { - if APIGatewayHttpEvent::is_match(&payload_value) { - if let Some(t) = APIGatewayHttpEvent::new(payload_value) { - let mut span = Span { - span_id: Self::generate_span_id(), - ..Default::default() - }; - - t.enrich_span(&mut span); - span.meta.extend([ - ( - FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(), - "api_gateway".to_string(), - ), - ( - FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(), - t.get_arn(&aws_config.region), - ), - ]); - - self.is_async_span = t.is_async(); - self.inferred_span = Some(span); - } - } else if APIGatewayRestEvent::is_match(&payload_value) { - if let Some(t) = APIGatewayRestEvent::new(payload_value) { - let mut span = Span { - span_id: Self::generate_span_id(), - ..Default::default() - }; - - t.enrich_span(&mut span); - span.meta.extend([ - ( - FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(), - "api_gateway".to_string(), - ), - ( - FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(), - t.get_arn(&aws_config.region), - ), - ]); - - self.is_async_span = t.is_async(); - self.inferred_span = Some(span); - } - } else { - debug!("Unable to infer span from payload"); + if APIGatewayHttpEvent::is_match(payload_value) { + if let Some(t) = APIGatewayHttpEvent::new(payload_value.clone()) { + let mut span = Span { + span_id: Self::generate_span_id(), + ..Default::default() + }; + + t.enrich_span(&mut span); + span.meta.extend([ + ( + FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(), + "api_gateway".to_string(), + ), + ( + FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(), + t.get_arn(&aws_config.region), + ), + ]); + + self.carrier = Some(t.get_carrier()); + self.is_async_span = t.is_async(); + self.inferred_span = Some(span); + } + } else if APIGatewayRestEvent::is_match(payload_value) { + if let Some(t) = APIGatewayRestEvent::new(payload_value.clone()) { + let mut span = Span { + span_id: Self::generate_span_id(), + ..Default::default() + }; + + t.enrich_span(&mut span); + span.meta.extend([ + ( + FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(), + "api_gateway".to_string(), + ), + ( + FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(), + t.get_arn(&aws_config.region), + ), + ]); + + self.carrier = Some(t.get_carrier()); + self.is_async_span = t.is_async(); + self.inferred_span = Some(span); } } else { - debug!("Unable to serialize payload"); + debug!("Unable to infer span from payload"); } } @@ -101,6 +103,12 @@ impl SpanInferrer { } } + pub fn extend_meta(&mut self, iter: HashMap) { + if let Some(s) = &mut self.inferred_span { + s.meta.extend(iter); + } + } + pub fn set_status_code(&mut self, status_code: String) { if let Some(s) = &mut self.inferred_span { s.meta.insert("http.status_code".to_string(), status_code); @@ -136,4 +144,9 @@ impl SpanInferrer { pub fn get_inferred_span(&self) -> &Option { &self.inferred_span } + + #[must_use] + pub fn get_carrier(&self) -> Option> { + self.carrier.clone() + } } From 7f155728cd2ae9d1ec65a86d692ed7d5b6e055b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 17:07:06 -0400 Subject: [PATCH 07/12] add `headers_to_map` function --- bottlecap/src/lifecycle/listener.rs | 52 +++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/bottlecap/src/lifecycle/listener.rs b/bottlecap/src/lifecycle/listener.rs index 14dfbd0cd..5f8a6b42b 100644 --- a/bottlecap/src/lifecycle/listener.rs +++ b/bottlecap/src/lifecycle/listener.rs @@ -13,6 +13,9 @@ 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_TRACE_ID_KEY, +}; const HELLO_PATH: &str = "/lambda/hello"; const START_INVOCATION_PATH: &str = "/lambda/start-invocation"; @@ -90,24 +93,31 @@ impl Listener { let body = b.to_vec(); let mut processor = invocation_processor.lock().await; - // HeaderMap to HashMap - let headers: HashMap = parts - .headers - .iter() - .map(|(k, v)| { - ( - k.as_str().to_lowercase().to_string(), - v.to_str().unwrap_or_default().to_string(), - ) - }) - .collect(); + let headers = Self::headers_to_map(parts.headers); processor.on_invocation_start(headers, body); let mut response = Response::builder().status(200); - if processor.span.trace_id != 0 { - response = - response.header("x-datadog-trace-id", processor.span.trace_id.to_string()); + + // If a `SpanContext` exists, then tell the tracer to use it. + // todo: update this whole code with DatadogHeaderPropagator::inject + // since this logic looks messy + if let Some(sp) = &processor.extracted_span_context { + response = response.header(DATADOG_TRACE_ID_KEY, sp.trace_id.to_string()); + if let Some(priority) = sp.sampling.and_then(|s| s.priority) { + response = + response.header(DATADOG_SAMPLING_PRIORITY_KEY, priority.to_string()); + } + + // Handle 128 bit trace ids + if let Some(trace_id_higher_order_bits) = + sp.tags.get(DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY) + { + response = response.header( + DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, + trace_id_higher_order_bits, + ); + } } drop(processor); @@ -141,6 +151,8 @@ impl Listener { let mut processor = invocation_processor.lock().await; + // todo: fix this, code is a copy of the existing logic in Go, not accounting + // when a 128 bit trace id exist let mut trace_id = 0; if let Some(header) = headers.get("x-datadog-trace-id") { if let Ok(header_value) = header.to_str() { @@ -176,4 +188,16 @@ impl Listener { .status(200) .body(Body::from(json!({}).to_string())) } + + fn headers_to_map(headers: http::HeaderMap) -> HashMap { + headers + .iter() + .map(|(k, v)| { + ( + k.as_str().to_string(), + v.to_str().unwrap_or_default().to_string(), + ) + }) + .collect() + } } From 2362a539c22e587de4667818e97a49071a991ab6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 17:12:10 -0400 Subject: [PATCH 08/12] reparent spans I suspect there might be something wrong here, the code in Go is quite convoluted --- .../src/lifecycle/invocation/processor.rs | 58 +++++++++++++++++-- 1 file changed, 54 insertions(+), 4 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index 9b9cc98ef..6bd58b854 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -7,6 +7,7 @@ use std::{ use chrono::{DateTime, Utc}; 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 tracing::debug; @@ -14,7 +15,11 @@ use crate::{ config::{self, AwsConfig}, lifecycle::invocation::{context::ContextBuffer, span_inferrer::SpanInferrer}, tags::provider, - traces::trace_processor, + traces::{ + context::SpanContext, + propagation::{DatadogCompositePropagator, Propagator}, + trace_processor, + }, }; pub const MS_TO_NS: f64 = 1_000_000.0; @@ -23,6 +28,9 @@ pub struct Processor { pub context_buffer: ContextBuffer, inferrer: SpanInferrer, pub span: Span, + pub extracted_span_context: Option, + // Used to extract the trace context from an inferred span + propagator: DatadogCompositePropagator, aws_config: AwsConfig, tracer_detected: bool, } @@ -39,6 +47,8 @@ impl Processor { .get_canonical_resource_name() .unwrap_or("aws_lambda".to_string()); + let propagator = DatadogCompositePropagator::new(Arc::clone(&config)); + Processor { context_buffer: ContextBuffer::default(), inferrer: SpanInferrer::default(), @@ -58,6 +68,8 @@ impl Processor { meta_struct: HashMap::new(), span_links: Vec::new(), }, + extracted_span_context: None, + propagator, aws_config: aws_config.clone(), tracer_detected: false, } @@ -166,7 +178,7 @@ impl Processor { /// If this method is called, it means that we are operating in a Universally Instrumented /// runtime. Therefore, we need to set the `tracer_detected` flag to `true`. /// - pub fn on_invocation_start(&mut self, payload: Vec) { + pub fn on_invocation_start(&mut self, headers: HashMap, payload: Vec) { self.tracer_detected = true; // Reset trace context @@ -174,13 +186,52 @@ impl Processor { self.span.parent_id = 0; self.span.span_id = 0; - self.inferrer.infer_span(&payload, &self.aws_config); + let payload_value = match serde_json::from_slice::(&payload) { + Ok(value) => value, + Err(_) => json!({}), + }; + + self.extract_span_context(&headers, &payload_value); + self.inferrer.infer_span(&payload_value, &self.aws_config); + + if let Some(sp) = &self.extracted_span_context { + self.span.trace_id = sp.trace_id; + self.span.parent_id = sp.span_id; + + if self.inferrer.get_inferred_span().is_some() { + self.inferrer.set_parent_id(sp.span_id); + self.inferrer.extend_meta(sp.tags.clone()); + } else { + self.span.meta.extend(sp.tags.clone()); + } + } if let Some(inferred_span) = self.inferrer.get_inferred_span() { self.span.parent_id = inferred_span.span_id; } } + fn extract_span_context(&mut self, headers: &HashMap, payload_value: &Value) { + if let Some(carrier) = self.inferrer.get_carrier() { + if let Some(sc) = self.propagator.extract(&carrier) { + debug!("Extracted trace context from inferred span"); + self.extracted_span_context = Some(sc); + } + } + + if let Some(payload_headers) = payload_value.get("headers") { + if let Some(sc) = self.propagator.extract(payload_headers) { + debug!("Extracted trace context from event headers"); + self.extracted_span_context = Some(sc); + } + } + + if let Some(sc) = self.propagator.extract(headers) { + debug!("Extracted trace context from headers"); + self.extracted_span_context = Some(sc); + } + } + /// Given trace context information, set it to the current span. /// pub fn on_invocation_end( @@ -194,7 +245,6 @@ impl Processor { self.span.span_id = span_id; if self.inferrer.get_inferred_span().is_some() { - self.inferrer.set_parent_id(parent_id); if let Some(status_code) = status_code { self.inferrer.set_status_code(status_code); } From 63a66ac4c7f423aadec0eac01406767afb0e45b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Tue, 22 Oct 2024 17:15:10 -0400 Subject: [PATCH 09/12] make some variables public --- bottlecap/src/traces/propagation/text_map_propagator.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bottlecap/src/traces/propagation/text_map_propagator.rs b/bottlecap/src/traces/propagation/text_map_propagator.rs index 42b4a17fe..947893cff 100644 --- a/bottlecap/src/traces/propagation/text_map_propagator.rs +++ b/bottlecap/src/traces/propagation/text_map_propagator.rs @@ -12,9 +12,9 @@ use crate::traces::propagation::{ }; // Datadog Keys -const DATADOG_TRACE_ID_KEY: &str = "x-datadog-trace-id"; +pub const DATADOG_TRACE_ID_KEY: &str = "x-datadog-trace-id"; const DATADOG_PARENT_ID_KEY: &str = "x-datadog-parent-id"; -const DATADOG_SAMPLING_PRIORITY_KEY: &str = "x-datadog-sampling-priority"; +pub const DATADOG_SAMPLING_PRIORITY_KEY: &str = "x-datadog-sampling-priority"; const DATADOG_ORIGIN_KEY: &str = "x-datadog-origin"; const DATADOG_TAGS_KEY: &str = "x-datadog-tags"; From 8a411492d8dd84a940bdab502674b8bd75351bfc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Wed, 23 Oct 2024 10:59:32 -0400 Subject: [PATCH 10/12] fix to return early on `extract_span_context` --- bottlecap/src/lifecycle/invocation/processor.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index 6bd58b854..254042d75 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -191,7 +191,7 @@ impl Processor { Err(_) => json!({}), }; - self.extract_span_context(&headers, &payload_value); + self.extracted_span_context = self.extract_span_context(&headers, &payload_value); self.inferrer.infer_span(&payload_value, &self.aws_config); if let Some(sp) = &self.extracted_span_context { @@ -211,25 +211,31 @@ impl Processor { } } - fn extract_span_context(&mut self, headers: &HashMap, payload_value: &Value) { + fn extract_span_context( + &mut self, + headers: &HashMap, + payload_value: &Value, + ) -> Option { if let Some(carrier) = self.inferrer.get_carrier() { if let Some(sc) = self.propagator.extract(&carrier) { debug!("Extracted trace context from inferred span"); - self.extracted_span_context = Some(sc); + return Some(sc); } } if let Some(payload_headers) = payload_value.get("headers") { if let Some(sc) = self.propagator.extract(payload_headers) { debug!("Extracted trace context from event headers"); - self.extracted_span_context = Some(sc); + return Some(sc); } } if let Some(sc) = self.propagator.extract(headers) { debug!("Extracted trace context from headers"); - self.extracted_span_context = Some(sc); + return Some(sc); } + + None } /// Given trace context information, set it to the current span. From cd486692615a68d27644f1a1cd091536274202dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Wed, 23 Oct 2024 13:21:45 -0400 Subject: [PATCH 11/12] fix how 128 bit is handled also updated some variable names --- bottlecap/src/lifecycle/invocation/processor.rs | 14 ++++++++------ bottlecap/src/lifecycle/listener.rs | 7 ++++--- .../src/traces/propagation/text_map_propagator.rs | 2 +- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index 254042d75..6384ea3b5 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -194,15 +194,17 @@ impl Processor { self.extracted_span_context = self.extract_span_context(&headers, &payload_value); self.inferrer.infer_span(&payload_value, &self.aws_config); - if let Some(sp) = &self.extracted_span_context { - self.span.trace_id = sp.trace_id; - self.span.parent_id = sp.span_id; + if let Some(sc) = &self.extracted_span_context { + self.span.trace_id = sc.trace_id; + self.span.parent_id = sc.span_id; + // Set the right data to the correct root level span, + // If there's an inferred span, then that should be the root. if self.inferrer.get_inferred_span().is_some() { - self.inferrer.set_parent_id(sp.span_id); - self.inferrer.extend_meta(sp.tags.clone()); + self.inferrer.set_parent_id(sc.span_id); + self.inferrer.extend_meta(sc.tags.clone()); } else { - self.span.meta.extend(sp.tags.clone()); + self.span.meta.extend(sc.tags.clone()); } } diff --git a/bottlecap/src/lifecycle/listener.rs b/bottlecap/src/lifecycle/listener.rs index 5f8a6b42b..a4d39310b 100644 --- a/bottlecap/src/lifecycle/listener.rs +++ b/bottlecap/src/lifecycle/listener.rs @@ -14,7 +14,8 @@ 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_TRACE_ID_KEY, + DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, DATADOG_SAMPLING_PRIORITY_KEY, DATADOG_TAGS_KEY, + DATADOG_TRACE_ID_KEY, }; const HELLO_PATH: &str = "/lambda/hello"; @@ -114,8 +115,8 @@ impl Listener { sp.tags.get(DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY) { response = response.header( - DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY, - trace_id_higher_order_bits, + DATADOG_TAGS_KEY, + format!("{DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY}={trace_id_higher_order_bits}"), ); } } diff --git a/bottlecap/src/traces/propagation/text_map_propagator.rs b/bottlecap/src/traces/propagation/text_map_propagator.rs index 947893cff..1a0803aac 100644 --- a/bottlecap/src/traces/propagation/text_map_propagator.rs +++ b/bottlecap/src/traces/propagation/text_map_propagator.rs @@ -16,7 +16,7 @@ pub const DATADOG_TRACE_ID_KEY: &str = "x-datadog-trace-id"; const DATADOG_PARENT_ID_KEY: &str = "x-datadog-parent-id"; pub const DATADOG_SAMPLING_PRIORITY_KEY: &str = "x-datadog-sampling-priority"; const DATADOG_ORIGIN_KEY: &str = "x-datadog-origin"; -const DATADOG_TAGS_KEY: &str = "x-datadog-tags"; +pub const DATADOG_TAGS_KEY: &str = "x-datadog-tags"; pub const DATADOG_HIGHER_ORDER_TRACE_ID_BITS_KEY: &str = "_dd.p.tid"; const DATADOG_PROPAGATION_ERROR_KEY: &str = "_dd.propagation_error"; From 40e84a905f92fbc43f39f15717090e1425167b17 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Wed, 23 Oct 2024 13:29:28 -0400 Subject: [PATCH 12/12] update comment --- bottlecap/src/lifecycle/invocation/processor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bottlecap/src/lifecycle/invocation/processor.rs b/bottlecap/src/lifecycle/invocation/processor.rs index 6384ea3b5..d7a7cad66 100644 --- a/bottlecap/src/lifecycle/invocation/processor.rs +++ b/bottlecap/src/lifecycle/invocation/processor.rs @@ -29,7 +29,7 @@ pub struct Processor { inferrer: SpanInferrer, pub span: Span, pub extracted_span_context: Option, - // Used to extract the trace context from an inferred span + // Used to extract the trace context from inferred span, headers, or payload propagator: DatadogCompositePropagator, aws_config: AwsConfig, tracer_detected: bool,