feat: pass ec2 instance type from imds in EC2 instance metadata

Author: Ariel Ben-Yehuda

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(

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

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