diff --git a/crates/datadog-trace-agent/src/config.rs b/crates/datadog-trace-agent/src/config.rs index 0ea40508..31459257 100644 --- a/crates/datadog-trace-agent/src/config.rs +++ b/crates/datadog-trace-agent/src/config.rs @@ -3,8 +3,10 @@ use ddcommon::Endpoint; use std::borrow::Cow; +use std::collections::HashMap; use std::env; use std::str::FromStr; +use std::sync::OnceLock; use datadog_trace_obfuscation::obfuscation_config; use datadog_trace_utils::config_utils::{ @@ -15,6 +17,59 @@ use datadog_trace_utils::trace_utils; const DEFAULT_DOGSTATSD_PORT: u16 = 8125; +#[derive(Debug)] +pub struct Tags { + tags: HashMap, + function_tags_string: OnceLock, +} + +impl Tags { + pub fn from_env_string(env_tags: &str) -> Self { + let mut tags = HashMap::new(); + + // Space-separated key:value tags are the standard for tagging. For compatibility reasons + // we also support comma-separated key:value tags as well. + let normalized = env_tags.replace(',', " "); + + for kv in normalized.split_whitespace() { + let parts = kv.split(':').collect::>(); + if parts.len() == 2 { + tags.insert(parts[0].to_string(), parts[1].to_string()); + } + } + Self { + tags, + function_tags_string: OnceLock::new(), + } + } + + pub fn new() -> Self { + Self { + tags: HashMap::new(), + function_tags_string: OnceLock::new(), + } + } + + pub fn tags(&self) -> &HashMap { + &self.tags + } + + pub fn function_tags(&self) -> Option<&str> { + if self.tags.is_empty() { + return None; + } + Some(self.function_tags_string.get_or_init(|| { + let mut kvs = self + .tags + .iter() + .map(|(k, v)| format!("{k}:{v}")) + .collect::>(); + kvs.sort(); + kvs.join(",") + })) + } +} + #[derive(Debug)] pub struct Config { pub dd_site: String, @@ -24,6 +79,7 @@ pub struct Config { pub max_request_content_length: usize, pub obfuscation_config: obfuscation_config::ObfuscationConfig, pub os: String, + pub tags: Tags, /// how often to flush stats, in seconds pub stats_flush_interval: u64, /// how often to flush traces, in seconds @@ -69,6 +125,12 @@ impl Config { ) })?; + let tags = if let Ok(env_tags) = env::var("DD_TAGS") { + Tags::from_env_string(&env_tags) + } else { + Tags::new() + }; + #[allow(clippy::unwrap_used)] Ok(Config { app_name: Some(app_name), @@ -94,6 +156,7 @@ impl Config { proxy_url: env::var("DD_PROXY_HTTPS") .or_else(|_| env::var("HTTPS_PROXY")) .ok(), + tags, }) } } @@ -102,6 +165,7 @@ impl Config { mod tests { use duplicate::duplicate_item; use serial_test::serial; + use std::collections::HashMap; use std::env; use crate::config; @@ -250,4 +314,81 @@ mod tests { env::remove_var("ASCSVCRT_SPRING__APPLICATION__NAME"); env::remove_var("DD_DOGSTATSD_PORT"); } + + fn test_config_with_dd_tags(dd_tags: &str) -> config::Config { + env::set_var("DD_API_KEY", "_not_a_real_key_"); + env::set_var("ASCSVCRT_SPRING__APPLICATION__NAME", "test-spring-app"); + env::set_var("DD_TAGS", dd_tags); + let config_res = config::Config::new(); + assert!(config_res.is_ok()); + let config = config_res.unwrap(); + env::remove_var("DD_API_KEY"); + env::remove_var("ASCSVCRT_SPRING__APPLICATION__NAME"); + env::remove_var("DD_TAGS"); + config + } + + #[test] + #[serial] + fn test_dd_tags_comma_separated() { + let config = test_config_with_dd_tags("some:tag,another:thing,invalid:thing:here"); + let expected_tags = HashMap::from([ + ("some".to_string(), "tag".to_string()), + ("another".to_string(), "thing".to_string()), + ]); + assert_eq!(config.tags.tags(), &expected_tags); + assert_eq!(config.tags.function_tags(), Some("another:thing,some:tag")); + } + + #[test] + #[serial] + fn test_dd_tags_space_separated() { + let config = test_config_with_dd_tags("some:tag another:thing invalid:thing:here"); + let expected_tags = HashMap::from([ + ("some".to_string(), "tag".to_string()), + ("another".to_string(), "thing".to_string()), + ]); + assert_eq!(config.tags.tags(), &expected_tags); + assert_eq!(config.tags.function_tags(), Some("another:thing,some:tag")); + } + + #[test] + #[serial] + fn test_dd_tags_mixed_separators() { + let config = test_config_with_dd_tags("some:tag,another:thing extra:value"); + let expected_tags = HashMap::from([ + ("some".to_string(), "tag".to_string()), + ("another".to_string(), "thing".to_string()), + ("extra".to_string(), "value".to_string()), + ]); + assert_eq!(config.tags.tags(), &expected_tags); + assert_eq!( + config.tags.function_tags(), + Some("another:thing,extra:value,some:tag") + ); + } + + #[test] + #[serial] + fn test_dd_tags_no_valid_tags() { + // Test with only invalid tags + let config = test_config_with_dd_tags("invalid:thing:here,also-bad"); + assert_eq!(config.tags.tags(), &HashMap::new()); + assert_eq!(config.tags.function_tags(), None); + + // Test with empty string + let config = test_config_with_dd_tags(""); + assert_eq!(config.tags.tags(), &HashMap::new()); + assert_eq!(config.tags.function_tags(), None); + + // Test with just whitespace + let config = test_config_with_dd_tags(" "); + assert_eq!(config.tags.tags(), &HashMap::new()); + assert_eq!(config.tags.function_tags(), None); + + // Test with just commas and spaces + let config = test_config_with_dd_tags(" , , "); + assert_eq!(config.tags.tags(), &HashMap::new()); + assert_eq!(config.tags.function_tags(), None); + } } diff --git a/crates/datadog-trace-agent/src/trace_processor.rs b/crates/datadog-trace-agent/src/trace_processor.rs index afe3f4c9..dbdf65d1 100644 --- a/crates/datadog-trace-agent/src/trace_processor.rs +++ b/crates/datadog-trace-agent/src/trace_processor.rs @@ -13,13 +13,15 @@ use datadog_trace_obfuscation::obfuscate::obfuscate_span; use datadog_trace_protobuf::pb; use datadog_trace_utils::trace_utils::{self}; use datadog_trace_utils::trace_utils::{EnvironmentType, SendData}; -use datadog_trace_utils::tracer_payload::TraceChunkProcessor; +use datadog_trace_utils::tracer_payload::{TraceChunkProcessor, TracerPayloadCollection}; use crate::{ config::Config, http_utils::{self, log_and_create_http_response, log_and_create_traces_success_http_response}, }; +const TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY: &str = "_dd.tags.function"; + #[async_trait] pub trait TraceProcessor { /// Deserializes traces from a hyper request body and sends them through the provided tokio mpsc @@ -104,7 +106,7 @@ impl TraceProcessor for ServerlessTraceProcessor { ); } - let payload = match trace_utils::collect_pb_trace_chunks( + let mut payload = match trace_utils::collect_pb_trace_chunks( traces, &tracer_header_tags, &mut ChunkProcessor { @@ -122,6 +124,18 @@ impl TraceProcessor for ServerlessTraceProcessor { } }; + // Add function_tags to payload if we can + if let Some(function_tags) = config.tags.function_tags() { + if let TracerPayloadCollection::V07(ref mut tracer_payloads) = payload { + for tracer_payload in tracer_payloads { + tracer_payload.tags.insert( + TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY.to_string(), + function_tags.to_string(), + ); + } + } + } + let send_data = SendData::new(body_size, payload, tracer_header_tags, &config.trace_intake); // send trace payload to our trace flusher @@ -150,8 +164,8 @@ mod tests { use tokio::sync::mpsc::{self, Receiver, Sender}; use crate::{ - config::Config, - trace_processor::{self, TraceProcessor}, + config::{Config, Tags}, + trace_processor::{self, TraceProcessor, TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY}, }; use datadog_trace_protobuf::pb; use datadog_trace_utils::test_utils::{create_test_gcp_json_span, create_test_gcp_span}; @@ -188,6 +202,7 @@ mod tests { os: "linux".to_string(), obfuscation_config: ObfuscationConfig::new().unwrap(), proxy_url: None, + tags: Tags::from_env_string("env:test,service:my-service"), } } @@ -251,7 +266,10 @@ mod tests { tags: HashMap::new(), dropped_trace: false, }], - tags: HashMap::new(), + tags: HashMap::from([( + TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY.to_string(), + "env:test,service:my-service".to_string(), + )]), env: "test-env".to_string(), hostname: "".to_string(), app_version: "".to_string(), @@ -324,7 +342,10 @@ mod tests { tags: HashMap::new(), dropped_trace: false, }], - tags: HashMap::new(), + tags: HashMap::from([( + TRACER_PAYLOAD_FUNCTION_TAGS_TAG_KEY.to_string(), + "env:test,service:my-service".to_string(), + )]), env: "test-env".to_string(), hostname: "".to_string(), app_version: "".to_string(),