Skip to content

Commit 1215937

Browse files
authored
Merge pull request #111 from arielb1/ec2-instance-type
feat: pass ec2 instance type from imds in EC2 instance metadata
2 parents ede48a5 + 3375cd9 commit 1215937

File tree

3 files changed

+31
-7
lines changed

3 files changed

+31
-7
lines changed

src/metadata/aws.rs

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,8 @@ impl super::AgentMetadata {
140140
aws_account_id: imds_ec2_instance_metadata.account_id,
141141
aws_region_id: imds_ec2_instance_metadata.region,
142142
ec2_instance_id: imds_ec2_instance_metadata.instance_id,
143+
#[cfg(feature = "__unstable-fargate-cpu-count")]
144+
ec2_instance_type: imds_ec2_instance_metadata.instance_type,
143145
}
144146
}
145147

@@ -253,14 +255,15 @@ mod tests {
253255
let agent_metadata =
254256
AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata);
255257

256-
assert_eq!(
257-
agent_metadata,
258-
AgentMetadata::Ec2AgentMetadata {
259-
aws_account_id: "123456789012".to_owned(),
260-
aws_region_id: "eu-west-1".to_owned(),
261-
ec2_instance_id: "i-092eba08c089f6325".to_owned(),
262-
}
258+
let expected = AgentMetadata::ec2_agent_metadata(
259+
"123456789012".to_owned(),
260+
"eu-west-1".to_owned(),
261+
"i-092eba08c089f6325".to_owned(),
263262
)
263+
.with_ec2_instance_type("c5.4xlarge".to_owned())
264+
.build();
265+
266+
assert_eq!(agent_metadata, expected);
264267
}
265268

266269
#[test_case(

src/metadata/mod.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,9 @@ pub enum AgentMetadata {
5151
aws_region_id: String,
5252
/// The EC2 instance id
5353
ec2_instance_id: String,
54+
#[cfg(feature = "__unstable-fargate-cpu-count")]
55+
/// The EC2 instance type
56+
ec2_instance_type: String,
5457
},
5558
/// Metadata for a [Fargate] task running on AWS.
5659
///
@@ -130,6 +133,7 @@ impl AgentMetadata {
130133
aws_account_id,
131134
aws_region_id,
132135
ec2_instance_id,
136+
ec2_instance_type: None,
133137
}
134138
}
135139

@@ -184,15 +188,24 @@ pub struct Ec2AgentMetadataBuilder {
184188
aws_account_id: String,
185189
aws_region_id: String,
186190
ec2_instance_id: String,
191+
ec2_instance_type: Option<String>,
187192
}
188193

189194
impl Ec2AgentMetadataBuilder {
195+
/// Set the EC2 instance type
196+
pub fn with_ec2_instance_type(mut self, ec2_instance_type: String) -> Self {
197+
self.ec2_instance_type = Some(ec2_instance_type);
198+
self
199+
}
200+
190201
/// Build the AgentMetadata
191202
pub fn build(self) -> AgentMetadata {
192203
AgentMetadata::Ec2AgentMetadata {
193204
aws_account_id: self.aws_account_id,
194205
aws_region_id: self.aws_region_id,
195206
ec2_instance_id: self.ec2_instance_id,
207+
#[cfg(feature = "__unstable-fargate-cpu-count")]
208+
ec2_instance_type: self.ec2_instance_type.unwrap_or_default(),
196209
}
197210
}
198211
}

src/profiler.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1217,6 +1217,8 @@ mod tests {
12171217
aws_account_id: "0".into(),
12181218
aws_region_id: "us-east-1".into(),
12191219
ec2_instance_id: "i-fake".into(),
1220+
#[cfg(feature = "__unstable-fargate-cpu-count")]
1221+
ec2_instance_type: "t3.micro".into(),
12201222
})
12211223
.build();
12221224
(agent, rx)
@@ -1228,6 +1230,8 @@ mod tests {
12281230
aws_account_id: "0".into(),
12291231
aws_region_id: "us-east-1".into(),
12301232
ec2_instance_id: "i-fake".into(),
1233+
#[cfg(feature = "__unstable-fargate-cpu-count")]
1234+
ec2_instance_type: "t3.micro".into(),
12311235
};
12321236
let (agent, mut rx) = make_mock_profiler();
12331237
agent
@@ -1251,6 +1255,8 @@ mod tests {
12511255
aws_account_id: "0".into(),
12521256
aws_region_id: "us-east-1".into(),
12531257
ec2_instance_id: "i-fake".into(),
1258+
#[cfg(feature = "__unstable-fargate-cpu-count")]
1259+
ec2_instance_type: "t3.micro".into(),
12541260
};
12551261
let (agent, mut rx) = make_mock_profiler();
12561262
let rt = tokio::runtime::Builder::new_current_thread()
@@ -1317,6 +1323,8 @@ mod tests {
13171323
aws_account_id: "0".into(),
13181324
aws_region_id: "us-east-1".into(),
13191325
ec2_instance_id: "i-fake".into(),
1326+
#[cfg(feature = "__unstable-fargate-cpu-count")]
1327+
ec2_instance_type: "t3.micro".into(),
13201328
};
13211329
let (agent, mut rx) = make_mock_profiler();
13221330
let profiler_ref = agent

0 commit comments

Comments
 (0)