diff --git a/src/metadata/aws.rs b/src/metadata/aws.rs index 0ee4693..c4eb558 100644 --- a/src/metadata/aws.rs +++ b/src/metadata/aws.rs @@ -140,6 +140,8 @@ impl super::AgentMetadata { aws_account_id: imds_ec2_instance_metadata.account_id, aws_region_id: imds_ec2_instance_metadata.region, ec2_instance_id: imds_ec2_instance_metadata.instance_id, + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: imds_ec2_instance_metadata.instance_type, } } @@ -253,14 +255,15 @@ mod tests { let agent_metadata = AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata); - assert_eq!( - agent_metadata, - AgentMetadata::Ec2AgentMetadata { - aws_account_id: "123456789012".to_owned(), - aws_region_id: "eu-west-1".to_owned(), - ec2_instance_id: "i-092eba08c089f6325".to_owned(), - } + let expected = AgentMetadata::ec2_agent_metadata( + "123456789012".to_owned(), + "eu-west-1".to_owned(), + "i-092eba08c089f6325".to_owned(), ) + .with_ec2_instance_type("c5.4xlarge".to_owned()) + .build(); + + assert_eq!(agent_metadata, expected); } #[test_case( diff --git a/src/metadata/mod.rs b/src/metadata/mod.rs index 8248b69..5d55497 100644 --- a/src/metadata/mod.rs +++ b/src/metadata/mod.rs @@ -51,6 +51,9 @@ pub enum AgentMetadata { aws_region_id: String, /// The EC2 instance id ec2_instance_id: String, + #[cfg(feature = "__unstable-fargate-cpu-count")] + /// The EC2 instance type + ec2_instance_type: String, }, /// Metadata for a [Fargate] task running on AWS. /// @@ -130,6 +133,7 @@ impl AgentMetadata { aws_account_id, aws_region_id, ec2_instance_id, + ec2_instance_type: None, } } @@ -184,15 +188,24 @@ pub struct Ec2AgentMetadataBuilder { aws_account_id: String, aws_region_id: String, ec2_instance_id: String, + ec2_instance_type: Option, } impl Ec2AgentMetadataBuilder { + /// Set the EC2 instance type + pub fn with_ec2_instance_type(mut self, ec2_instance_type: String) -> Self { + self.ec2_instance_type = Some(ec2_instance_type); + self + } + /// Build the AgentMetadata pub fn build(self) -> AgentMetadata { AgentMetadata::Ec2AgentMetadata { aws_account_id: self.aws_account_id, aws_region_id: self.aws_region_id, ec2_instance_id: self.ec2_instance_id, + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: self.ec2_instance_type.unwrap_or_default(), } } } diff --git a/src/profiler.rs b/src/profiler.rs index 9bd0792..ad2eae7 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -1217,6 +1217,8 @@ mod tests { aws_account_id: "0".into(), aws_region_id: "us-east-1".into(), ec2_instance_id: "i-fake".into(), + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: "t3.micro".into(), }) .build(); (agent, rx) @@ -1228,6 +1230,8 @@ mod tests { aws_account_id: "0".into(), aws_region_id: "us-east-1".into(), ec2_instance_id: "i-fake".into(), + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: "t3.micro".into(), }; let (agent, mut rx) = make_mock_profiler(); agent @@ -1251,6 +1255,8 @@ mod tests { aws_account_id: "0".into(), aws_region_id: "us-east-1".into(), ec2_instance_id: "i-fake".into(), + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: "t3.micro".into(), }; let (agent, mut rx) = make_mock_profiler(); let rt = tokio::runtime::Builder::new_current_thread() @@ -1317,6 +1323,8 @@ mod tests { aws_account_id: "0".into(), aws_region_id: "us-east-1".into(), ec2_instance_id: "i-fake".into(), + #[cfg(feature = "__unstable-fargate-cpu-count")] + ec2_instance_type: "t3.micro".into(), }; let (agent, mut rx) = make_mock_profiler(); let profiler_ref = agent