[SVLS-3545] Add Serverless metric origins to dogstatsd package

dogstatsd/src/aggregator.rs

@@ -4,15 +4,18 @@
 //! The aggregation of metrics.
 
 use crate::constants;
-use crate::datadog::{self, Metric as MetricToShip, Series};
+use crate::datadog::{
+    self, Metadata as MetadataToShip, Metric as MetricToShip, Origin as OriginToShip, Series,
+};
 use crate::errors;
 use crate::metric::{self, Metric, MetricValue, SortedTags};
+use crate::origin::find_metric_origin;
 use std::time;
 
-use datadog_protos::metrics::{Dogsketch, Sketch, SketchPayload};
+use datadog_protos::metrics::{Dogsketch, Metadata, Sketch, SketchPayload};
 use ddsketch_agent::DDSketch;
 use hashbrown::hash_table;
-use protobuf::Message;
+use protobuf::{Message, MessageField, SpecialFields};
 use tracing::{error, warn};
 use ustr::Ustr;
 
@@ -268,6 +271,14 @@ fn build_sketch(now: i64, entry: &Metric, mut base_tag_vec: SortedTags) -> Optio
         base_tag_vec.extend(&tags);
     }
     sketch.set_tags(base_tag_vec.to_chars());
+
+    if let Some(origin) = find_metric_origin(entry, base_tag_vec) {
+        sketch.set_metadata(Metadata {
+            origin: MessageField::some(origin),
+            special_fields: SpecialFields::default(),
+        });
+    }
+
     Some(sketch)
 }
 
@@ -296,12 +307,22 @@ fn build_metric(entry: &Metric, mut base_tag_vec: SortedTags) -> Option<MetricTo
         base_tag_vec.extend(&tags);
     }
 
+    let origin = find_metric_origin(entry, base_tag_vec.clone());
+    let metadata = origin.map(|o| MetadataToShip {
+        origin: Some(OriginToShip {
+            origin_product: o.origin_product,
+            origin_sub_product: o.origin_category,
+            origin_product_detail: o.origin_service,
+        }),
+    });
+
     Some(MetricToShip {
         metric: entry.name.as_str(),
         resources,
         kind,
         points: [point; 1],
         tags: base_tag_vec.to_strings(),
+        metadata,
     })
 }
 
@@ -318,7 +339,7 @@ pub mod tests {
 
     const PRECISION: f64 = 0.000_000_01;
 
-    const SINGLE_METRIC_SIZE: usize = 193; // taken from the test, size of a serialized metric with one tag and 1 digit counter value
+    const SINGLE_METRIC_SIZE: usize = 209; // taken from the test, size of a serialized metric with one tag and 1 digit counter value
     const SINGLE_DISTRIBUTION_SIZE: u64 = 140;
     const DEFAULT_TAGS: &str =
         "dd_extension_version:63-next,architecture:x86_64,_dd.compute_stats:1";
@@ -664,7 +685,7 @@ pub mod tests {
     fn consume_metrics_batch_bytes() {
         let expected_metrics_per_batch = 2;
         let total_number_of_metrics = 5;
-        let two_metrics_size = 374;
+        let two_metrics_size = 406;
         let max_bytes = SINGLE_METRIC_SIZE * expected_metrics_per_batch + 13;
         let mut aggregator = Aggregator {
             tags: to_sorted_tags(),

dogstatsd/src/datadog.rs

@@ -283,6 +283,20 @@ pub(crate) struct Metric {
     /// The kind of metric
     pub(crate) kind: DdMetricKind,
     pub(crate) tags: Vec<String>,
+    /// Optional metadata associated with the metric
+    pub(crate) metadata: Option<Metadata>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct Metadata {
+    pub(crate) origin: Option<Origin>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct Origin {
+    pub(crate) origin_product: u32,
+    pub(crate) origin_sub_product: u32,
+    pub(crate) origin_product_detail: u32,
 }
 
 #[derive(Debug, Serialize)]

dogstatsd/src/lib.rs

@@ -8,3 +8,4 @@ pub mod dogstatsd;
 pub mod errors;
 pub mod flusher;
 pub mod metric;
+pub mod origin;

dogstatsd/src/metric.rs

@@ -117,6 +117,13 @@ impl SortedTags {
         tags_as_vec
     }
 
+    pub fn find_all(&self, tag_key: &str) -> Vec<&Ustr> {
+        self.values
+            .iter()
+            .filter_map(|(k, v)| if k == tag_key { Some(v) } else { None })
+            .collect()
+    }
+
     pub(crate) fn to_resources(&self) -> Vec<datadog::Resource> {
         let mut resources = Vec::with_capacity(constants::MAX_TAGS);
         for (key, val) in &self.values {
@@ -561,4 +568,13 @@ mod tests {
         assert_eq!(first_element.0, Ustr::from("a"));
         assert_eq!(first_element.1, Ustr::from("a1"));
     }
+
+    #[test]
+    fn sorted_tags_find_all() {
+        let tags = SortedTags::parse("a,a:1,b:2,c:3").unwrap();
+        assert_eq!(tags.find_all("a"), vec![&Ustr::from(""), &Ustr::from("1")]);
+        assert_eq!(tags.find_all("b"), vec![&Ustr::from("2")]);
+        assert_eq!(tags.find_all("c"), vec![&Ustr::from("3")]);
+        assert_eq!(tags.find_all("d"), Vec::<&Ustr>::new());
+    }
 }

dogstatsd/src/origin.rs

@@ -0,0 +1,244 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::metric::{Metric, SortedTags};
+use datadog_protos::metrics::Origin;
+
+// Metric tag keys
+const DD_ORIGIN_TAG_KEY: &str = "origin";
+const AWS_LAMBDA_TAG_KEY: &str = "function_arn";
+const AWS_STEP_FUNCTIONS_TAG_KEY: &str = "statemachinearn";
+
+// Metric tag values
+const GOOGLE_CLOUD_RUN_TAG_VALUE: &str = "cloudrun";
+const AZURE_APP_SERVICES_TAG_VALUE: &str = "appservice";
+const AZURE_CONTAINER_APP_TAG_VALUE: &str = "containerapp";
+const AZURE_FUNCTIONS_TAG_VALUE: &str = "azurefunction";
+
+// Metric prefixes
+const DATADOG_PREFIX: &str = "datadog";
+const GOOGLE_CLOUD_RUN_PREFIX: &str = "gcp.run";
+const AZURE_APP_SERVICES_PREFIX: &str = "azure.app_services";
+const AZURE_CONTAINER_APP_PREFIX: &str = "azure.app_containerapps";
+const AZURE_FUNCTIONS_PREFIX: &str = "azure.functions";
+const AWS_LAMBDA_PREFIX: &str = "aws.lambda";
+const AWS_STEP_FUNCTIONS_PREFIX: &str = "aws.states";
+
+/// Represents the product origin of a metric.
+/// The full enum is exhaustive so we only include what we need. Please reference the corresponding
+/// enum for all possible values https://github.com/DataDog/dd-source/blob/573dee9b5f7ee13935cb3ad11b16dde970528983/domains/metrics/shared/libs/proto/origin/origin.proto#L161
+pub enum OriginProduct {
+    Other = 0,
+    Serverless = 1,
+}
+
+impl From<OriginProduct> for u32 {
+    fn from(product: OriginProduct) -> u32 {
+        product as u32
+    }
+}
+
+/// Represents the category origin of a metric.
+/// The full enum is exhaustive so we only include what we need. Please reference the corresponding
+/// enum for all possible values https://github.com/DataDog/dd-source/blob/573dee9b5f7ee13935cb3ad11b16dde970528983/domains/metrics/shared/libs/proto/origin/origin.proto#L276
+pub enum OriginCategory {
+    Other = 0,
+    AppServicesMetrics = 35,
+    CloudRunMetrics = 36,
+    ContainerAppMetrics = 37,
+    LambdaMetrics = 38,
+    StepFunctionsMetrics = 41,
+    AzureFunctionsMetrics = 71,
+}
+
+impl From<OriginCategory> for u32 {
+    fn from(category: OriginCategory) -> u32 {
+        category as u32
+    }
+}
+
+/// Represents the service origin of a metric.
+/// The full enum is exhaustive so we only include what we need. Please reference the corresponding
+/// enum for all possible values https://github.com/DataDog/dd-source/blob/573dee9b5f7ee13935cb3ad11b16dde970528983/domains/metrics/shared/libs/proto/origin/origin.proto#L417
+pub enum OriginService {
+    Other = 0,
+}
+
+impl From<OriginService> for u32 {
+    fn from(service: OriginService) -> u32 {
+        service as u32
+    }
+}
+
+pub fn find_metric_origin(metric: &Metric, tags: SortedTags) -> Option<Origin> {
+    let name = metric.name.to_string();
+    let prefix = name.split('.').take(2).collect::<Vec<&str>>().join(".");
+
+    if is_datadog_metric(&prefix) {
+        return None;
+    }
+    if is_azure_app_services(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::AppServicesMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+    if is_google_cloud_run(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::CloudRunMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+    if is_azure_container_app(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::ContainerAppMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+    if is_azure_functions(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::AzureFunctionsMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+    if is_aws_lambda(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::LambdaMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+    if is_aws_step_functions(&tags, &prefix) {
+        return Some(Origin {
+            origin_product: OriginProduct::Serverless.into(),
+            origin_category: OriginCategory::StepFunctionsMetrics.into(),
+            origin_service: OriginService::Other.into(),
+            ..Default::default()
+        });
+    }
+
+    None
+}
+
+fn get_first_tag_value<'a>(tags: &'a SortedTags, key: &str) -> Option<&'a str> {
+    tags.find_all(key)
+        .iter()
+        .filter_map(|value| {
+            if !value.is_empty() {
+                Some(value.as_str())
+            } else {
+                None
+            }
+        })
+        .next()
+}
+
+fn is_datadog_metric(prefix: &str) -> bool {
+    prefix == DATADOG_PREFIX
+}
+
+fn is_google_cloud_run(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, DD_ORIGIN_TAG_KEY) == Some(GOOGLE_CLOUD_RUN_TAG_VALUE)
+        && prefix != GOOGLE_CLOUD_RUN_PREFIX
+}
+
+fn is_azure_app_services(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, DD_ORIGIN_TAG_KEY) == Some(AZURE_APP_SERVICES_TAG_VALUE)
+        && prefix != AZURE_APP_SERVICES_PREFIX
+}
+
+fn is_azure_container_app(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, DD_ORIGIN_TAG_KEY) == Some(AZURE_CONTAINER_APP_TAG_VALUE)
+        && prefix != AZURE_CONTAINER_APP_PREFIX
+}
+
+fn is_azure_functions(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, DD_ORIGIN_TAG_KEY) == Some(AZURE_FUNCTIONS_TAG_VALUE)
+        && prefix != AZURE_FUNCTIONS_PREFIX
+}
+
+fn is_aws_lambda(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, AWS_LAMBDA_TAG_KEY).is_some() && prefix != AWS_LAMBDA_PREFIX
+}
+
+fn is_aws_step_functions(tags: &SortedTags, prefix: &str) -> bool {
+    get_first_tag_value(tags, AWS_STEP_FUNCTIONS_TAG_KEY).is_some()
+        && prefix != AWS_STEP_FUNCTIONS_PREFIX
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::metric::MetricValue;
+
+    use super::*;
+
+    #[test]
+    fn test_origin_product() {
+        let origin_product: u32 = OriginProduct::Serverless.into();
+        assert_eq!(origin_product, 1);
+    }
+
+    #[test]
+    fn test_origin_category() {
+        let origin_category: u32 = OriginCategory::LambdaMetrics.into();
+        assert_eq!(origin_category, 38);
+    }
+
+    #[test]
+    fn test_origin_service() {
+        let origin_service: u32 = OriginService::Other.into();
+        assert_eq!(origin_service, 0);
+    }
+
+    #[test]
+    fn test_find_metric_origin_aws_lambda_standard_metric() {
+        let tags = SortedTags::parse("function_arn:hello123").unwrap();
+        let metric = Metric {
+            id: 0,
+            name: "aws.lambda.enhanced.invocations".into(),
+            value: MetricValue::Gauge(1.0),
+            tags: Some(tags.clone()),
+        };
+        let origin = find_metric_origin(&metric, tags);
+        assert_eq!(origin, None);
+    }
+
+    #[test]
+    fn test_find_metric_origin_aws_lambda_custom_metric() {
+        let tags = SortedTags::parse("function_arn:hello123").unwrap();
+        let metric = Metric {
+            id: 0,
+            name: "my.custom.aws.lambda.invocations".into(),
+            value: MetricValue::Gauge(1.0),
+            tags: Some(tags.clone()),
+        };
+        let origin = find_metric_origin(&metric, tags);
+        assert_eq!(
+            origin,
+            Some(Origin {
+                origin_product: OriginProduct::Serverless.into(),
+                origin_category: OriginCategory::LambdaMetrics.into(),
+                origin_service: OriginService::Other.into(),
+                ..Default::default()
+            })
+        );
+    }
+
+    #[test]
+    fn test_get_first_tag_value() {
+        let tags = SortedTags::parse("a,a:1,b:2,c:3").unwrap();
+        assert_eq!(get_first_tag_value(&tags, "a"), Some("1"));
+        assert_eq!(get_first_tag_value(&tags, "b"), Some("2"));
+        assert_eq!(get_first_tag_value(&tags, "c"), Some("3"));
+        assert_eq!(get_first_tag_value(&tags, "d"), None);
+    }
+}