aws_manager/ec2/
mod.rs

1pub mod disk;
2pub mod metadata;
3pub mod plugins;
4
5use std::{
6    collections::{HashMap, HashSet},
7    fs::{self, File},
8    io::{self, Write},
9    path::Path,
10    str::FromStr,
11};
12
13use crate::errors::{self, Error, Result};
14use aws_sdk_ec2::{
15    operation::delete_key_pair::DeleteKeyPairError,
16    types::{
17        Address, AttachmentStatus, Filter, Image, ImageState, Instance, InstanceState,
18        InstanceStateName, KeyFormat, KeyType, ResourceType, SecurityGroup, Subnet, Tag,
19        TagSpecification, Volume, VolumeAttachmentState, VolumeState, Vpc,
20    },
21    Client,
22};
23use aws_smithy_client::SdkError;
24use aws_types::SdkConfig as AwsSdkConfig;
25use chrono::{DateTime, NaiveDateTime, Utc};
26use serde::{Deserialize, Serialize};
27use tokio::time::{sleep, Duration, Instant};
28
29/// Defines the Arch type.
30#[derive(
31    Deserialize,
32    Serialize,
33    std::clone::Clone,
34    std::cmp::Eq,
35    std::cmp::Ord,
36    std::cmp::PartialEq,
37    std::cmp::PartialOrd,
38    std::fmt::Debug,
39    std::hash::Hash,
40)]
41pub enum ArchType {
42    #[serde(rename = "amd64")]
43    Amd64,
44    #[serde(rename = "arm64")]
45    Arm64,
46
47    /// For p4 instances.
48    /// ref. <https://aws.amazon.com/ec2/instance-types/p4/>
49    #[serde(rename = "amd64-gpu-p4-nvidia-tesla-a100")]
50    Amd64GpuP4NvidiaTeslaA100,
51    /// For g3 instances.
52    /// ref. <https://aws.amazon.com/ec2/instance-types/g3/>
53    #[serde(rename = "amd64-gpu-g3-nvidia-tesla-m60")]
54    Amd64GpuG3NvidiaTeslaM60,
55    /// For g4dn instances.
56    /// ref. <https://aws.amazon.com/ec2/instance-types/g4/>
57    #[serde(rename = "amd64-gpu-g4dn-nvidia-t4")]
58    Amd64GpuG4dnNvidiaT4,
59    /// For g4 instances.
60    /// ref. <https://aws.amazon.com/ec2/instance-types/g4/>
61    #[serde(rename = "amd64-gpu-g4ad-radeon")]
62    Amd64GpuG4adRadeon,
63    /// For g5 instances.
64    /// ref. <https://aws.amazon.com/ec2/instance-types/g5/>
65    #[serde(rename = "amd64-gpu-g5-nvidia-a10g")]
66    Amd64GpuG5NvidiaA10G,
67
68    /// For inf1 instances.
69    /// ref. <https://aws.amazon.com/ec2/instance-types/inf1/>
70    #[serde(rename = "amd64-gpu-inf1")]
71    Amd64GpuInf1,
72
73    /// For trn1 instances.
74    /// ref. <https://aws.amazon.com/ec2/instance-types/trn1/>
75    #[serde(rename = "amd64-gpu-trn1")]
76    Amd64GpuTrn1,
77
78    Unknown(String),
79}
80
81impl std::convert::From<&str> for ArchType {
82    fn from(s: &str) -> Self {
83        match s {
84            "amd64" => ArchType::Amd64,
85            "arm64" => ArchType::Arm64,
86            "amd64-gpu-p4-nvidia-tesla-a100" => ArchType::Amd64GpuP4NvidiaTeslaA100,
87            "amd64-gpu-g3-nvidia-tesla-m60" => ArchType::Amd64GpuG3NvidiaTeslaM60,
88            "amd64-gpu-g4dn-nvidia-t4" => ArchType::Amd64GpuG4dnNvidiaT4,
89            "amd64-gpu-g4ad-radeon" => ArchType::Amd64GpuG4adRadeon,
90            "amd64-gpu-g5-nvidia-a10g" => ArchType::Amd64GpuG5NvidiaA10G,
91            "amd64-gpu-inf1" => ArchType::Amd64GpuInf1,
92            "amd64-gpu-trn1" => ArchType::Amd64GpuTrn1,
93            other => ArchType::Unknown(other.to_owned()),
94        }
95    }
96}
97
98impl std::str::FromStr for ArchType {
99    type Err = std::convert::Infallible;
100
101    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
102        Ok(ArchType::from(s))
103    }
104}
105
106impl ArchType {
107    /// Returns the `&str` value of the enum member.
108    pub fn as_str(&self) -> &str {
109        match self {
110            ArchType::Amd64 => "amd64",
111            ArchType::Arm64 => "arm64",
112            ArchType::Amd64GpuP4NvidiaTeslaA100 => "amd64-gpu-p4-nvidia-tesla-a100",
113            ArchType::Amd64GpuG3NvidiaTeslaM60 => "amd64-gpu-g3-nvidia-tesla-m60",
114            ArchType::Amd64GpuG4dnNvidiaT4 => "amd64-gpu-g4dn-nvidia-t4",
115            ArchType::Amd64GpuG4adRadeon => "amd64-gpu-g4ad-radeon",
116            ArchType::Amd64GpuG5NvidiaA10G => "amd64-gpu-g5-nvidia-a10g",
117            ArchType::Amd64GpuInf1 => "amd64-gpu-inf1",
118            ArchType::Amd64GpuTrn1 => "amd64-gpu-trn1",
119            ArchType::Unknown(s) => s.as_ref(),
120        }
121    }
122
123    /// Returns all the `&str` values of the enum members.
124    pub fn values() -> &'static [&'static str] {
125        &[
126            "amd64",                          //
127            "arm64",                          //
128            "amd64-gpu-p4-nvidia-tesla-a100", //
129            "amd64-gpu-g3-nvidia-tesla-m60",  //
130            "amd64-gpu-g4dn-nvidia-t4",       //
131            "amd64-gpu-g4ad-radeon",          //
132            "amd64-gpu-g5-nvidia-a10g",       //
133            "amd64-gpu-inf1",                 //
134            "amd64-gpu-trn1",                 //
135        ]
136    }
137
138    pub fn is_nvidia(&self) -> bool {
139        matches!(
140            self,
141            ArchType::Amd64GpuP4NvidiaTeslaA100
142                | ArchType::Amd64GpuG3NvidiaTeslaM60
143                | ArchType::Amd64GpuG4dnNvidiaT4
144                | ArchType::Amd64GpuG5NvidiaA10G
145        )
146    }
147}
148
149impl AsRef<str> for ArchType {
150    fn as_ref(&self) -> &str {
151        self.as_str()
152    }
153}
154
155/// ref. <https://docs.aws.amazon.com/linux/al2023/ug/compare-with-al2.html>
156/// ref. <https://us-west-2.console.aws.amazon.com/systems-manager/parameters/?region=us-west-2&tab=PublicTable#public_parameter_service=canonical>
157pub fn default_image_id_ssm_parameter(arch: &str, os: &str) -> io::Result<String> {
158    let arch_type = ArchType::from_str(arch).map_err(|e| {
159        io::Error::new(
160            io::ErrorKind::InvalidInput,
161            format!("failed ArchType::from_str '{arch}' with {}", e),
162        )
163    })?;
164    let os_type = OsType::from_str(os).map_err(|e| {
165        io::Error::new(
166            io::ErrorKind::Other,
167            format!("failed OsType::from_str '{os}' with {}", e),
168        )
169    })?;
170
171    match (arch_type, os_type) {
172        (
173            ArchType::Amd64
174            | ArchType::Amd64GpuP4NvidiaTeslaA100
175            | ArchType::Amd64GpuG3NvidiaTeslaM60
176            | ArchType::Amd64GpuG4dnNvidiaT4
177            | ArchType::Amd64GpuG4adRadeon
178            | ArchType::Amd64GpuG5NvidiaA10G
179            | ArchType::Amd64GpuInf1
180            | ArchType::Amd64GpuTrn1,
181            OsType::Al2023,
182        ) => {
183            Ok("/aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-x86_64".to_string())
184        }
185        (ArchType::Arm64, OsType::Al2023) => {
186            Ok("/aws/service/ami-amazon-linux-latest/al2023-ami-kernel-default-arm64".to_string())
187        }
188
189        (
190            ArchType::Amd64
191            | ArchType::Amd64GpuP4NvidiaTeslaA100
192            | ArchType::Amd64GpuG3NvidiaTeslaM60
193            | ArchType::Amd64GpuG4dnNvidiaT4
194            | ArchType::Amd64GpuG4adRadeon
195            | ArchType::Amd64GpuG5NvidiaA10G
196            | ArchType::Amd64GpuInf1
197            | ArchType::Amd64GpuTrn1,
198            OsType::Ubuntu2004,
199        ) => Ok(
200            "/aws/service/canonical/ubuntu/server/20.04/stable/current/amd64/hvm/ebs-gp2/ami-id"
201                .to_string(),
202        ),
203        (ArchType::Arm64, OsType::Ubuntu2004) => Ok(
204            "/aws/service/canonical/ubuntu/server/20.04/stable/current/arm64/hvm/ebs-gp2/ami-id"
205                .to_string(),
206        ),
207
208        (
209            ArchType::Amd64
210            | ArchType::Amd64GpuP4NvidiaTeslaA100
211            | ArchType::Amd64GpuG3NvidiaTeslaM60
212            | ArchType::Amd64GpuG4dnNvidiaT4
213            | ArchType::Amd64GpuG4adRadeon
214            | ArchType::Amd64GpuG5NvidiaA10G
215            | ArchType::Amd64GpuInf1
216            | ArchType::Amd64GpuTrn1,
217            OsType::Ubuntu2204,
218        ) => Ok(
219            "/aws/service/canonical/ubuntu/server/22.04/stable/current/amd64/hvm/ebs-gp2/ami-id"
220                .to_string(),
221        ),
222        (ArchType::Arm64, OsType::Ubuntu2204) => Ok(
223            "/aws/service/canonical/ubuntu/server/22.04/stable/current/arm64/hvm/ebs-gp2/ami-id"
224                .to_string(),
225        ),
226
227        _ => Err(io::Error::new(
228            io::ErrorKind::InvalidInput,
229            format!("unknown arch '{arch}' or os_type '{os}'"),
230        )),
231    }
232}
233
234/// Returns default instance types.
235/// ref. <https://instances.vantage.sh/?min_memory=8&min_vcpus=4&region=us-west-2&cost_duration=monthly&selected=t4g.xlarge,c5.xlarge>
236/// ref. <https://docs.aws.amazon.com/dlami/latest/devguide/gpu.html>
237///
238/// TODO: add Graviton 3 (in preview)
239/// ref. <https://aws.amazon.com/ec2/instance-types/c7g/>
240pub fn default_instance_types(
241    region: &str,
242    arch: &str,
243    instance_size: &str,
244) -> io::Result<Vec<String>> {
245    let arch_type = ArchType::from_str(arch).map_err(|e| {
246        io::Error::new(
247            io::ErrorKind::InvalidInput,
248            format!("failed ArchType::from_str '{arch}' with {}", e),
249        )
250    })?;
251
252    // NOTE
253    // ICN/* regions do not support c6a/m6a yet
254    // ICN/* regions do not support a1 yet
255    // r6g more expensive than c5...
256    // ref. <https://aws.amazon.com/ec2/instance-types/r6g>
257    match (region, arch_type, instance_size) {
258        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, "4xlarge") => Ok(vec![
259            format!("m5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m5>
260            format!("c5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c5>
261        ]),
262        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, "8xlarge") => Ok(vec![
263            format!("m5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m5>
264            format!("c5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c5>
265        ]),
266        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, "12xlarge") => Ok(vec![
267            format!("m5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m5>
268            format!("c5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c5>
269        ]),
270        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, "16xlarge") => Ok(vec![
271            format!("m5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m5>
272            format!("c5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c5>
273        ]),
274        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, "24xlarge") => Ok(vec![
275            format!("m5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m5>
276            format!("c5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c5>
277        ]),
278        ("ap-northeast-2" | "ap-northeast-3", ArchType::Amd64, _) => Ok(vec![
279            format!("t3a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/t3>
280            format!("t3.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/t3>
281            format!("t2.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/t2>
282            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
283            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
284        ]),
285
286        ("ap-northeast-2" | "ap-northeast-3", ArchType::Arm64, "4xlarge") => Ok(vec![
287            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
288            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
289        ]),
290        ("ap-northeast-2" | "ap-northeast-3", ArchType::Arm64, "8xlarge") => Ok(vec![
291            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
292            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
293        ]),
294        ("ap-northeast-2" | "ap-northeast-3", ArchType::Arm64, "12xlarge") => Ok(vec![
295            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
296            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
297        ]),
298        ("ap-northeast-2" | "ap-northeast-3", ArchType::Arm64, "16xlarge") => Ok(vec![
299            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
300            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
301        ]),
302        ("ap-northeast-2" | "ap-northeast-3", ArchType::Arm64, _) => Ok(vec![
303            format!("t4g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/t4g>
304            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
305            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
306        ]),
307
308        (_, ArchType::Amd64, "4xlarge") => Ok(vec![
309            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
310            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
311            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
312            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
313        ]),
314        (_, ArchType::Amd64, "8xlarge") => Ok(vec![
315            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
316            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
317            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
318            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
319        ]),
320        (_, ArchType::Amd64, "12xlarge") => Ok(vec![
321            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
322            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
323            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
324            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
325        ]),
326        (_, ArchType::Amd64, "16xlarge") => Ok(vec![
327            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
328            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
329            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
330            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
331        ]),
332        (_, ArchType::Amd64, "24xlarge") => Ok(vec![
333            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
334            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
335            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
336            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
337        ]),
338        (_, ArchType::Amd64, _) => Ok(vec![
339            format!("t3a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/t3>
340            format!("t3.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/t3>
341            format!("t2.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/t2>
342            format!("c6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6a>
343            format!("m6a.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6a>
344            format!("m5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/m5>
345            format!("c5.{instance_size}"),  // ref. <https://aws.amazon.com/ec2/instance-types/c5>
346        ]),
347
348        (_, ArchType::Arm64, "4xlarge") => Ok(vec![
349            format!("a1.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/a1>
350            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
351            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
352        ]),
353        (_, ArchType::Arm64, "8xlarge") => Ok(vec![
354            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
355            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
356        ]),
357        (_, ArchType::Arm64, _) => Ok(vec![
358            format!("a1.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/a1>
359            format!("t4g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/t4g>
360            format!("c6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/c6g>
361            format!("m6g.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/m6g>
362        ]),
363
364        (_, ArchType::Amd64GpuP4NvidiaTeslaA100, "24xlarge") => Ok(vec![
365            format!("p4d.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/p4>
366        ]),
367
368        (_, ArchType::Amd64GpuG3NvidiaTeslaM60, "xlarge") => {
369            Ok(vec![
370                format!("g3s.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/g3>
371            ])
372        }
373        (_, ArchType::Amd64GpuG3NvidiaTeslaM60, "4xlarge" | "8xlarge" | "16xlarge") => {
374            Ok(vec![
375                format!("g3.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/g3>
376            ])
377        }
378
379        (
380            _,
381            ArchType::Amd64GpuG4dnNvidiaT4,
382            "xlarge" | "2xlarge" | "4xlarge" | "8xlarge" | "12xlarge" | "16xlarge",
383        ) => Ok(vec![
384            format!("g4dn.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/g4>
385        ]),
386        (
387            _,
388            ArchType::Amd64GpuG4adRadeon,
389            "xlarge" | "2xlarge" | "4xlarge" | "8xlarge" | "16xlarge",
390        ) => {
391            Ok(vec![
392                format!("g4ad.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/g4>
393            ])
394        }
395
396        (
397            _,
398            ArchType::Amd64GpuG5NvidiaA10G,
399            "xlarge" | "2xlarge" | "4xlarge" | "8xlarge" | "12xlarge" | "16xlarge" | "24xlarge"
400            | "48xlarge",
401        ) => Ok(vec![
402            format!("g5.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/g5>
403        ]),
404
405        (_, ArchType::Amd64GpuInf1, "xlarge" | "2xlarge" | "6xlarge" | "24xlarge") => Ok(vec![
406            format!("inf1.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/inf1>
407        ]),
408        (_, ArchType::Amd64GpuTrn1, "2xlarge" | "32xlarge") => Ok(vec![
409            format!("trn1.{instance_size}"), // ref. <https://aws.amazon.com/ec2/instance-types/trn1>
410        ]),
411
412        _ => Err(io::Error::new(
413            io::ErrorKind::InvalidInput,
414            format!("unknown region '{region}' and arch '{arch}'"),
415        )),
416    }
417}
418
419/// Returns a set of valid instance types.
420/// Empty if not known.
421pub fn valid_instance_types(arch_type: ArchType) -> HashSet<String> {
422    match arch_type {
423        ArchType::Amd64GpuP4NvidiaTeslaA100 => {
424            // ref. <https://aws.amazon.com/ec2/instance-types/p4>
425            let mut s = HashSet::new();
426            s.insert("p4d.24xlarge".to_string());
427            s
428        }
429        ArchType::Amd64GpuG3NvidiaTeslaM60 => {
430            // ref. <https://aws.amazon.com/ec2/instance-types/g3>
431            let mut s = HashSet::new();
432            s.insert("g3s.xlarge".to_string());
433            s.insert("g3.4xlarge".to_string());
434            s.insert("g3.8xlarge".to_string());
435            s.insert("g3.16xlarge".to_string());
436            s
437        }
438        ArchType::Amd64GpuG4dnNvidiaT4 => {
439            // ref. <https://aws.amazon.com/ec2/instance-types/g4>
440            let mut s = HashSet::new();
441            s.insert("g4dn.xlarge".to_string());
442            s.insert("g4dn.2xlarge".to_string());
443            s.insert("g4dn.4xlarge".to_string());
444            s.insert("g4dn.8xlarge".to_string());
445            s.insert("g4dn.16xlarge".to_string());
446            s.insert("g4dn.12xlarge".to_string());
447            s.insert("g4dn.metal".to_string());
448            s
449        }
450        ArchType::Amd64GpuG4adRadeon => {
451            // ref. <https://aws.amazon.com/ec2/instance-types/g4>
452            let mut s = HashSet::new();
453            s.insert("g4ad.xlarge".to_string());
454            s.insert("g4ad.2xlarge".to_string());
455            s.insert("g4ad.4xlarge".to_string());
456            s.insert("g4ad.8xlarge".to_string());
457            s.insert("g4ad.16xlarge".to_string());
458            s
459        }
460        ArchType::Amd64GpuG5NvidiaA10G => {
461            // ref. <https://aws.amazon.com/ec2/instance-types/g5>
462            let mut s = HashSet::new();
463            s.insert("g5.xlarge".to_string());
464            s.insert("g5.2xlarge".to_string());
465            s.insert("g5.4xlarge".to_string());
466            s.insert("g5.8xlarge".to_string());
467            s.insert("g5.16xlarge".to_string());
468            s.insert("g5.12xlarge".to_string());
469            s.insert("g5.24xlarge".to_string());
470            s.insert("g5.48xlarge".to_string());
471            s
472        }
473        ArchType::Amd64GpuInf1 => {
474            // ref. <https://aws.amazon.com/ec2/instance-types/inf1>
475            let mut s = HashSet::new();
476            s.insert("inf1.xlarge".to_string());
477            s.insert("inf1.2xlarge".to_string());
478            s.insert("inf1.6xlarge".to_string());
479            s.insert("inf1.24xlarge".to_string());
480            s
481        }
482        ArchType::Amd64GpuTrn1 => {
483            // ref. <https://aws.amazon.com/ec2/instance-types/trn1>
484            let mut s = HashSet::new();
485            s.insert("trn1.2xlarge".to_string());
486            s.insert("trn1.32xlarge".to_string());
487            s.insert("trn1n.32xlarge".to_string());
488            s
489        }
490        _ => HashSet::new(),
491    }
492}
493
494/// Defines the OS type.
495#[derive(
496    Deserialize,
497    Serialize,
498    std::clone::Clone,
499    std::cmp::Eq,
500    std::cmp::Ord,
501    std::cmp::PartialEq,
502    std::cmp::PartialOrd,
503    std::fmt::Debug,
504    std::hash::Hash,
505)]
506pub enum OsType {
507    #[serde(rename = "al2023")]
508    Al2023,
509    #[serde(rename = "ubuntu20.04")]
510    Ubuntu2004,
511    #[serde(rename = "ubuntu22.04")]
512    Ubuntu2204,
513
514    Unknown(String),
515}
516
517impl std::convert::From<&str> for OsType {
518    fn from(s: &str) -> Self {
519        match s {
520            "al2023" => OsType::Al2023,
521            "ubuntu20.04" => OsType::Ubuntu2004,
522            "ubuntu-20.04" => OsType::Ubuntu2004,
523            "ubuntu22.04" => OsType::Ubuntu2204,
524            "ubuntu-22.04" => OsType::Ubuntu2204,
525            other => OsType::Unknown(other.to_owned()),
526        }
527    }
528}
529
530impl std::str::FromStr for OsType {
531    type Err = std::convert::Infallible;
532
533    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
534        Ok(OsType::from(s))
535    }
536}
537
538impl OsType {
539    /// Returns the `&str` value of the enum member.
540    pub fn as_str(&self) -> &str {
541        match self {
542            OsType::Al2023 => "al2023",
543            OsType::Ubuntu2004 => "ubuntu20.04",
544            OsType::Ubuntu2204 => "ubuntu22.04",
545            OsType::Unknown(s) => s.as_ref(),
546        }
547    }
548
549    /// Returns all the `&str` values of the enum members.
550    pub fn values() -> &'static [&'static str] {
551        &[
552            "al2023",      //
553            "ubuntu20.04", //
554            "ubuntu22.04", //
555        ]
556    }
557}
558
559impl AsRef<str> for OsType {
560    fn as_ref(&self) -> &str {
561        self.as_str()
562    }
563}
564
565/// ref. <https://docs.aws.amazon.com/linux/al2023/ug/compare-with-al2.html>
566pub fn default_user_name(os_type: &str) -> io::Result<String> {
567    match OsType::from_str(os_type).map_err(|e| {
568        io::Error::new(
569            io::ErrorKind::InvalidInput,
570            format!("failed OsType::from_str '{os_type}' with {}", e),
571        )
572    })? {
573        OsType::Al2023 => Ok("ec2-user".to_string()),
574        OsType::Ubuntu2004 | OsType::Ubuntu2204 => Ok("ubuntu".to_string()),
575        OsType::Unknown(v) => Err(io::Error::new(
576            io::ErrorKind::InvalidInput,
577            format!("unknown os_type '{v}'"),
578        )),
579    }
580}
581
582/// Implements AWS EC2 manager.
583#[derive(Debug, Clone)]
584pub struct Manager {
585    pub region: String,
586    pub cli: Client,
587}
588
589impl Manager {
590    pub fn new(shared_config: &AwsSdkConfig) -> Self {
591        Self {
592            region: shared_config.region().unwrap().to_string(),
593            cli: Client::new(shared_config),
594        }
595    }
596
597    /// Imports a public key to EC2 key.
598    pub async fn import_key(&self, key_name: &str, pubkey_path: &str) -> Result<String> {
599        let path = Path::new(pubkey_path);
600        if !path.exists() {
601            return Err(Error::Other {
602                message: format!("public key path {pubkey_path} does not exist"),
603                retryable: false,
604            });
605        }
606        let pubkey_raw = fs::read(pubkey_path).map_err(|e| Error::Other {
607            message: format!("failed to read {} {:?}", pubkey_path, e),
608            retryable: false,
609        })?;
610        let pubkey_material = aws_smithy_types::Blob::new(pubkey_raw);
611
612        log::info!(
613            "importing a public key '{pubkey_path}' with key name '{key_name}' in region '{}'",
614            self.region
615        );
616
617        let out = self
618            .cli
619            .import_key_pair()
620            .key_name(key_name)
621            .public_key_material(pubkey_material)
622            .send()
623            .await
624            .map_err(|e| Error::API {
625                message: format!("failed import_key_pair {} {:?}", pubkey_path, e),
626                retryable: errors::is_sdk_err_retryable(&e),
627            })?;
628
629        let key_pair_id = out.key_pair_id().unwrap().clone();
630        log::info!("imported key pair id '{key_pair_id}' -- describing");
631
632        let out = self
633            .cli
634            .describe_key_pairs()
635            .key_pair_ids(key_pair_id)
636            .send()
637            .await
638            .map_err(|e| Error::API {
639                message: format!("failed describe_key_pairs {} {:?}", pubkey_path, e),
640                retryable: errors::is_sdk_err_retryable(&e),
641            })?;
642        if let Some(kps) = out.key_pairs() {
643            if kps.len() != 1 {
644                return Err(Error::API {
645                    message: format!("unexpected {} key pairs from describe_key_pairs", kps.len()),
646                    retryable: false,
647                });
648            }
649
650            let described_key_name = kps[0].key_name().clone().unwrap().to_string();
651            let described_key_pair_id = kps[0].key_pair_id().clone().unwrap().to_string();
652            log::info!("described imported key name {described_key_name} and key pair id {described_key_pair_id}");
653
654            if described_key_name != key_name {
655                return Err(Error::API {
656                    message: format!(
657                        "unexpected described key name {} != {}",
658                        described_key_name, key_name
659                    ),
660                    retryable: false,
661                });
662            }
663            if described_key_pair_id != key_pair_id {
664                return Err(Error::API {
665                    message: format!(
666                        "unexpected described key pair id {} != {}",
667                        described_key_pair_id, key_pair_id
668                    ),
669                    retryable: false,
670                });
671            }
672        } else {
673            return Err(Error::API {
674                message: format!("unexpected empty key pair from describe_key_pairs"),
675                retryable: false,
676            });
677        }
678
679        log::info!(
680            "successfully imported the key {key_name} with the public key file {pubkey_path}"
681        );
682        Ok(key_pair_id.to_string())
683    }
684
685    /// Creates an AWS EC2 key-pair and saves the private key to disk.
686    /// It overwrites "key_path" file with the newly created key.
687    pub async fn create_key_pair(&self, key_name: &str, key_path: &str) -> Result<()> {
688        let path = Path::new(key_path);
689        if path.exists() {
690            return Err(Error::Other {
691                message: format!(
692                    "private key path {} already exists, can't overwrite with a new key",
693                    key_path
694                ),
695                retryable: false,
696            });
697        }
698
699        // "KeyType::Rsa" is the default
700        // "KeyFormat::Pem" is the default
701        log::info!(
702            "creating EC2 key-pair '{}' '{key_name}' in region '{}'",
703            KeyType::Rsa.as_str(),
704            self.region
705        );
706        let ret = self
707            .cli
708            .create_key_pair()
709            .key_name(key_name)
710            .key_type(KeyType::Rsa)
711            .key_format(KeyFormat::Pem)
712            .send()
713            .await;
714        let resp = match ret {
715            Ok(v) => v,
716            Err(e) => {
717                return Err(Error::API {
718                    message: format!("failed create_key_pair {:?}", e),
719                    retryable: errors::is_sdk_err_retryable(&e),
720                });
721            }
722        };
723
724        log::info!(
725            "persisting the created EC2 key-pair '{}' in '{}'",
726            key_name,
727            key_path
728        );
729        let key_material = resp.key_material().unwrap();
730
731        let mut f = match File::create(&path) {
732            Ok(f) => f,
733            Err(e) => {
734                return Err(Error::Other {
735                    message: format!("failed to create file {:?}", e),
736                    retryable: false,
737                });
738            }
739        };
740        match f.write_all(key_material.as_bytes()) {
741            Ok(_) => {}
742            Err(e) => {
743                return Err(Error::Other {
744                    message: format!("failed to write file {:?}", e),
745                    retryable: false,
746                });
747            }
748        }
749
750        Ok(())
751    }
752
753    /// Deletes the AWS EC2 key-pair.
754    pub async fn delete_key_pair(&self, key_name: &str) -> Result<()> {
755        log::info!(
756            "deleting EC2 key-pair '{key_name}' in region '{}'",
757            self.region
758        );
759        let ret = self.cli.delete_key_pair().key_name(key_name).send().await;
760        match ret {
761            Ok(_) => {}
762            Err(e) => {
763                if !is_err_does_not_exist_delete_key_pair(&e) {
764                    return Err(Error::API {
765                        message: format!("failed delete_key_pair {:?}", e),
766                        retryable: errors::is_sdk_err_retryable(&e),
767                    });
768                }
769                log::warn!("key already deleted ({})", e);
770            }
771        };
772
773        Ok(())
774    }
775
776    /// Describes an AWS EC2 VPC.
777    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeVpcs.html>
778    pub async fn describe_vpc(&self, vpc_id: &str) -> Result<Vpc> {
779        log::info!("describing VPC '{vpc_id}' in region '{}'", self.region);
780        let ret = self.cli.describe_vpcs().vpc_ids(vpc_id).send().await;
781        let vpcs = match ret {
782            Ok(out) => {
783                if let Some(vpcs) = out.vpcs() {
784                    vpcs.to_vec()
785                } else {
786                    return Err(Error::API {
787                        message: "no vpc found".to_string(),
788                        retryable: false,
789                    });
790                }
791            }
792            Err(e) => {
793                return Err(Error::API {
794                    message: format!("failed describe_vpcs {:?}", e),
795                    retryable: errors::is_sdk_err_retryable(&e),
796                });
797            }
798        };
799        if vpcs.len() != 1 {
800            return Err(Error::API {
801                message: format!("expected 1 VPC, got {} VPCs", vpcs.len()),
802                retryable: false,
803            });
804        }
805
806        Ok(vpcs[0].to_owned())
807    }
808
809    /// Describes security groups by VPC Id.
810    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeSecurityGroups.html>
811    pub async fn describe_security_groups_by_vpc(
812        &self,
813        vpc_id: &str,
814    ) -> Result<Vec<SecurityGroup>> {
815        log::info!(
816            "describing security groups for '{vpc_id}' in region '{}'",
817            self.region
818        );
819        let ret = self
820            .cli
821            .describe_security_groups()
822            .filters(
823                Filter::builder()
824                    .set_name(Some(String::from("vpc-id")))
825                    .set_values(Some(vec![vpc_id.to_string()]))
826                    .build(),
827            )
828            .send()
829            .await;
830        match ret {
831            Ok(out) => {
832                if let Some(sgs) = out.security_groups() {
833                    Ok(sgs.to_vec())
834                } else {
835                    return Err(Error::API {
836                        message: "no security group found".to_string(),
837                        retryable: false,
838                    });
839                }
840            }
841            Err(e) => {
842                return Err(Error::API {
843                    message: format!("failed describe_security_groups {:?}", e),
844                    retryable: errors::is_sdk_err_retryable(&e),
845                });
846            }
847        }
848    }
849
850    /// Describes subnets by VPC Id.
851    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeSubnets.html>
852    pub async fn describe_subnets_by_vpc(&self, vpc_id: &str) -> Result<Vec<Subnet>> {
853        log::info!(
854            "describing subnets for '{vpc_id}' in region '{}'",
855            self.region
856        );
857        let ret = self
858            .cli
859            .describe_subnets()
860            .filters(
861                Filter::builder()
862                    .set_name(Some(String::from("vpc-id")))
863                    .set_values(Some(vec![vpc_id.to_string()]))
864                    .build(),
865            )
866            .send()
867            .await;
868        match ret {
869            Ok(out) => {
870                if let Some(ss) = out.subnets() {
871                    Ok(ss.to_vec())
872                } else {
873                    return Err(Error::API {
874                        message: "no subnet found".to_string(),
875                        retryable: false,
876                    });
877                }
878            }
879            Err(e) => {
880                return Err(Error::API {
881                    message: format!("failed describe_subnets {:?}", e),
882                    retryable: errors::is_sdk_err_retryable(&e),
883                });
884            }
885        }
886    }
887
888    /// Describes the EBS volumes by filters.
889    /// ref. https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeVolumes.html
890    pub async fn describe_volumes(&self, filters: Option<Vec<Filter>>) -> Result<Vec<Volume>> {
891        log::info!("describing volumes in region '{}'", self.region);
892        let resp = match self
893            .cli
894            .describe_volumes()
895            .set_filters(filters)
896            .send()
897            .await
898        {
899            Ok(r) => r,
900            Err(e) => {
901                return Err(Error::API {
902                    message: format!("failed describe_volumes {:?}", e),
903                    retryable: errors::is_sdk_err_retryable(&e),
904                });
905            }
906        };
907
908        let volumes = if let Some(vols) = resp.volumes {
909            vols
910        } else {
911            Vec::new()
912        };
913
914        log::info!(
915            "described {} volumes in region '{}'",
916            volumes.len(),
917            self.region
918        );
919        Ok(volumes)
920    }
921
922    /// Polls the EBS volume by its state.
923    pub async fn poll_volume_state(
924        &self,
925        ebs_volume_id: String,
926        desired_state: VolumeState,
927        timeout: Duration,
928        interval: Duration,
929    ) -> Result<Option<Volume>> {
930        let start = Instant::now();
931        let mut cnt: u128 = 0;
932        loop {
933            let elapsed = start.elapsed();
934            if elapsed.gt(&timeout) {
935                break;
936            }
937
938            let itv = {
939                if cnt == 0 {
940                    // first poll with no wait
941                    Duration::from_secs(1)
942                } else {
943                    interval
944                }
945            };
946            sleep(itv).await;
947
948            let volumes = self
949                .describe_volumes(Some(vec![Filter::builder()
950                    .set_name(Some(String::from("volume-id")))
951                    .set_values(Some(vec![ebs_volume_id.clone()]))
952                    .build()]))
953                .await?;
954            if volumes.is_empty() {
955                if desired_state.eq(&VolumeState::Deleted) {
956                    log::info!("volume already deleted");
957                    return Ok(None);
958                }
959
960                log::warn!("no volume found");
961                continue;
962            }
963            if volumes.len() != 1 {
964                log::warn!("unexpected {} volumes found", volumes.len());
965                continue;
966            }
967            let volume = volumes[0].clone();
968
969            let current_state = {
970                if let Some(v) = volume.state() {
971                    v.clone()
972                } else {
973                    VolumeState::from("not found")
974                }
975            };
976            log::info!(
977                "poll (current volume state {:?}, elapsed {:?})",
978                current_state,
979                elapsed
980            );
981
982            if current_state.eq(&desired_state) {
983                return Ok(Some(volume));
984            }
985
986            cnt += 1;
987        }
988
989        Err(Error::Other {
990            message: format!(
991                "failed to poll volume state for '{}' in time",
992                ebs_volume_id
993            ),
994            retryable: true,
995        })
996    }
997
998    /// Describes the attached volume by the volume Id and EBS device name.
999    /// The "local_ec2_instance_id" is only set to bypass extra EC2 metadata
1000    /// service API calls.
1001    /// The region used for API call is inherited from the EC2 client SDK.
1002    ///
1003    /// e.g.,
1004    ///
1005    /// aws ec2 describe-volumes \
1006    /// --region ${AWS::Region} \
1007    /// --filters \
1008    ///   Name=attachment.instance-id,Values=$INSTANCE_ID \
1009    ///   Name=attachment.device,Values=/dev/xvdb \
1010    /// --query Volumes[].Attachments[].State \
1011    /// --output text
1012    ///
1013    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeVolumes.html>
1014    /// ref. <https://github.com/ava-labs/avalanche-ops/blob/fcbac87a219a8d3d6d3c38a1663fe1dafe78e04e/bin/avalancheup-aws/cfn-templates/asg_amd64_ubuntu.yaml#L397-L409>
1015    pub async fn describe_local_volumes(
1016        &self,
1017        ebs_volume_id: Option<String>,
1018        ebs_device_name: String,
1019        local_ec2_instance_id: Option<String>,
1020    ) -> Result<Vec<Volume>> {
1021        let mut filters: Vec<Filter> = vec![];
1022
1023        if let Some(v) = ebs_volume_id {
1024            log::info!(
1025                "filtering volumes via volume Id '{}' in region '{}'",
1026                v,
1027                self.region
1028            );
1029            filters.push(
1030                Filter::builder()
1031                    .set_name(Some(String::from("volume-id")))
1032                    .set_values(Some(vec![v]))
1033                    .build(),
1034            );
1035        }
1036
1037        let device = if ebs_device_name.starts_with("/dev/") {
1038            ebs_device_name
1039        } else {
1040            format!("/dev/{}", ebs_device_name.clone()).to_string()
1041        };
1042        log::info!(
1043            "filtering volumes via EBS device name '{}' in region '{}'",
1044            device,
1045            self.region
1046        );
1047        filters.push(
1048            Filter::builder()
1049                .set_name(Some(String::from("attachment.device")))
1050                .set_values(Some(vec![device]))
1051                .build(),
1052        );
1053
1054        let ec2_instance_id = if let Some(v) = local_ec2_instance_id {
1055            v
1056        } else {
1057            metadata::fetch_instance_id().await?
1058        };
1059        log::info!("filtering volumes via instance Id {}", ec2_instance_id);
1060        filters.push(
1061            Filter::builder()
1062                .set_name(Some(String::from("attachment.instance-id")))
1063                .set_values(Some(vec![ec2_instance_id]))
1064                .build(),
1065        );
1066
1067        self.describe_volumes(Some(filters)).await
1068    }
1069
1070    /// Polls the EBS volume attachment state.
1071    /// For instance, the "device_name" can be either "/dev/xvdb" or "xvdb" (for the secondary volume).
1072    pub async fn poll_local_volume_by_attachment_state(
1073        &self,
1074        ebs_volume_id: Option<String>,
1075        ebs_device_name: String,
1076        desired_attachment_state: VolumeAttachmentState,
1077        timeout: Duration,
1078        interval: Duration,
1079    ) -> Result<Volume> {
1080        let local_ec2_instance_id = metadata::fetch_instance_id().await?;
1081        let start = Instant::now();
1082        let mut cnt: u128 = 0;
1083        loop {
1084            let elapsed = start.elapsed();
1085            if elapsed.gt(&timeout) {
1086                break;
1087            }
1088
1089            let itv = {
1090                if cnt == 0 {
1091                    // first poll with no wait
1092                    Duration::from_secs(1)
1093                } else {
1094                    interval
1095                }
1096            };
1097            sleep(itv).await;
1098
1099            let volumes = self
1100                .describe_local_volumes(
1101                    ebs_volume_id.clone(),
1102                    ebs_device_name.clone(),
1103                    Some(local_ec2_instance_id.clone()),
1104                )
1105                .await?;
1106            if volumes.is_empty() {
1107                log::warn!("no volume found");
1108                continue;
1109            }
1110            if volumes.len() != 1 {
1111                log::warn!("unexpected {} volumes found", volumes.len());
1112                continue;
1113            }
1114            let volume = volumes[0].clone();
1115            if volume.attachments().is_none() {
1116                log::warn!("no attachment found");
1117                continue;
1118            }
1119            let attachments = volume.attachments().unwrap();
1120            if attachments.is_empty() {
1121                log::warn!("no attachment found");
1122                continue;
1123            }
1124            if attachments.len() != 1 {
1125                log::warn!("unexpected attachment found {}", attachments.len());
1126                continue;
1127            }
1128            let current_attachment_state = attachments[0].state().unwrap();
1129            log::info!(
1130                "poll (current volume attachment state {:?}, elapsed {:?})",
1131                current_attachment_state,
1132                elapsed
1133            );
1134
1135            if current_attachment_state.eq(&desired_attachment_state) {
1136                return Ok(volume);
1137            }
1138
1139            cnt += 1;
1140        }
1141
1142        Err(Error::Other {
1143            message: format!(
1144                "failed to poll volume attachment state for '{}' in time",
1145                local_ec2_instance_id
1146            ),
1147            retryable: true,
1148        })
1149    }
1150
1151    /// Fetches all tags for the specified instance.
1152    ///
1153    /// "If a single piece of data must be accessible from more than one task
1154    /// concurrently, then it must be shared using synchronization primitives such as Arc."
1155    /// ref. https://tokio.rs/tokio/tutorial/spawning
1156    pub async fn fetch_tags(&self, instance_id: &str) -> Result<Vec<Tag>> {
1157        log::info!("fetching tags for '{}'", instance_id);
1158        let ret = self
1159            .cli
1160            .describe_instances()
1161            .instance_ids(instance_id)
1162            .send()
1163            .await;
1164        let resp = match ret {
1165            Ok(r) => r,
1166            Err(e) => {
1167                return Err(Error::API {
1168                    message: format!("failed describe_instances {:?}", e),
1169                    retryable: errors::is_sdk_err_retryable(&e),
1170                });
1171            }
1172        };
1173
1174        let reservations = match resp.reservations {
1175            Some(rvs) => rvs,
1176            None => {
1177                return Err(Error::API {
1178                    message: String::from("empty reservation from describe_instances response"),
1179                    retryable: false,
1180                });
1181            }
1182        };
1183        if reservations.len() != 1 {
1184            return Err(Error::API {
1185                message: format!(
1186                    "expected only 1 reservation from describe_instances response but got {}",
1187                    reservations.len()
1188                ),
1189                retryable: false,
1190            });
1191        }
1192
1193        let rvs = reservations.get(0).unwrap();
1194        let instances = rvs.instances.to_owned().unwrap();
1195        if instances.len() != 1 {
1196            return Err(Error::API {
1197                message: format!(
1198                    "expected only 1 instance from describe_instances response but got {}",
1199                    instances.len()
1200                ),
1201                retryable: false,
1202            });
1203        }
1204
1205        let instance = instances.get(0).unwrap();
1206        let tags = match instance.tags.to_owned() {
1207            Some(ss) => ss,
1208            None => {
1209                return Err(Error::API {
1210                    message: String::from("empty tags from describe_instances response"),
1211                    retryable: false,
1212                });
1213            }
1214        };
1215        log::info!("fetched {} tags for '{}'", tags.len(), instance_id);
1216
1217        Ok(tags)
1218    }
1219
1220    /// Lists instances by the Auto Scaling Groups name.
1221    pub async fn list_asg(&self, asg_name: &str) -> Result<Vec<Droplet>> {
1222        log::info!("listing asg '{asg_name}' for the region '{}'", self.region);
1223
1224        let filter = Filter::builder()
1225            .set_name(Some(String::from("tag:aws:autoscaling:groupName")))
1226            .set_values(Some(vec![String::from(asg_name)]))
1227            .build();
1228        let resp = match self
1229            .cli
1230            .describe_instances()
1231            .set_filters(Some(vec![filter]))
1232            .send()
1233            .await
1234        {
1235            Ok(r) => r,
1236            Err(e) => {
1237                return Err(Error::API {
1238                    message: format!("failed describe_instances {:?}", e),
1239                    retryable: errors::is_sdk_err_retryable(&e),
1240                });
1241            }
1242        };
1243
1244        let reservations = match resp.reservations {
1245            Some(rvs) => rvs,
1246            None => {
1247                log::warn!("empty reservation from describe_instances response");
1248                return Ok(vec![]);
1249            }
1250        };
1251
1252        let mut droplets: Vec<Droplet> = Vec::new();
1253        for rsv in reservations.iter() {
1254            let instances = rsv.instances().unwrap();
1255            for instance in instances {
1256                let instance_id = instance.instance_id().unwrap();
1257                log::info!("instance {}", instance_id);
1258                droplets.push(Droplet::new(instance));
1259            }
1260        }
1261
1262        Ok(droplets)
1263    }
1264
1265    /// Allocates an EIP and returns the allocation Id and the public Ip.
1266    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_AllocateAddress.html>
1267    pub async fn allocate_eip(&self, tags: HashMap<String, String>) -> Result<Eip> {
1268        log::info!("allocating elastic IP with tags {:?}", tags);
1269
1270        let mut eip_tags = TagSpecification::builder().resource_type(ResourceType::ElasticIp);
1271        for (k, v) in tags.iter() {
1272            eip_tags = eip_tags.tags(Tag::builder().key(k).value(v).build());
1273        }
1274
1275        let resp = match self
1276            .cli
1277            .allocate_address()
1278            .tag_specifications(eip_tags.build())
1279            .send()
1280            .await
1281        {
1282            Ok(r) => r,
1283            Err(e) => {
1284                return Err(Error::API {
1285                    message: format!("failed allocate_address {:?}", e),
1286                    retryable: errors::is_sdk_err_retryable(&e),
1287                });
1288            }
1289        };
1290
1291        let allocation_id = resp
1292            .allocation_id
1293            .to_owned()
1294            .unwrap_or_else(|| String::from(""));
1295        let public_ip = resp
1296            .public_ip
1297            .to_owned()
1298            .unwrap_or_else(|| String::from(""));
1299        log::info!("successfully allocated elastic IP {public_ip} with {allocation_id}");
1300
1301        Ok(Eip {
1302            allocation_id,
1303            public_ip,
1304        })
1305    }
1306
1307    /// Associates the elastic Ip with an EC2 instance.
1308    /// ref. https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_AssociateAddress.html
1309    pub async fn associate_eip(&self, allocation_id: &str, instance_id: &str) -> Result<String> {
1310        log::info!("associating elastic IP {allocation_id} with EC2 instance {instance_id}");
1311        let resp = match self
1312            .cli
1313            .associate_address()
1314            .allocation_id(allocation_id)
1315            .instance_id(instance_id)
1316            .send()
1317            .await
1318        {
1319            Ok(r) => r,
1320            Err(e) => {
1321                return Err(Error::API {
1322                    message: format!("failed associate_address {:?}", e),
1323                    retryable: errors::is_sdk_err_retryable(&e),
1324                });
1325            }
1326        };
1327
1328        let association_id = resp
1329            .association_id
1330            .to_owned()
1331            .unwrap_or_else(|| String::from(""));
1332        log::info!("successfully associated elastic IP {allocation_id} with association Id {association_id}");
1333
1334        Ok(association_id)
1335    }
1336
1337    /// Describes the elastic IP addresses with the instance Id.
1338    /// ref. https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeAddresses.html
1339    pub async fn describe_eips_by_instance_id(&self, instance_id: &str) -> Result<Vec<Address>> {
1340        log::info!("describing elastic IP addresses for EC2 instance {instance_id}");
1341
1342        let resp = match self
1343            .cli
1344            .describe_addresses()
1345            .set_filters(Some(vec![Filter::builder()
1346                .set_name(Some(String::from("instance-id")))
1347                .set_values(Some(vec![instance_id.to_string()]))
1348                .build()]))
1349            .send()
1350            .await
1351        {
1352            Ok(r) => r,
1353            Err(e) => {
1354                return Err(Error::API {
1355                    message: format!("failed describe_addresses {:?}", e),
1356                    retryable: errors::is_sdk_err_retryable(&e),
1357                });
1358            }
1359        };
1360        let addrs = if let Some(addrs) = resp.addresses() {
1361            addrs.to_vec()
1362        } else {
1363            Vec::new()
1364        };
1365
1366        log::info!("successfully described addresses: {:?}", addrs);
1367        Ok(addrs)
1368    }
1369
1370    /// Describes the elastic IP addresses with the tags.
1371    /// ref. <https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeAddresses.html>
1372    pub async fn describe_eips_by_tags(
1373        &self,
1374        tags: HashMap<String, String>,
1375    ) -> Result<Vec<Address>> {
1376        log::info!("describing elastic IP addresses with tags {:?}", tags);
1377
1378        let mut filters = Vec::new();
1379        for (k, v) in tags.iter() {
1380            filters.push(
1381                Filter::builder()
1382                    .set_name(Some(format!("tag:{}", k)))
1383                    .set_values(Some(vec![v.clone()]))
1384                    .build(),
1385            );
1386        }
1387        let resp = match self
1388            .cli
1389            .describe_addresses()
1390            .set_filters(Some(filters))
1391            .send()
1392            .await
1393        {
1394            Ok(r) => r,
1395            Err(e) => {
1396                return Err(Error::API {
1397                    message: format!("failed describe_addresses {:?}", e),
1398                    retryable: errors::is_sdk_err_retryable(&e),
1399                });
1400            }
1401        };
1402        let addrs = if let Some(addrs) = resp.addresses() {
1403            addrs.to_vec()
1404        } else {
1405            Vec::new()
1406        };
1407
1408        log::info!("successfully described addresses: {:?}", addrs);
1409        Ok(addrs)
1410    }
1411
1412    /// Polls the elastic Ip for its describe address state,
1413    /// until the elastic Ip becomes attached to the instance.
1414    /// ref. https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeAddresses.html
1415    pub async fn poll_eip_by_describe_addresses(
1416        &self,
1417        association_id: &str,
1418        instance_id: &str,
1419        timeout: Duration,
1420        interval: Duration,
1421    ) -> Result<Vec<Address>> {
1422        log::info!(
1423            "describing elastic IP association Id {association_id} for EC2 instance {instance_id}"
1424        );
1425
1426        let filters = vec![
1427            Filter::builder()
1428                .set_name(Some(String::from("association-id")))
1429                .set_values(Some(vec![association_id.to_string()]))
1430                .build(),
1431            Filter::builder()
1432                .set_name(Some(String::from("instance-id")))
1433                .set_values(Some(vec![instance_id.to_string()]))
1434                .build(),
1435        ];
1436
1437        let start = Instant::now();
1438        let mut cnt: u128 = 0;
1439        loop {
1440            let elapsed = start.elapsed();
1441            if elapsed.gt(&timeout) {
1442                break;
1443            }
1444
1445            let itv = {
1446                if cnt == 0 {
1447                    // first poll with no wait
1448                    Duration::from_secs(1)
1449                } else {
1450                    interval
1451                }
1452            };
1453            sleep(itv).await;
1454
1455            let resp = match self
1456                .cli
1457                .describe_addresses()
1458                .set_filters(Some(filters.clone()))
1459                .send()
1460                .await
1461            {
1462                Ok(r) => r,
1463                Err(e) => {
1464                    return Err(Error::API {
1465                        message: format!("failed describe_addresses {:?}", e),
1466                        retryable: errors::is_sdk_err_retryable(&e),
1467                    });
1468                }
1469            };
1470            let addrs = if let Some(addrs) = resp.addresses() {
1471                addrs.to_vec()
1472            } else {
1473                Vec::new()
1474            };
1475            log::info!("successfully described addresses: {:?}", addrs);
1476            if !addrs.is_empty() {
1477                return Ok(addrs);
1478            }
1479
1480            cnt += 1;
1481        }
1482
1483        Err(Error::Other {
1484            message: format!(
1485                "failed to poll describe_address elastic IP association Id {association_id} for EC2 instance {instance_id} in time",
1486            ),
1487            retryable: true,
1488        })
1489    }
1490
1491    /// Creates an image and returns the AMI ID.
1492    pub async fn create_image(
1493        &self,
1494        instance_id: &str,
1495        image_name: &str,
1496        tags: HashMap<String, String>,
1497    ) -> Result<String> {
1498        log::info!("creating an image '{image_name}' in instance '{instance_id}'");
1499
1500        let mut ami_tags = TagSpecification::builder().resource_type(ResourceType::Image);
1501        for (k, v) in tags.iter() {
1502            ami_tags = ami_tags.tags(Tag::builder().key(k).value(v).build());
1503        }
1504
1505        let ami = self
1506            .cli
1507            .create_image()
1508            .instance_id(instance_id)
1509            .name(image_name)
1510            .tag_specifications(ami_tags.build())
1511            .send()
1512            .await
1513            .map_err(|e| Error::API {
1514                message: format!("failed create_image {:?}", e),
1515                retryable: errors::is_sdk_err_retryable(&e),
1516            })?;
1517
1518        let ami_id = ami.image_id().clone().unwrap().to_string();
1519        log::info!("created AMI '{ami_id}' from the instance '{instance_id}'");
1520
1521        Ok(ami_id)
1522    }
1523
1524    /// Polls the image until the state is "Available".
1525    pub async fn poll_image_until_available(
1526        &self,
1527        image_id: &str,
1528        timeout: Duration,
1529        interval: Duration,
1530    ) -> Result<Image> {
1531        log::info!("describing AMI {image_id} until available");
1532
1533        let start = Instant::now();
1534        let mut cnt: u128 = 0;
1535        loop {
1536            let elapsed = start.elapsed();
1537            if elapsed.gt(&timeout) {
1538                break;
1539            }
1540
1541            let itv = {
1542                if cnt == 0 {
1543                    // first poll with no wait
1544                    Duration::from_secs(1)
1545                } else {
1546                    interval
1547                }
1548            };
1549            sleep(itv).await;
1550
1551            let resp = match self.cli.describe_images().image_ids(image_id).send().await {
1552                Ok(r) => r,
1553                Err(e) => {
1554                    return Err(Error::API {
1555                        message: format!("failed describe_images {:?}", e),
1556                        retryable: errors::is_sdk_err_retryable(&e),
1557                    });
1558                }
1559            };
1560            let images = if let Some(images) = resp.images() {
1561                images.to_vec()
1562            } else {
1563                Vec::new()
1564            };
1565            if images.len() != 1 {
1566                return Err(Error::Other {
1567                    message: format!(
1568                        "unexpected output from describe_images, expected 1 image but got {}",
1569                        images.len()
1570                    ),
1571                    retryable: false,
1572                });
1573            }
1574            let state = images[0].state().clone().unwrap();
1575            if state.eq(&ImageState::Available) {
1576                return Ok(images[0].clone());
1577            }
1578
1579            log::info!(
1580                "image {image_id} is still {} (elapsed {:?})",
1581                state.as_str(),
1582                elapsed
1583            );
1584
1585            cnt += 1;
1586        }
1587
1588        Err(Error::Other {
1589            message: format!("failed to poll image state {image_id} in time",),
1590            retryable: true,
1591        })
1592    }
1593}
1594
1595/// Represents the underlying EC2 instance.
1596#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
1597#[serde(rename_all = "snake_case")]
1598pub struct Droplet {
1599    pub instance_id: String,
1600    /// Represents the data format in RFC3339.
1601    /// ref. https://serde.rs/custom-date-format.html
1602    #[serde(with = "rfc_manager::serde_format::rfc_3339")]
1603    pub launched_at_utc: DateTime<Utc>,
1604    pub instance_state_code: i32,
1605    pub instance_state_name: String,
1606    pub availability_zone: String,
1607    pub public_hostname: String,
1608    pub public_ipv4: String,
1609
1610    pub block_device_mappings: Vec<BlockDeviceMapping>,
1611}
1612
1613#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
1614#[serde(rename_all = "snake_case")]
1615pub struct BlockDeviceMapping {
1616    pub device_name: String,
1617    pub volume_id: String,
1618    pub attachment_status: String,
1619}
1620
1621impl Droplet {
1622    pub fn new(inst: &Instance) -> Self {
1623        let instance_id = match inst.instance_id.to_owned() {
1624            Some(v) => v,
1625            None => String::new(),
1626        };
1627        let launch_time = inst.launch_time().unwrap();
1628        let native_dt = NaiveDateTime::from_timestamp_opt(launch_time.secs(), 0).unwrap();
1629        let launched_at_utc = DateTime::<Utc>::from_utc(native_dt, Utc);
1630
1631        let instance_state = match inst.state.to_owned() {
1632            Some(v) => v,
1633            None => InstanceState::builder().build(),
1634        };
1635        let instance_state_code = instance_state.code.unwrap_or(0);
1636        let instance_state_name = instance_state
1637            .name
1638            .unwrap_or_else(|| InstanceStateName::from("unknown"));
1639        let instance_state_name = instance_state_name.as_str().to_string();
1640
1641        let availability_zone = match inst.placement.to_owned() {
1642            Some(v) => match v.availability_zone {
1643                Some(v2) => v2,
1644                None => String::new(),
1645            },
1646            None => String::new(),
1647        };
1648
1649        let public_hostname = inst
1650            .public_dns_name
1651            .to_owned()
1652            .unwrap_or_else(|| String::from(""));
1653        let public_ipv4 = inst
1654            .public_ip_address
1655            .to_owned()
1656            .unwrap_or_else(|| String::from(""));
1657
1658        let mut block_device_mappings = Vec::new();
1659        if let Some(mappings) = inst.block_device_mappings() {
1660            for block_device_mapping in mappings.iter() {
1661                let device_name = block_device_mapping
1662                    .device_name
1663                    .to_owned()
1664                    .unwrap_or_else(|| String::from(""));
1665
1666                let (volume_id, attachment_status) = if let Some(ebs) = block_device_mapping.ebs() {
1667                    let volume_id = ebs.volume_id.to_owned().unwrap_or_else(|| String::from(""));
1668                    let attachment_status = ebs
1669                        .status
1670                        .to_owned()
1671                        .unwrap_or_else(|| AttachmentStatus::from(""));
1672                    (volume_id, attachment_status.as_str().to_string())
1673                } else {
1674                    (String::new(), String::new())
1675                };
1676
1677                block_device_mappings.push(BlockDeviceMapping {
1678                    device_name,
1679                    volume_id,
1680                    attachment_status,
1681                });
1682            }
1683        }
1684
1685        Self {
1686            instance_id,
1687            launched_at_utc,
1688            instance_state_code,
1689            instance_state_name,
1690            availability_zone,
1691            public_hostname,
1692            public_ipv4,
1693            block_device_mappings,
1694        }
1695    }
1696}
1697
1698/// EC2 does not return any error for non-existing key deletes, just in case...
1699#[inline]
1700fn is_err_does_not_exist_delete_key_pair(
1701    e: &SdkError<DeleteKeyPairError, aws_smithy_runtime_api::client::orchestrator::HttpResponse>,
1702) -> bool {
1703    match e {
1704        SdkError::ServiceError(err) => {
1705            let msg = format!("{:?}", err);
1706            msg.contains("does not exist")
1707        }
1708        _ => false,
1709    }
1710}
1711
1712/// Represents Elastic IP spec for management.
1713#[derive(Debug, Serialize, Deserialize, Eq, PartialEq, Clone)]
1714pub struct Eip {
1715    pub allocation_id: String,
1716    pub public_ip: String,
1717}
1718
1719impl Eip {
1720    /// Saves to disk overwriting the file, if any.
1721    pub fn sync(&self, file_path: &str) -> io::Result<()> {
1722        log::info!("syncing Eip spec to '{}'", file_path);
1723        let path = Path::new(file_path);
1724        let parent_dir = path.parent().expect("unexpected None parent");
1725        fs::create_dir_all(parent_dir)?;
1726
1727        let d = serde_yaml::to_string(self).map_err(|e| {
1728            io::Error::new(
1729                io::ErrorKind::Other,
1730                format!("failed to serialize Eip spec info to YAML {}", e),
1731            )
1732        })?;
1733
1734        let mut f = File::create(file_path)?;
1735        f.write_all(d.as_bytes())?;
1736
1737        Ok(())
1738    }
1739
1740    /// Loads the Eip spec from disk.
1741    pub fn load(file_path: &str) -> io::Result<Self> {
1742        log::info!("loading Eip spec from {}", file_path);
1743
1744        if !Path::new(file_path).exists() {
1745            return Err(io::Error::new(
1746                io::ErrorKind::NotFound,
1747                format!("file {} does not exists", file_path),
1748            ));
1749        }
1750
1751        let f = File::open(file_path).map_err(|e| {
1752            io::Error::new(
1753                io::ErrorKind::Other,
1754                format!("failed to open {} ({})", file_path, e),
1755            )
1756        })?;
1757
1758        serde_yaml::from_reader(f).map_err(|e| {
1759            io::Error::new(io::ErrorKind::InvalidInput, format!("invalid YAML: {}", e))
1760        })
1761    }
1762}
1763
1764/// RUST_LOG=debug cargo test --package aws-manager --lib -- ec2::test_eip --exact --show-output
1765#[test]
1766fn test_eip() {
1767    let d = r#"
1768allocation_id: test
1769public_ip: 1.2.3.4
1770
1771"#;
1772    let mut f = tempfile::NamedTempFile::new().unwrap();
1773    let ret = f.write_all(d.as_bytes());
1774    assert!(ret.is_ok());
1775    let eip_path = f.path().to_str().unwrap();
1776
1777    let ret = Eip::load(eip_path);
1778    assert!(ret.is_ok());
1779    let eip = ret.unwrap();
1780
1781    let ret = eip.sync(eip_path);
1782    assert!(ret.is_ok());
1783
1784    let orig = Eip {
1785        allocation_id: String::from("test"),
1786        public_ip: String::from("1.2.3.4"),
1787    };
1788    assert_eq!(eip, orig);
1789}