feat: add parsing of CPU and Memory fields

Author: Ariel Ben-Yehuda

Cargo.lock

@@ -17,6 +17,7 @@ chrono = "0.4"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
 libloading = "0.8"
 reqwest = { version = "0.12", default-features = false, optional = true, features = ["charset", "http2"] }
+ordered-float = { version = "5", features=["serde"]}
 serde_json = "1"
 serde = { version = "1", features = ["derive"] }
 tempfile = "3"
@@ -40,13 +41,16 @@ humantime = "2"
 name = 'simple'
 
 [features]
-default = ["s3", "aws-metadata"]
+default = ["s3", "aws-metadata", "__unstable-fargate-cpu-count"]
 s3 = ["s3-no-defaults", "aws-config/default", "aws-sdk-s3/default"]
 # A version of the s3 feature that does not enable AWS default features
 s3-no-defaults = ["dep:aws-config", "dep:aws-sdk-s3"]
 aws-metadata = ["aws-metadata-no-defaults", "aws-config/default", "reqwest/rustls-tls"]
 # A version of the aws-metadata feature that does not enable AWS default features
 aws-metadata-no-defaults = ["dep:reqwest", "dep:aws-config", "dep:aws-arn"]
+# *Unstable* feature to allow Fargate CPU count. The feature will be stabilized
+# and removed soon.
+__unstable-fargate-cpu-count = []
 
 [package.metadata.docs.rs]
 all-features = true

Cargo.toml

@@ -9,6 +9,8 @@ use reqwest::Method;
 use serde::Deserialize;
 use thiserror::Error;
 
+use crate::metadata::OrderedF64;
+
 use super::AgentMetadata;
 
 /// An error converting Fargate IMDS metadata to Agent metadata. This error
@@ -80,12 +82,25 @@ async fn read_ec2_metadata() -> Result<ImdsEc2InstanceMetadata, AwsProfilerMetad
     Ok(serde_json::from_str(imds_document.as_ref())?)
 }
 
+#[derive(Deserialize, Debug, PartialEq, Eq)]
+struct FargateLimits {
+    #[serde(rename = "CPU")]
+    cpu: Option<OrderedF64>,
+    #[serde(rename = "Memory")]
+    memory: Option<u64>,
+}
+
 #[derive(Deserialize, Debug, PartialEq, Eq)]
 struct FargateMetadata {
     #[serde(rename = "Cluster")]
     cluster: String,
     #[serde(rename = "TaskARN")]
     task_arn: String,
+    // According to <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint-v4-fargate-response.html>
+    // Limits: The resource limits specified at the task levels such as CPU (expressed in vCPUs) and memory.
+    // This parameter is omitted if no resource limits are defined.
+    #[serde(rename = "Limits")]
+    limits: Option<FargateLimits>,
 }
 
 async fn read_fargate_metadata(
@@ -144,6 +159,10 @@ impl super::AgentMetadata {
                 .to_string(),
             ecs_task_arn: fargate_metadata.task_arn,
             ecs_cluster_arn: fargate_metadata.cluster,
+            #[cfg(feature = "__unstable-fargate-cpu-count")]
+            cpu_limit: fargate_metadata.limits.as_ref().and_then(|limits| limits.cpu),
+            #[cfg(feature = "__unstable-fargate-cpu-count")]
+            memory_limit: fargate_metadata.limits.as_ref().and_then(|limits| limits.memory),
         })
     }
 }
@@ -173,6 +192,7 @@ pub async fn load_agent_metadata() -> Result<AgentMetadata, AwsProfilerMetadataE
 #[cfg(test)]
 mod tests {
     use super::*;
+    use test_case::test_case;
 
     // these constants are "anonymized" (aka randomly generated in that format)
 
@@ -237,10 +257,29 @@ mod tests {
         )
     }
 
-    #[test]
-    fn test_fargate_metadata() {
-        let json_str = r#"
-{
+    #[test_case(
+        r#"{
+    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
+    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f"
+}"#,
+        None,
+        None,
+        None
+        ; "no_limits"
+    )]
+    #[test_case(
+        r#"{
+    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
+    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
+    "Limits": {}
+}"#,
+        Some(FargateLimits { cpu: None, memory: None }),
+        None,
+        None
+        ; "empty_limits"
+    )]
+    #[test_case(
+        r#"{
     "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
     "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
     "Family": "profiler-metadata",
@@ -316,17 +355,27 @@ mod tests {
         "Utilized": 208,
         "Reserved": 20496
     }
-}
-"#;
-
-        let fargate_metadata: FargateMetadata = serde_json::from_str(&json_str).unwrap();
+}"#,
+        Some(FargateLimits { cpu: Some(0.25.into()), memory: Some(2048) }),
+        Some(0.25.into()),
+        Some(2048)
+        ; "with_limits"
+    )]
+    fn test_fargate_metadata(
+        json_str: &str,
+        expected_limits: Option<FargateLimits>,
+        _expected_cpu_limit: Option<OrderedF64>,
+        _expected_memory_limit: Option<u64>,
+    ) {
+        let fargate_metadata: FargateMetadata = serde_json::from_str(json_str).unwrap();
 
         assert_eq!(
             fargate_metadata,
             FargateMetadata {
                 cluster: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster"
                     .to_owned(),
                 task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
+                limits: expected_limits,
             }
         );
 
@@ -339,6 +388,10 @@ mod tests {
                 aws_region_id: "us-east-1".to_owned(),
                 ecs_task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
                 ecs_cluster_arn: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster".to_owned(),
+                #[cfg(feature = "__unstable-fargate-cpu-count")]
+                cpu_limit: _expected_cpu_limit,
+                #[cfg(feature = "__unstable-fargate-cpu-count")]
+                memory_limit: _expected_memory_limit,
             }
         )
     }

src/metadata/aws.rs

@@ -3,8 +3,38 @@
 
 //! This module includes ways to get metadata attached to profiling reports.
 
+use ordered_float::OrderedFloat;
 pub use std::time::Duration;
 
+/// An ordered float. Individual type to avoid public API dependencies.
+#[derive(Debug, Copy, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+#[serde(transparent)]
+pub struct OrderedF64(OrderedFloat<f64>);
+
+impl OrderedF64 {
+    /// Create a new `OrderedF64`
+    ///
+    /// ```rust
+    /// # use async_profiler_agent::metadata::OrderedF64;
+    /// let _v = OrderedF64::new(0.0);
+    /// ```
+    pub const fn new(value: f64) -> Self {
+        Self(OrderedFloat(value))
+    }
+}
+
+impl From<f64> for OrderedF64 {
+    fn from(value: f64) -> Self {
+        Self(OrderedFloat(value))
+    }
+}
+
+impl From<OrderedF64> for f64 {
+    fn from(value: OrderedF64) -> Self {
+        value.0 .0
+    }
+}
+
 /// Host Metadata, which describes a host that runs a profiling agent. The current set of supported agent metadata is
 /// AWS-specific. If you are not running on AWS, you can use [AgentMetadata::NoMetadata].
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -13,6 +43,7 @@ pub enum AgentMetadata {
     /// Metadata for an [EC2] instance running on AWS
     ///
     /// [EC2]: https://aws.amazon.com/ec2
+    #[cfg_attr(feature = "__unstable-fargate-cpu-count", non_exhaustive)]
     Ec2AgentMetadata {
         /// The AWS account id
         aws_account_id: String,
@@ -24,6 +55,7 @@ pub enum AgentMetadata {
     /// Metadata for a [Fargate] task running on AWS.
     ///
     /// [Fargate]: https://aws.amazon.com/fargate
+    #[cfg_attr(feature = "__unstable-fargate-cpu-count", non_exhaustive)]
     FargateAgentMetadata {
         /// The AWS account id
         aws_account_id: String,
@@ -41,6 +73,20 @@ pub enum AgentMetadata {
         ///
         /// See the ECS documentation for more details
         ecs_cluster_arn: String,
+        /// The CPU limit for the Fargate cluster
+        ///
+        /// For example, `Some(0.25)`. This will be `None` if the CPU limit is not specified.
+        ///
+        /// See the ECS documentation for more details
+        #[cfg(feature = "__unstable-fargate-cpu-count")]
+        cpu_limit: Option<OrderedF64>,
+        /// The memory limit for the Fargate cluster (in megabytes)
+        ///
+        /// For example, `Some(2048)`. This will be `None` if the memory limit is not specified.
+        ///
+        /// See the ECS documentation for more details
+        #[cfg(feature = "__unstable-fargate-cpu-count")]
+        memory_limit: Option<u64>,
     },
     /// Metadata for a host that is neither an EC2 nor a Fargate
     #[deprecated = "Use AgentMetadata::NoMetadata"]
@@ -99,6 +145,10 @@ impl AgentMetadata {
             aws_region_id,
             ecs_task_arn,
             ecs_cluster_arn,
+            #[cfg(feature = "__unstable-fargate-cpu-count")]
+            cpu_limit: None,
+            #[cfg(feature = "__unstable-fargate-cpu-count")]
+            memory_limit: None,
         }
     }
 }

src/metadata/mod.rs

@@ -132,6 +132,7 @@ fn make_s3_file_name(
             aws_account_id: _,
             aws_region_id: _,
             ec2_instance_id,
+            ..
         } => {
             let ec2_instance_id = ec2_instance_id.replace("/", "-").replace("_", "-");
             format!("ec2_{ec2_instance_id}_")
@@ -141,6 +142,7 @@ fn make_s3_file_name(
             aws_region_id: _,
             ecs_task_arn,
             ecs_cluster_arn: _,
+            ..
         } => {
             let task_arn = ecs_task_arn.replace("/", "-").replace("_", "-");
             format!("ecs_{task_arn}_")
@@ -256,17 +258,17 @@ mod test {
 
     #[test_case(#[allow(deprecated)] { AgentMetadata::Other }, "profile_pg_onprem___<pid>_<time>.zip"; "other")]
     #[test_case(AgentMetadata::NoMetadata, "profile_pg_unknown___<pid>_<time>.zip"; "no-metadata")]
-    #[test_case(AgentMetadata::Ec2AgentMetadata {
-        aws_account_id: "1".into(),
-        aws_region_id: "us-east-1".into(),
-        ec2_instance_id: "i-0".into()
-    }, "profile_pg_ec2_i-0__<pid>_<time>.zip"; "ec2")]
-    #[test_case(AgentMetadata::FargateAgentMetadata {
-        aws_account_id: "1".into(),
-        aws_region_id: "us-east-1".into(),
-        ecs_task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".into(),
-        ecs_cluster_arn: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster".into()
-    }, "profile_pg_ecs_arn:aws:ecs:us-east-1:123456789012:task-profiler-metadata-cluster-5261e761e0e2a3d92da3f02c8e5bab1f__<pid>_<time>.zip"; "ecs")]
+    #[test_case(AgentMetadata::ec2_agent_metadata_v1(
+        "1".into(),
+        "us-east-1".into(),
+        "i-0".into()
+    ), "profile_pg_ec2_i-0__<pid>_<time>.zip"; "ec2")]
+    #[test_case(AgentMetadata::fargate_agent_metadata_v1(
+        "1".into(),
+        "us-east-1".into(),
+        "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".into(),
+        "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster".into(),
+    ), "profile_pg_ecs_arn:aws:ecs:us-east-1:123456789012:task-profiler-metadata-cluster-5261e761e0e2a3d92da3f02c8e5bab1f__<pid>_<time>.zip"; "ecs")]
     fn test_make_s3_file_name(metadata: AgentMetadata, expected: &str) {
         let file_name = super::make_s3_file_name(&metadata, "pg", SystemTime::UNIX_EPOCH);
         assert_eq!(