Add _dd.tags.function to top level trace tags hashmap. (#516)

* Add _dd.tags.function to top level trace tags hashmap.

* Tests for get_function_tags_map.

Author: purple4reina

bottlecap/src/tags/lambda/tags.rs

@@ -44,6 +44,8 @@ const SERVICE_KEY: &str = "service";
 const COMPUTE_STATS_KEY: &str = "_dd.compute_stats";
 // ComputeStatsValue is the tag value indicating trace stats should be computed
 const COMPUTE_STATS_VALUE: &str = "1";
+// FunctionTagsKey is the tag key for a function's tags to be set on the top level tracepayload
+const FUNCTION_TAGS_KEY: &str = "_dd.tags.function";
 // TODO(astuyve) decide what to do with the version
 const EXTENSION_VERSION_KEY: &str = "dd_extension_version";
 // TODO(duncanista) figure out a better way to not hardcode this
@@ -249,6 +251,17 @@ impl Lambda {
     pub fn get_tags_map(&self) -> &hash_map::HashMap<String, String> {
         &self.tags_map
     }
+
+    #[must_use]
+    pub fn get_function_tags_map(&self) -> hash_map::HashMap<String, String> {
+        let tags = self
+            .tags_map
+            .iter()
+            .map(|(k, v)| format!("{k}:{v}"))
+            .collect::<Vec<String>>()
+            .join(",");
+        hash_map::HashMap::from_iter([(FUNCTION_TAGS_KEY.to_string(), tags)])
+    }
 }
 
 #[cfg(test)]
@@ -371,4 +384,43 @@ mod tests {
         fs::remove_file(path).unwrap();
         assert_eq!(runtime, "provided.al2023");
     }
+
+    #[test]
+    fn test_get_function_tags_map() {
+        let mut metadata = hash_map::HashMap::new();
+        metadata.insert(
+            FUNCTION_ARN_KEY.to_string(),
+            "arn:aws:lambda:us-west-2:123456789012:function:my-function".to_string(),
+        );
+        let config = Arc::new(Config {
+            service: Some("my-service".to_string()),
+            tags: Some("key1:value1,key2:value2".to_string()),
+            env: Some("test".to_string()),
+            version: Some("1.0.0".to_string()),
+            ..Config::default()
+        });
+        let tags = Lambda::new_from_config(config, &metadata);
+        let function_tags = tags.get_function_tags_map();
+        assert_eq!(function_tags.len(), 1);
+        let fn_tags_map: hash_map::HashMap<String, String> = hash_map::HashMap::from_iter(
+            function_tags
+                .get(FUNCTION_TAGS_KEY)
+                .unwrap()
+                .split(',')
+                .map(|tag| {
+                    let parts = tag.split(':').collect::<Vec<&str>>();
+                    (parts[0].to_string(), parts[1].to_string())
+                }),
+        );
+        assert_eq!(fn_tags_map.len(), 14);
+        assert_eq!(fn_tags_map.get("key1").unwrap(), "value1");
+        assert_eq!(fn_tags_map.get("key2").unwrap(), "value2");
+        assert_eq!(fn_tags_map.get(ACCOUNT_ID_KEY).unwrap(), "123456789012");
+        assert_eq!(fn_tags_map.get(ENV_KEY).unwrap(), "test");
+        assert_eq!(fn_tags_map.get(FUNCTION_ARN_KEY).unwrap(), "arn");
+        assert_eq!(fn_tags_map.get(FUNCTION_NAME_KEY).unwrap(), "my-function");
+        assert_eq!(fn_tags_map.get(REGION_KEY).unwrap(), "us-west-2");
+        assert_eq!(fn_tags_map.get(SERVICE_KEY).unwrap(), "my-service");
+        assert_eq!(fn_tags_map.get(VERSION_KEY).unwrap(), "1.0.0");
+    }
 }

bottlecap/src/tags/provider.rs

@@ -56,13 +56,19 @@ impl Provider {
     pub fn get_tags_map(&self) -> &hash_map::HashMap<String, String> {
         self.tag_provider.get_tags_map()
     }
+
+    #[must_use]
+    pub fn get_function_tags_map(&self) -> hash_map::HashMap<String, String> {
+        self.tag_provider.get_function_tags_map()
+    }
 }
 
 trait GetTags {
     fn get_tags_vec(&self) -> Vec<String>;
     fn get_canonical_id(&self) -> Option<String>;
     fn get_canonical_resource_name(&self) -> Option<String>;
     fn get_tags_map(&self) -> &hash_map::HashMap<String, String>;
+    fn get_function_tags_map(&self) -> hash_map::HashMap<String, String>;
 }
 
 impl GetTags for TagProvider {
@@ -89,6 +95,12 @@ impl GetTags for TagProvider {
             TagProvider::Lambda(lambda_tags) => lambda_tags.get_tags_map(),
         }
     }
+
+    fn get_function_tags_map(&self) -> hash_map::HashMap<String, String> {
+        match self {
+            TagProvider::Lambda(lambda_tags) => lambda_tags.get_function_tags_map(),
+        }
+    }
 }
 
 #[cfg(test)]

bottlecap/src/traces/trace_processor.rs

@@ -7,7 +7,9 @@ use datadog_trace_obfuscation::obfuscation_config;
 use datadog_trace_protobuf::pb;
 use datadog_trace_utils::config_utils::trace_intake_url;
 use datadog_trace_utils::tracer_header_tags;
-use datadog_trace_utils::tracer_payload::{TraceChunkProcessor, TraceCollection::V07};
+use datadog_trace_utils::tracer_payload::{
+    TraceChunkProcessor, TraceCollection::V07, TracerPayloadCollection,
+};
 use ddcommon::Endpoint;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -136,7 +138,7 @@ impl TraceProcessor for ServerlessTraceProcessor {
         body_size: usize,
         span_pointers: Option<Vec<SpanPointer>>,
     ) -> SendData {
-        let payload = trace_utils::collect_trace_chunks(
+        let mut payload = trace_utils::collect_trace_chunks(
             V07(traces),
             &header_tags,
             &mut ChunkProcessor {
@@ -146,6 +148,16 @@ impl TraceProcessor for ServerlessTraceProcessor {
             },
             true,
         );
+        match payload {
+            TracerPayloadCollection::V04(_) => {}
+            TracerPayloadCollection::V07(ref mut collection) => {
+                // add function tags to all payloads in this TracerPayloadCollection
+                let tags = tags_provider.get_function_tags_map();
+                for tracer_payload in collection.iter_mut() {
+                    tracer_payload.tags.extend(tags.clone());
+                }
+            }
+        }
         let intake_url = trace_intake_url(&config.site);
         let endpoint = Endpoint {
             url: hyper::Uri::from_str(&intake_url).expect("can't parse trace intake URL, exiting"),
@@ -296,7 +308,7 @@ mod tests {
                 tags: HashMap::new(),
                 dropped_trace: false,
             }],
-            tags: HashMap::new(),
+            tags: tags_provider.get_function_tags_map(),
             env: "test-env".to_string(),
             hostname: String::new(),
             app_version: String::new(),