async_profiler_agent/metadata/
aws.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Contains functions for getting host metadata from [IMDS]
5//!
6//! [IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
7
8use reqwest::Method;
9use serde::Deserialize;
10use thiserror::Error;
11
12use crate::metadata::OrderedF64;
13
14use super::AgentMetadata;
15
16/// An error converting Fargate IMDS metadata to Agent metadata. This error
17/// should probably not happen except in case of a bug in either this crate or IMDS.
18#[derive(Error, Debug)]
19pub enum FargateMetadataToAgentMetadataError {
20    /// unable to parse task ARN as a valid ARN
21    #[error("unable to parse task ARN as a valid ARN")]
22    TaskArnInvalid(#[from] aws_arn::Error),
23    /// AWS account id not found in Fargate metadata
24    #[error("AWS account id not found in Fargate metadata")]
25    AccountIdNotFound,
26    /// AWS region not found in Fargate metadata
27    #[error("AWS region not found in Fargate metadata")]
28    AwsRegionNotFound,
29}
30
31/// An error getting IMDS metadata
32#[derive(Error, Debug)]
33#[error("profiler metadata error: {0}")]
34pub enum AwsProfilerMetadataError {
35    /// Internal IO error
36    #[error("failed to create profiler metadata file: {0}")]
37    FailedToCreateFile(#[from] std::io::Error),
38
39    /// Error parsing IMDS metadata. Should normally not happen except in case of a bug
40    #[error("failed fetching valid Fargate metadata: {0}")]
41    FargateMetadataToAgentMetadataError(#[from] FargateMetadataToAgentMetadataError),
42
43    /// Invalid endpoint URI in `ECS_CONTAINER_METADATA_URI_V4`
44    #[error("retrieved invalid endpoint URI from ECS_CONTAINER_METADATA_URI_V4: {0}")]
45    InvalidUri(String),
46
47    /// Failed to fetch metadata from FarGate endpoint
48    #[error("failed to fetch metadata from endpoint over HTTP: {0}")]
49    FailedToFetchMetadataFromEndpoint(reqwest::Error),
50
51    /// Failed to fetch metadata from IMDS
52    #[error("failed to fetch metadata from IMDS endpoint over HTTP: {0}")]
53    FailedToFetchMetadataFromImds(#[from] aws_config::imds::client::error::ImdsError),
54
55    /// Failed to parse metadata from IMDS - this indicates a bug in this crate
56    /// or in IMDS
57    #[error("failed to parse metadata as valid UTF-8 from endpoint over HTTP: {0}")]
58    FailedToParseMetadataFromEndpoint(reqwest::Error),
59
60    /// Failed to serialize metadata file - this indicates a bug in this crate
61    /// or in IMDS
62    #[error("failed to serialize metadata file: {0}")]
63    FailedToSerializeMetadataFile(#[from] serde_json::Error),
64}
65
66#[derive(Debug, Deserialize, PartialEq, Eq)]
67#[serde(rename_all = "camelCase")]
68struct ImdsEc2InstanceMetadata {
69    account_id: String,
70    region: String,
71    #[allow(dead_code)]
72    instance_type: String,
73    instance_id: String,
74}
75
76async fn read_ec2_metadata() -> Result<ImdsEc2InstanceMetadata, AwsProfilerMetadataError> {
77    let imds = aws_config::imds::Client::builder().build();
78    let imds_document = imds
79        .get("/latest/dynamic/instance-identity/document")
80        .await?;
81
82    Ok(serde_json::from_str(imds_document.as_ref())?)
83}
84
85#[derive(Deserialize, Debug, PartialEq, Eq)]
86struct FargateLimits {
87    #[serde(rename = "CPU")]
88    cpu: Option<OrderedF64>,
89    #[serde(rename = "Memory")]
90    memory: Option<u64>,
91}
92
93#[derive(Deserialize, Debug, PartialEq, Eq)]
94struct FargateMetadata {
95    #[serde(rename = "Cluster")]
96    cluster: String,
97    #[serde(rename = "TaskARN")]
98    task_arn: String,
99    // According to <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint-v4-fargate-response.html>
100    // Limits: The resource limits specified at the task levels such as CPU (expressed in vCPUs) and memory.
101    // This parameter is omitted if no resource limits are defined.
102    #[serde(rename = "Limits")]
103    limits: Option<FargateLimits>,
104}
105
106async fn read_fargate_metadata(
107    http_client: &reqwest::Client,
108) -> Result<FargateMetadata, AwsProfilerMetadataError> {
109    let Ok(md_uri) = std::env::var("ECS_CONTAINER_METADATA_URI_V4") else {
110        return Err(AwsProfilerMetadataError::InvalidUri(
111            "not running on fargate".into(),
112        ));
113    };
114    // Only available if running on Fargate. Something like
115    // `http://169.254.232.106/v4/5261e761e0e2a3d92da3f02c8e5bab1f-3356750833`.
116    let uri = format!("{md_uri}/task",);
117
118    let req = http_client
119        .request(Method::GET, uri.clone())
120        .build()
121        // The only thing that can be invalid about this request is necessarily the URI.
122        .map_err(|_e| AwsProfilerMetadataError::InvalidUri(uri))?;
123    let res = http_client
124        .execute(req)
125        .await
126        .map_err(AwsProfilerMetadataError::FailedToFetchMetadataFromEndpoint)?;
127    let body_str = res
128        .text()
129        .await
130        .map_err(AwsProfilerMetadataError::FailedToParseMetadataFromEndpoint)?;
131
132    Ok(serde_json::from_str(&body_str)?)
133}
134
135impl super::AgentMetadata {
136    fn from_imds_ec2_instance_metadata(
137        imds_ec2_instance_metadata: ImdsEc2InstanceMetadata,
138    ) -> Self {
139        Self::Ec2AgentMetadata {
140            aws_account_id: imds_ec2_instance_metadata.account_id,
141            aws_region_id: imds_ec2_instance_metadata.region,
142            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,
145        }
146    }
147
148    fn try_from_fargate_metadata(
149        fargate_metadata: FargateMetadata,
150    ) -> Result<Self, FargateMetadataToAgentMetadataError> {
151        let ecs_task_arn: aws_arn::ResourceName = fargate_metadata.task_arn.parse()?;
152
153        Ok(Self::FargateAgentMetadata {
154            aws_account_id: ecs_task_arn
155                .account_id
156                .ok_or(FargateMetadataToAgentMetadataError::AccountIdNotFound)?
157                .to_string(),
158            aws_region_id: ecs_task_arn
159                .region
160                .ok_or(FargateMetadataToAgentMetadataError::AwsRegionNotFound)?
161                .to_string(),
162            ecs_task_arn: fargate_metadata.task_arn,
163            ecs_cluster_arn: fargate_metadata.cluster,
164            #[cfg(feature = "__unstable-fargate-cpu-count")]
165            cpu_limit: fargate_metadata
166                .limits
167                .as_ref()
168                .and_then(|limits| limits.cpu),
169            #[cfg(feature = "__unstable-fargate-cpu-count")]
170            memory_limit: fargate_metadata
171                .limits
172                .as_ref()
173                .and_then(|limits| limits.memory),
174        })
175    }
176}
177
178/// Load agent metadata from [Fargate] or [IMDS].
179///
180/// This will return an error if this machine does not appear to be a [Fargate] or [EC2].
181///
182/// [Fargate]: https://aws.amazon.com/fargate
183/// [IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
184/// [EC2]: https://aws.amazon.com/ec2
185pub async fn load_agent_metadata() -> Result<AgentMetadata, AwsProfilerMetadataError> {
186    let agent_metadata: AgentMetadata = match read_ec2_metadata().await {
187        Ok(imds_ec2_instance_metadata) => {
188            AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata)
189        }
190        Err(_) => {
191            let http_client = reqwest::Client::new();
192            let fargate_metadata = read_fargate_metadata(&http_client).await?;
193
194            AgentMetadata::try_from_fargate_metadata(fargate_metadata)?
195        }
196    };
197    Ok(agent_metadata)
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use test_case::test_case;
204
205    // these constants are "anonymized" (aka randomly generated in that format)
206
207    // AMI_ID = 52d4b310b0a459ff
208    // INSTANCE_ID = 92eba08c089f6325
209    // ACCOUNT_ID = 123456789012
210    // IP = 10.65.149.216
211    // DATE = 2025-03-20T16:41:24Z
212    // DATE1 = 2025-03-20T16:41:24.713942268Z
213    // DATE2 = 2025-03-20T16:41:25.623883595Z
214    // TASK_ARN = 5261e761e0e2a3d92da3f02c8e5bab1f
215    // DOCKER_ID = 3356750833
216    // PRIV_IP = 169.254.232.106
217    // IMAGE_ID = ad9d89a36c31afef34c79e05263b06087ad354796cfd90c66ced30f40ea2dbf4
218    // TASK_UUID = f4094744-1b40-4701-9f26-ad84ebb709d7
219    // CLOCK_ERROR = 0.3148955
220
221    #[test]
222    fn test_imds_ec2_metadata() {
223        let json_str = r#"
224{
225    "accountId" : "123456789012",
226    "architecture" : "x86_64",
227    "availabilityZone" : "eu-west-1b",
228    "billingProducts" : null,
229    "devpayProductCodes" : null,
230    "marketplaceProductCodes" : null,
231    "imageId" : "ami-052d4b310b0a459ff",
232    "instanceId" : "i-092eba08c089f6325",
233    "instanceType" : "c5.4xlarge",
234    "kernelId" : null,
235    "pendingTime" : "2025-03-20T16:41:24Z",
236    "privateIp" : "10.65.149.216",
237    "ramdiskId" : null,
238    "region" : "eu-west-1",
239    "version" : "2017-09-30"
240}"#;
241
242        let imds_ec2_instance_metadata: ImdsEc2InstanceMetadata =
243            serde_json::from_str(&json_str).unwrap();
244
245        assert_eq!(
246            imds_ec2_instance_metadata,
247            ImdsEc2InstanceMetadata {
248                account_id: "123456789012".to_owned(),
249                region: "eu-west-1".to_owned(),
250                instance_type: "c5.4xlarge".to_owned(),
251                instance_id: "i-092eba08c089f6325".to_owned(),
252            }
253        );
254
255        let agent_metadata =
256            AgentMetadata::from_imds_ec2_instance_metadata(imds_ec2_instance_metadata);
257
258        let expected = AgentMetadata::ec2_agent_metadata(
259            "123456789012".to_owned(),
260            "eu-west-1".to_owned(),
261            "i-092eba08c089f6325".to_owned(),
262        )
263        .with_ec2_instance_type("c5.4xlarge".to_owned())
264        .build();
265
266        assert_eq!(agent_metadata, expected);
267    }
268
269    #[test_case(
270        r#"{
271    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
272    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f"
273}"#,
274        None,
275        None,
276        None
277        ; "no_limits"
278    )]
279    #[test_case(
280        r#"{
281    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
282    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
283    "Limits": {}
284}"#,
285        Some(FargateLimits { cpu: None, memory: None }),
286        None,
287        None
288        ; "empty_limits"
289    )]
290    #[test_case(
291        r#"{
292    "Cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
293    "TaskARN": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
294    "Family": "profiler-metadata",
295    "Revision": "1",
296    "DesiredStatus": "RUNNING",
297    "KnownStatus": "NONE",
298    "Limits": {
299        "CPU": 0.25,
300        "Memory": 2048
301    },
302    "PullStartedAt": "2025-03-20T16:41:24.713942268Z",
303    "PullStoppedAt": "2025-03-20T16:41:25.623883595Z",
304    "AvailabilityZone": "us-east-1f",
305    "LaunchType": "FARGATE",
306    "Containers": [
307        {
308            "DockerId": "5261e761e0e2a3d92da3f02c8e5bab1f-3356750833",
309            "Name": "profiler-metadata",
310            "DockerName": "profiler-metadata",
311            "Image": "123456789012.dkr.ecr.us-east-1.amazonaws.com/profiler-metadata",
312            "ImageID": "sha256:ad9d89a36c31afef34c79e05263b06087ad354796cfd90c66ced30f40ea2dbf4",
313            "Labels": {
314                "com.amazonaws.ecs.cluster": "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster",
315                "com.amazonaws.ecs.container-name": "profiler-metadata",
316                "com.amazonaws.ecs.task-arn": "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f",
317                "com.amazonaws.ecs.task-definition-family": "profiler-metadata",
318                "com.amazonaws.ecs.task-definition-version": "1"
319            },
320            "DesiredStatus": "RUNNING",
321            "KnownStatus": "PULLED",
322            "Limits": {
323                "CPU": 0
324            },
325            "Type": "NORMAL",
326            "LogDriver": "awslogs",
327            "LogOptions": {
328                "awslogs-create-group": "true",
329                "awslogs-group": "/ecs/profiler-metadata",
330                "awslogs-region": "us-east-1",
331                "awslogs-stream": "ecs/profiler-metadata/5261e761e0e2a3d92da3f02c8e5bab1f",
332                "max-buffer-size": "25m",
333                "mode": "non-blocking"
334            },
335            "ContainerARN": "arn:aws:ecs:us-east-1:123456789012:container/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f/f4094744-1b40-4701-9f26-ad84ebb709d7",
336            "Networks": [
337                {
338                    "NetworkMode": "awsvpc",
339                    "IPv4Addresses": [
340                        "172.31.233.169"
341                    ],
342                    "AttachmentIndex": 0,
343                    "MACAddress": "16:ff:d6:e1:dc:99",
344                    "IPv4SubnetCIDRBlock": "172.31.192.0/20",
345                    "DomainNameServers": [
346                        "172.31.0.2"
347                    ],
348                    "DomainNameSearchList": [
349                        "ec2.internal"
350                    ],
351                    "PrivateDNSName": "ip-172-31-233-169.ec2.internal",
352                    "SubnetGatewayIpv4Address": "172.31.192.1/20"
353                }
354            ],
355            "Snapshotter": "overlayfs"
356        }
357    ],
358    "ClockDrift": {
359        "ClockErrorBound": 0.3148955,
360        "ReferenceTimestamp": "2025-03-20T16:41:24Z",
361        "ClockSynchronizationStatus": "SYNCHRONIZED"
362    },
363    "EphemeralStorageMetrics": {
364        "Utilized": 208,
365        "Reserved": 20496
366    }
367}"#,
368        Some(FargateLimits { cpu: Some(0.25.into()), memory: Some(2048) }),
369        Some(0.25.into()),
370        Some(2048)
371        ; "with_limits"
372    )]
373    fn test_fargate_metadata(
374        json_str: &str,
375        expected_limits: Option<FargateLimits>,
376        _expected_cpu_limit: Option<OrderedF64>,
377        _expected_memory_limit: Option<u64>,
378    ) {
379        let fargate_metadata: FargateMetadata = serde_json::from_str(json_str).unwrap();
380
381        assert_eq!(
382            fargate_metadata,
383            FargateMetadata {
384                cluster: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster"
385                    .to_owned(),
386                task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
387                limits: expected_limits,
388            }
389        );
390
391        let agent_metadata = AgentMetadata::try_from_fargate_metadata(fargate_metadata).unwrap();
392
393        assert_eq!(
394            agent_metadata,
395            AgentMetadata::FargateAgentMetadata {
396                aws_account_id: "123456789012".to_owned(),
397                aws_region_id: "us-east-1".to_owned(),
398                ecs_task_arn: "arn:aws:ecs:us-east-1:123456789012:task/profiler-metadata-cluster/5261e761e0e2a3d92da3f02c8e5bab1f".to_owned(),
399                ecs_cluster_arn: "arn:aws:ecs:us-east-1:123456789012:cluster/profiler-metadata-cluster".to_owned(),
400                #[cfg(feature = "__unstable-fargate-cpu-count")]
401                cpu_limit: _expected_cpu_limit,
402                #[cfg(feature = "__unstable-fargate-cpu-count")]
403                memory_limit: _expected_memory_limit,
404            }
405        )
406    }
407}
async_profiler_agent/metadata/aws.rs

async_profiler_agent/metadata/
aws.rs
