Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions src/metadata/aws.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down Expand Up @@ -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(
Expand Down
13 changes: 13 additions & 0 deletions src/metadata/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -130,6 +133,7 @@ impl AgentMetadata {
aws_account_id,
aws_region_id,
ec2_instance_id,
ec2_instance_type: None,
}
}

Expand Down Expand Up @@ -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(),
}
}
}
Expand Down
8 changes: 8 additions & 0 deletions src/profiler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down
Loading