From deceaa72efcf86ecb32f980ade4cb68c07560385 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 2 Apr 2026 13:20:36 -0400 Subject: [PATCH 1/5] init commit --- .../src/lifecycle/invocation/span_inferrer.rs | 33 +++ bottlecap/src/traces/trace_processor.rs | 213 +++++++++++++++++- 2 files changed, 244 insertions(+), 2 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs index 10abf56dd..3ed872418 100644 --- a/bottlecap/src/lifecycle/invocation/span_inferrer.rs +++ b/bottlecap/src/lifecycle/invocation/span_inferrer.rs @@ -112,6 +112,14 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; + wrapped_inferred_span.meta.insert( + "_inferred_span.tag_source".to_string(), + "self".to_string(), + ); + wrapped_inferred_span.meta.insert( + "_inferred_span.synchronicity".to_string(), + "async".to_string(), + ); return Some(wrapped_inferred_span); } else if let Ok(event_bridge_entity) = @@ -131,6 +139,14 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; + wrapped_inferred_span.meta.insert( + "_inferred_span.tag_source".to_string(), + "self".to_string(), + ); + wrapped_inferred_span.meta.insert( + "_inferred_span.synchronicity".to_string(), + "async".to_string(), + ); return Some(wrapped_inferred_span); } @@ -158,6 +174,14 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; + wrapped_inferred_span.meta.insert( + "_inferred_span.tag_source".to_string(), + "self".to_string(), + ); + wrapped_inferred_span.meta.insert( + "_inferred_span.synchronicity".to_string(), + "async".to_string(), + ); return Some(wrapped_inferred_span); } @@ -248,6 +272,15 @@ impl SpanInferrer { if should_skip_inferred_span { self.inferred_span = None; } else { + let synchronicity = if t.is_async() { "async" } else { "sync" }; + inferred_span.meta.insert( + "_inferred_span.tag_source".to_string(), + "self".to_string(), + ); + inferred_span.meta.insert( + "_inferred_span.synchronicity".to_string(), + synchronicity.to_string(), + ); self.inferred_span = Some(inferred_span); } } diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 3864e4c60..02bd32b2d 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -5,6 +5,7 @@ use crate::appsec::processor::Processor as AppSecProcessor; use crate::appsec::processor::context::HoldArguments; use crate::config; use crate::lifecycle::invocation::processor::S_TO_MS; +use crate::lifecycle::invocation::triggers::get_default_service_name; use crate::tags::lambda::tags::COMPUTE_STATS_KEY; use crate::tags::provider; use crate::traces::span_pointers::{SpanPointer, attach_span_pointers_to_meta}; @@ -70,8 +71,25 @@ impl TraceChunkProcessor for ChunkProcessor { span.service.clone_from(service); } - // Remove the _dd.base_service tag for unintentional service name override - span.meta.remove("_dd.base_service"); + // For inferred spans, set _dd.base_service to the resolved execution + // span service (same get_default_service_name used in processor.rs). + // This covers extension-only languages (Go, .NET, Java, Ruby). + if span.meta.contains_key("_inferred_span.tag_source") + && !span.meta.contains_key("_dd.base_service") + { + let base_service = get_default_service_name( + &self + .config + .service + .clone() + .or_else(|| self.tags_provider.get_canonical_resource_name()) + .unwrap_or_else(|| "aws.lambda".to_string()), + "aws.lambda", + self.config.trace_aws_service_representation_enabled, + ); + span.meta + .insert("_dd.base_service".to_string(), base_service); + } self.tags_provider.get_tags_map().iter().for_each(|(k, v)| { span.meta.insert(k.clone(), v.clone()); @@ -1546,4 +1564,195 @@ mod tests { // Both extension and tracer compute stats → tracer takes precedence, tag "0", no double-count. check_compute_stats_behavior(true, true).await; } + + fn create_inferred_span() -> pb::Span { + let mut meta = HashMap::new(); + meta.insert( + "_inferred_span.tag_source".to_string(), + "self".to_string(), + ); + pb::Span { + name: "aws.sqs".to_string(), + service: "sqs".to_string(), + resource: "my-queue".to_string(), + trace_id: 1, + span_id: 2, + parent_id: 0, + start: 1000, + duration: 500, + error: 0, + meta, + metrics: HashMap::new(), + r#type: "web".to_string(), + span_links: vec![], + meta_struct: HashMap::new(), + span_events: vec![], + } + } + + fn create_chunk_processor(config: Arc) -> ChunkProcessor { + let tags_provider = create_tags_provider(config.clone()); + ChunkProcessor { + config, + obfuscation_config: Arc::new( + ObfuscationConfig::new().expect("Failed to create ObfuscationConfig"), + ), + tags_provider, + span_pointers: None, + } + } + + #[test] + fn test_base_service_uses_function_name_when_no_dd_service() { + let config = Arc::new(Config { + service: None, + trace_aws_service_representation_enabled: true, + ..Config::default() + }); + let mut processor = create_chunk_processor(config); + + let inferred_span = create_inferred_span(); + let mut chunk = pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: vec![inferred_span], + tags: HashMap::new(), + dropped_trace: false, + }; + + processor.process(&mut chunk, 0); + + assert_eq!( + chunk.spans[0].meta.get("_dd.base_service").unwrap(), + "my-function", + "base_service should be the function name when DD_SERVICE is not set" + ); + } + + #[test] + fn test_base_service_uses_dd_service_when_set() { + let config = Arc::new(Config { + service: Some("my-payments-api".to_string()), + trace_aws_service_representation_enabled: true, + ..Config::default() + }); + let mut processor = create_chunk_processor(config); + + let inferred_span = create_inferred_span(); + let mut chunk = pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: vec![inferred_span], + tags: HashMap::new(), + dropped_trace: false, + }; + + processor.process(&mut chunk, 0); + + assert_eq!( + chunk.spans[0].meta.get("_dd.base_service").unwrap(), + "my-payments-api", + "base_service should be DD_SERVICE when set" + ); + } + + #[test] + fn test_base_service_is_aws_lambda_when_representation_disabled() { + let config = Arc::new(Config { + service: Some("my-payments-api".to_string()), + trace_aws_service_representation_enabled: false, + ..Config::default() + }); + let mut processor = create_chunk_processor(config); + + let inferred_span = create_inferred_span(); + let mut chunk = pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: vec![inferred_span], + tags: HashMap::new(), + dropped_trace: false, + }; + + processor.process(&mut chunk, 0); + + assert_eq!( + chunk.spans[0].meta.get("_dd.base_service").unwrap(), + "aws.lambda", + "base_service should be 'aws.lambda' when representation is disabled" + ); + } + + #[test] + fn test_base_service_not_set_on_non_inferred_spans() { + let config = Arc::new(Config { + service: Some("my-service".to_string()), + trace_aws_service_representation_enabled: true, + ..Config::default() + }); + let mut processor = create_chunk_processor(config); + + let regular_span = pb::Span { + name: "http.request".to_string(), + service: "my-service".to_string(), + resource: "GET /users".to_string(), + trace_id: 1, + span_id: 3, + parent_id: 0, + start: 1000, + duration: 500, + error: 0, + meta: HashMap::new(), + metrics: HashMap::new(), + r#type: "web".to_string(), + span_links: vec![], + meta_struct: HashMap::new(), + span_events: vec![], + }; + let mut chunk = pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: vec![regular_span], + tags: HashMap::new(), + dropped_trace: false, + }; + + processor.process(&mut chunk, 0); + + assert!( + !chunk.spans[0].meta.contains_key("_dd.base_service"), + "base_service should not be set on non-inferred spans" + ); + } + + #[test] + fn test_base_service_not_overwritten_when_already_set() { + let config = Arc::new(Config { + service: Some("my-service".to_string()), + trace_aws_service_representation_enabled: true, + ..Config::default() + }); + let mut processor = create_chunk_processor(config); + + let mut inferred_span = create_inferred_span(); + inferred_span.meta.insert( + "_dd.base_service".to_string(), + "tracer-set-value".to_string(), + ); + let mut chunk = pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: vec![inferred_span], + tags: HashMap::new(), + dropped_trace: false, + }; + + processor.process(&mut chunk, 0); + + assert_eq!( + chunk.spans[0].meta.get("_dd.base_service").unwrap(), + "tracer-set-value", + "base_service should not be overwritten when already set by the tracer" + ); + } } From 40ba259393a42e7188700a9db221359ba868263a Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 2 Apr 2026 13:29:14 -0400 Subject: [PATCH 2/5] format tests --- .../src/lifecycle/invocation/span_inferrer.rs | 28 ++++++++----------- bottlecap/src/traces/trace_processor.rs | 5 +--- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs index 3ed872418..0332e3f80 100644 --- a/bottlecap/src/lifecycle/invocation/span_inferrer.rs +++ b/bottlecap/src/lifecycle/invocation/span_inferrer.rs @@ -112,10 +112,9 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; - wrapped_inferred_span.meta.insert( - "_inferred_span.tag_source".to_string(), - "self".to_string(), - ); + wrapped_inferred_span + .meta + .insert("_inferred_span.tag_source".to_string(), "self".to_string()); wrapped_inferred_span.meta.insert( "_inferred_span.synchronicity".to_string(), "async".to_string(), @@ -139,10 +138,9 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; - wrapped_inferred_span.meta.insert( - "_inferred_span.tag_source".to_string(), - "self".to_string(), - ); + wrapped_inferred_span + .meta + .insert("_inferred_span.tag_source".to_string(), "self".to_string()); wrapped_inferred_span.meta.insert( "_inferred_span.synchronicity".to_string(), "async".to_string(), @@ -174,10 +172,9 @@ impl SpanInferrer { wrapped_inferred_span.duration = inferred_span.start - wrapped_inferred_span.start; - wrapped_inferred_span.meta.insert( - "_inferred_span.tag_source".to_string(), - "self".to_string(), - ); + wrapped_inferred_span + .meta + .insert("_inferred_span.tag_source".to_string(), "self".to_string()); wrapped_inferred_span.meta.insert( "_inferred_span.synchronicity".to_string(), "async".to_string(), @@ -273,10 +270,9 @@ impl SpanInferrer { self.inferred_span = None; } else { let synchronicity = if t.is_async() { "async" } else { "sync" }; - inferred_span.meta.insert( - "_inferred_span.tag_source".to_string(), - "self".to_string(), - ); + inferred_span + .meta + .insert("_inferred_span.tag_source".to_string(), "self".to_string()); inferred_span.meta.insert( "_inferred_span.synchronicity".to_string(), synchronicity.to_string(), diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 02bd32b2d..f17fc1689 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -1567,10 +1567,7 @@ mod tests { fn create_inferred_span() -> pb::Span { let mut meta = HashMap::new(); - meta.insert( - "_inferred_span.tag_source".to_string(), - "self".to_string(), - ); + meta.insert("_inferred_span.tag_source".to_string(), "self".to_string()); pb::Span { name: "aws.sqs".to_string(), service: "sqs".to_string(), From 12d610c1e18f2ebf54ee8e4eaf97ca5f827fc915 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 8 Apr 2026 11:23:25 -0400 Subject: [PATCH 3/5] fix: replace unwrap() with expect() in tests to satisfy clippy::unwrap_used --- bottlecap/src/traces/trace_processor.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index f17fc1689..8f6812bbd 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -1620,7 +1620,7 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").unwrap(), + chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), "my-function", "base_service should be the function name when DD_SERVICE is not set" ); @@ -1647,7 +1647,7 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").unwrap(), + chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), "my-payments-api", "base_service should be DD_SERVICE when set" ); @@ -1674,7 +1674,7 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").unwrap(), + chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), "aws.lambda", "base_service should be 'aws.lambda' when representation is disabled" ); @@ -1747,7 +1747,7 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").unwrap(), + chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), "tracer-set-value", "base_service should not be overwritten when already set by the tracer" ); From 5e75b5b375f00e443ba6661acee894c92be41982 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 8 Apr 2026 11:57:46 -0400 Subject: [PATCH 4/5] fix: apply rustfmt formatting to test assertions --- bottlecap/src/traces/trace_processor.rs | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 8f6812bbd..de753c75d 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -1620,7 +1620,10 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), + chunk.spans[0] + .meta + .get("_dd.base_service") + .expect("_dd.base_service should be present"), "my-function", "base_service should be the function name when DD_SERVICE is not set" ); @@ -1647,7 +1650,10 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), + chunk.spans[0] + .meta + .get("_dd.base_service") + .expect("_dd.base_service should be present"), "my-payments-api", "base_service should be DD_SERVICE when set" ); @@ -1674,7 +1680,10 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), + chunk.spans[0] + .meta + .get("_dd.base_service") + .expect("_dd.base_service should be present"), "aws.lambda", "base_service should be 'aws.lambda' when representation is disabled" ); @@ -1747,7 +1756,10 @@ mod tests { processor.process(&mut chunk, 0); assert_eq!( - chunk.spans[0].meta.get("_dd.base_service").expect("_dd.base_service should be present"), + chunk.spans[0] + .meta + .get("_dd.base_service") + .expect("_dd.base_service should be present"), "tracer-set-value", "base_service should not be overwritten when already set by the tracer" ); From 7ab0c7b2f1be63c6b804236d7e4fbad286b22efb Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 9 Apr 2026 10:33:12 -0400 Subject: [PATCH 5/5] fix: remove reverted compute_stats test code pulled in from merge --- bottlecap/src/traces/trace_processor.rs | 157 ------------------------ 1 file changed, 157 deletions(-) diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 38532f15f..481587384 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -6,7 +6,6 @@ use crate::appsec::processor::context::HoldArguments; use crate::config; use crate::lifecycle::invocation::processor::S_TO_MS; use crate::lifecycle::invocation::triggers::get_default_service_name; -use crate::tags::lambda::tags::COMPUTE_STATS_KEY; use crate::tags::provider; use crate::traces::span_pointers::{SpanPointer, attach_span_pointers_to_meta}; use crate::traces::{ @@ -1389,162 +1388,6 @@ mod tests { ); } - /// Shared helper for the four `_dd.compute_stats` / stats-generation combination tests. - /// - /// Asserts: - /// - `_dd.compute_stats` tag value in the trace payload - /// - whether the extension generates stats via `send_processed_traces` - /// - /// | Input: `compute_trace_stats_on_extension` | Input: `client_computed_stats` | Expected: `_dd.compute_stats` | Expected: Extension generates stats? | - /// |-------------------------------------------|--------------------------------|-------------------------------|--------------------------------------| - /// | `false` | `false` | `"1"` | No | - /// | `false` | `true` | `"0"` | No | - /// | `true` | `false` | `"0"` | Yes | - /// | `true` | `true` | `"0"` | No | - #[allow(clippy::unwrap_used)] - #[allow(clippy::too_many_lines)] - async fn check_compute_stats_behavior( - compute_trace_stats_on_extension: bool, - client_computed_stats: bool, - ) { - use crate::traces::stats_concentrator_service::StatsConcentratorHandle; - use crate::traces::stats_generator::StatsGenerator; - use libdd_trace_obfuscation::obfuscation_config::ObfuscationConfig; - use tokio::sync::mpsc; - - // "_dd.compute_stats" is "1" only when neither side computes stats (backend must do it). - let expected_tag = if !compute_trace_stats_on_extension && !client_computed_stats { - "1" - } else { - "0" - }; - // The extension generates stats only when it is configured to do so and the tracer hasn't. - let expect_stats = compute_trace_stats_on_extension && !client_computed_stats; - - let config = Arc::new(Config { - apm_dd_url: "https://trace.agent.datadoghq.com".to_string(), - compute_trace_stats_on_extension, - ..Config::default() - }); - let tags_provider = Arc::new(Provider::new( - config.clone(), - "lambda".to_string(), - &std::collections::HashMap::from([( - "function_arn".to_string(), - "test-arn".to_string(), - )]), - )); - let processor = Arc::new(ServerlessTraceProcessor { - obfuscation_config: Arc::new( - ObfuscationConfig::new().expect("Failed to create ObfuscationConfig"), - ), - }); - let header_tags = tracer_header_tags::TracerHeaderTags { - lang: "rust", - lang_version: "1.0", - lang_interpreter: "", - lang_vendor: "", - tracer_version: "1.0", - container_id: "", - client_computed_top_level: false, - client_computed_stats, - dropped_p0_traces: 0, - dropped_p0_spans: 0, - }; - let span = pb::Span { - trace_id: 1, - span_id: 1, - parent_id: 0, - service: "svc".to_string(), - name: "op".to_string(), - resource: "res".to_string(), - ..Default::default() - }; - - let (stats_tx, mut stats_rx) = mpsc::unbounded_channel(); - let stats_handle = StatsConcentratorHandle::new(stats_tx); - let stats_generator = Arc::new(StatsGenerator::new(stats_handle)); - let (trace_tx, mut trace_rx) = mpsc::channel(10); - let sending_processor = SendingTraceProcessor { - appsec: None, - processor, - trace_tx, - stats_generator, - }; - - sending_processor - .send_processed_traces( - config, - tags_provider, - header_tags, - vec![vec![span]], - 0, - None, - ) - .await - .unwrap(); - - if expect_stats { - assert!( - stats_rx.try_recv().is_ok(), - "extension must generate stats when compute_trace_stats_on_extension=true \ - and client_computed_stats=false" - ); - } else { - assert!( - stats_rx.try_recv().is_err(), - "extension must not generate stats (compute_trace_stats_on_extension={compute_trace_stats_on_extension}, \ - client_computed_stats={client_computed_stats})" - ); - } - - let payload_info = trace_rx - .try_recv() - .expect("expected payload in trace_tx channel"); - let send_data = payload_info.builder.build(); - let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { - panic!("expected V07"); - }; - for payload in payloads { - assert_eq!( - payload - .tags - .get(crate::tags::lambda::tags::COMPUTE_STATS_KEY), - Some(&expected_tag.to_string()), - "_dd.compute_stats must be {expected_tag} (compute_trace_stats_on_extension={compute_trace_stats_on_extension}, \ - client_computed_stats={client_computed_stats})" - ); - } - } - - #[tokio::test] - #[allow(clippy::unwrap_used)] - async fn test_compute_stats_tag_neither_side_computes() { - // Neither extension nor tracer computes stats → backend must compute → tag "1". - check_compute_stats_behavior(false, false).await; - } - - #[tokio::test] - #[allow(clippy::unwrap_used)] - async fn test_compute_stats_tag_tracer_computes() { - // Tracer computed stats (Datadog-Client-Computed-Stats header set) → tag "0". - check_compute_stats_behavior(false, true).await; - } - - #[tokio::test] - #[allow(clippy::unwrap_used)] - async fn test_compute_stats_tag_extension_computes() { - // Extension computes stats (DD_COMPUTE_TRACE_STATS_ON_EXTENSION=true) → tag "0". - check_compute_stats_behavior(true, false).await; - } - - #[tokio::test] - #[allow(clippy::unwrap_used)] - async fn test_compute_stats_tag_both_compute() { - // Both extension and tracer compute stats → tracer takes precedence, tag "0", no double-count. - check_compute_stats_behavior(true, true).await; - } - fn create_inferred_span() -> pb::Span { let mut meta = HashMap::new(); meta.insert("_inferred_span.tag_source".to_string(), "self".to_string());