feat: parse DD_TAGS and use them for the tracer payload function tags (#29)

* feat: add tags to config

* feat: precompute function_tags

* feat: inject function tags into the tracer payloads

* chore: better tags type

* feat: support space and comma separated tags

Author: apiarian-datadog

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<String, String>,
+    function_tags_string: OnceLock<String>,
+}
+
+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::<Vec<&str>>();
+            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<String, String> {
+        &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::<Vec<String>>();
+            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);
+    }
 }

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(),