iota_aws_orchestrator/client/

aws.rs

// Copyright (c) Mysten Labs, Inc.
// Modifications Copyright (c) 2024 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use std::{
    collections::HashMap,
    fmt::{Debug, Display},
};

use aws_runtime::env_config::file::{EnvConfigFileKind, EnvConfigFiles};
use aws_sdk_ec2::{
    config::Region,
    primitives::Blob,
    types::{
        BlockDeviceMapping, EbsBlockDevice, EphemeralNvmeSupport, Filter, ResourceType, Tag,
        TagSpecification, VolumeType, builders::FilterBuilder,
    },
};
use aws_smithy_runtime_api::client::{behavior_version::BehaviorVersion, result::SdkError};
use serde::Serialize;

use super::{Instance, InstanceRole, ServerProviderClient};
use crate::{
    error::{CloudProviderError, CloudProviderResult},
    settings::Settings,
};

// Make a request error from an AWS error message.
impl<T> From<SdkError<T, aws_smithy_runtime_api::client::orchestrator::HttpResponse>>
    for CloudProviderError
where
    T: Debug + std::error::Error + Send + Sync + 'static,
{
    fn from(e: SdkError<T, aws_smithy_runtime_api::client::orchestrator::HttpResponse>) -> Self {
        Self::Request(format!("{:?}", e.into_source()))
    }
}

/// A AWS client.
pub struct AwsClient {
    /// The settings of the testbed.
    settings: Settings,
    /// A list of clients, one per AWS region.
    clients: HashMap<String, aws_sdk_ec2::Client>,
}

impl Display for AwsClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "AWS EC2 client v{}", aws_sdk_ec2::meta::PKG_VERSION)
    }
}

impl AwsClient {
    const UBUNTU_NAME_PATTERN: &'static str =
        "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*";
    const CANONICAL_OWNER_ID: &'static str = "099720109477";

    /// Make a new AWS client.
    pub async fn new(settings: Settings) -> Self {
        let profile_files = EnvConfigFiles::builder()
            .with_file(EnvConfigFileKind::Credentials, &settings.token_file)
            .with_contents(EnvConfigFileKind::Config, "[default]\noutput=json")
            .build();

        let mut clients = HashMap::new();
        for region in settings.regions.clone() {
            let sdk_config = aws_config::defaults(BehaviorVersion::latest())
                .region(Region::new(region.clone()))
                .profile_files(profile_files.clone())
                .load()
                .await;
            let client = aws_sdk_ec2::Client::new(&sdk_config);
            clients.insert(region, client);
        }

        Self { settings, clients }
    }

    /// Parse an AWS response and ignore errors if they mean a request is a
    /// duplicate.
    fn check_but_ignore_duplicates<T, E>(
        response: Result<
            T,
            SdkError<E, aws_smithy_runtime_api::client::orchestrator::HttpResponse>,
        >,
    ) -> CloudProviderResult<()>
    where
        E: Debug + std::error::Error + Send + Sync + 'static,
    {
        if let Err(e) = response {
            let error_message = format!("{e:?}");
            if !error_message.to_lowercase().contains("duplicate") {
                return Err(e.into());
            }
        }
        Ok(())
    }
    fn get_tag_value(instance: &aws_sdk_ec2::types::Instance, key: &str) -> Option<String> {
        instance
            .tags()
            .iter()
            .find(|tag| tag.key().is_some_and(|k| k == key))
            .and_then(|tag| tag.value().map(|v| v.to_string()))
    }
    /// Convert an AWS instance into an orchestrator instance (used in the rest
    /// of the codebase).
    fn make_instance(
        &self,
        region: String,
        aws_instance: &aws_sdk_ec2::types::Instance,
    ) -> Instance {
        let role: InstanceRole = Self::get_tag_value(aws_instance, "Role")
            .expect("AWS instance should have a role")
            .as_str()
            .into();
        Instance {
            id: aws_instance
                .instance_id()
                .expect("AWS instance should have an id")
                .into(),
            region,
            main_ip: aws_instance
                .public_ip_address()
                .unwrap_or("0.0.0.0") // Stopped instances do not have an ip address.
                .parse()
                .expect("AWS instance should have a valid ip"),
            private_ip: aws_instance
                .private_ip_address()
                .unwrap_or("0.0.0.0") // Stopped instances do not have an ip address.
                .parse()
                .expect("AWS instance should have a valid ip"),
            tags: vec![self.settings.testbed_id.clone()],
            specs: format!(
                "{:?}",
                aws_instance
                    .instance_type()
                    .expect("AWS instance should have a type")
            ),
            status: format!(
                "{:?}",
                aws_instance
                    .state()
                    .expect("AWS instance should have a state")
                    .name()
                    .expect("AWS status should have a name")
            ),
            role,
        }
    }

    /// Query the image id determining the os of the instances.
    /// NOTE: The image id changes depending on the region.
    async fn find_image_id(&self, client: &aws_sdk_ec2::Client) -> CloudProviderResult<String> {
        // Use a more general filter that doesn't depend on specific build dates
        let filters = [
            // Filter for Ubuntu 24.04 LTS
            FilterBuilder::default()
                .name("name")
                .values(Self::UBUNTU_NAME_PATTERN)
                .build(),
            // Only look at images from Canonical
            FilterBuilder::default()
                .name("owner-id")
                .values(Self::CANONICAL_OWNER_ID)
                .build(),
            // Only want available images
            FilterBuilder::default()
                .name("state")
                .values("available")
                .build(),
        ];

        // Query images with these filters
        let request = client.describe_images().set_filters(Some(filters.to_vec()));
        let response = request.send().await?;

        // Sort images by creation date (newest first)
        let mut images = response.images().to_vec();
        images.sort_by(|a, b| {
            let a_date = a.creation_date().unwrap_or("");
            let b_date = b.creation_date().unwrap_or("");
            b_date.cmp(a_date) // Reverse order to get newest first
        });

        // Select the newest image
        let image = images
            .first()
            .ok_or_else(|| CloudProviderError::Request("Cannot find Ubuntu 24.04 image".into()))?;

        image
            .image_id
            .clone()
            .ok_or_else(|| CloudProviderError::UnexpectedResponse("Image without ID".into()))
    }

    /// Create a new security group for the instance (if it doesn't already
    /// exist).
    async fn create_security_group(&self, client: &aws_sdk_ec2::Client) -> CloudProviderResult<()> {
        // Create a security group (if it doesn't already exist).
        let request = client
            .create_security_group()
            .group_name(&self.settings.testbed_id)
            .description("Allow all traffic (used for benchmarks).");

        let response = request.send().await;
        Self::check_but_ignore_duplicates(response)?;

        // Authorize all traffic on the security group.
        for protocol in ["tcp", "udp", "icmp", "icmpv6"] {
            let mut request = client
                .authorize_security_group_ingress()
                .group_name(&self.settings.testbed_id)
                .ip_protocol(protocol)
                .cidr_ip("0.0.0.0/0"); // todo - allowing 0.0.0.0 seem a bit wild?
            if protocol == "icmp" || protocol == "icmpv6" {
                request = request.from_port(-1).to_port(-1);
            } else {
                request = request.from_port(0).to_port(65535);
            }

            let response = request.send().await;
            Self::check_but_ignore_duplicates(response)?;
        }
        Ok(())
    }

    /// Return the command to mount the first (standard) NVMe drive.
    fn nvme_mount_command(&self) -> Vec<String> {
        const DRIVE: &str = "nvme1n1";
        let directory = self.settings.working_dir.display();
        vec![
            format!("(sudo mkfs.ext4 -E nodiscard /dev/{DRIVE} || true)"),
            format!("(sudo mount /dev/{DRIVE} {directory} || true)"),
            format!("sudo chmod 777 -R {directory}"),
        ]
    }

    /// Check whether the instance type specified in the settings supports NVMe
    /// drives.
    async fn check_nvme_support(&self) -> CloudProviderResult<bool> {
        // Get the client for the first region. A given instance type should either have
        // NVMe support in all regions or in none.
        let client = match self
            .settings
            .regions
            .first()
            .and_then(|x| self.clients.get(x))
        {
            Some(client) => client,
            None => return Ok(false),
        };

        // Request storage details for the instance type specified in the settings.
        let request = client
            .describe_instance_types()
            .instance_types(self.settings.node_specs.as_str().into());

        // Send the request.
        let response = request.send().await?;

        // Return true if the response contains references to NVMe drives.
        if let Some(info) = response.instance_types().first() {
            if let Some(info) = info.instance_storage_info() {
                if info.nvme_support() == Some(&EphemeralNvmeSupport::Required) {
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }
}

#[async_trait::async_trait]
impl ServerProviderClient for AwsClient {
    const USERNAME: &'static str = "ubuntu";

    async fn list_instances_by_role(
        &self,
        role: InstanceRole,
    ) -> CloudProviderResult<Vec<Instance>> {
        let filter_name = Filter::builder()
            .name("tag:Name")
            .values(self.settings.testbed_id.clone())
            .build();
        let filter_role = Filter::builder()
            .name("tag:Role")
            .values(role.to_string())
            .build();
        let filter_state = Filter::builder()
            .name("instance-state-name")
            .values("pending")
            .values("running")
            .values("stopping")
            .values("stopped")
            .build();

        let mut instances = Vec::new();
        for (region, client) in &self.clients {
            let request = client.describe_instances().set_filters(Some(vec![
                filter_name.clone(),
                filter_role.clone(),
                filter_state.clone(),
            ]));
            let response = request.send().await?;
            for reservation in response.reservations() {
                for instance in reservation.instances() {
                    instances.push(self.make_instance(region.clone(), instance));
                }
            }
        }

        Ok(instances)
    }

    async fn list_instances_by_region_and_ids(
        &self,
        regions_and_ids: Vec<(String, String)>,
    ) -> CloudProviderResult<Vec<Instance>> {
        let mut instances = Vec::new();
        let ids_by_region: HashMap<String, Vec<String>> =
            regions_and_ids
                .into_iter()
                .fold(HashMap::new(), |mut acc, (k, v)| {
                    acc.entry(k).or_default().push(v);
                    acc
                });
        for (region, client) in &self.clients {
            let request = client
                .describe_instances()
                .set_instance_ids(ids_by_region.get(region).cloned());
            let response = request.send().await?;
            for reservation in response.reservations() {
                for instance in reservation.instances() {
                    instances.push(self.make_instance(region.clone(), instance));
                }
            }
        }

        Ok(instances)
    }

    async fn start_instances<'a, I>(&self, instances: I) -> CloudProviderResult<()>
    where
        I: Iterator<Item = &'a Instance> + Send,
    {
        let mut instance_ids = HashMap::new();
        for instance in instances {
            instance_ids
                .entry(&instance.region)
                .or_insert_with(Vec::new)
                .push(instance.id.clone());
        }

        for (region, client) in &self.clients {
            let ids = instance_ids.remove(&region.to_string());
            if ids.is_some() {
                client
                    .start_instances()
                    .set_instance_ids(ids)
                    .send()
                    .await?;
            }
        }
        Ok(())
    }

    async fn stop_instances<'a, I>(&self, instances: I) -> CloudProviderResult<()>
    where
        I: Iterator<Item = &'a Instance> + Send,
    {
        let mut instance_ids = HashMap::new();
        for instance in instances {
            instance_ids
                .entry(&instance.region)
                .or_insert_with(Vec::new)
                .push(instance.id.clone());
        }

        for (region, client) in &self.clients {
            let ids = instance_ids.remove(&region.to_string());
            if ids.is_some() {
                client.stop_instances().set_instance_ids(ids).send().await?;
            }
        }
        Ok(())
    }

    async fn create_instance<S>(
        &self,
        region: S,
        role: InstanceRole,
    ) -> CloudProviderResult<Instance>
    where
        S: Into<String> + Serialize + Send,
    {
        let region = region.into();
        let testbed_id = &self.settings.testbed_id;

        let client = self
            .clients
            .get(&region)
            .ok_or_else(|| CloudProviderError::Request(format!("Undefined region {region:?}")))?;

        // Create a security group (if needed).
        self.create_security_group(client).await?;

        // Query the image id.
        let image_id = self.find_image_id(client).await?;

        // Create a new instance.
        let tags = TagSpecification::builder()
            .resource_type(ResourceType::Instance)
            .tags(Tag::builder().key("Name").value(testbed_id).build())
            .tags(Tag::builder().key("Role").value(role.to_string()).build())
            .build();

        let storage = BlockDeviceMapping::builder()
            .device_name("/dev/sda1")
            .ebs(
                EbsBlockDevice::builder()
                    .delete_on_termination(true)
                    .volume_size(500)
                    .volume_type(VolumeType::Gp2)
                    .build(),
            )
            .build();
        let instance_type = match role {
            InstanceRole::Node => &self.settings.node_specs,
            InstanceRole::Metrics => &self.settings.metrics_specs,
            InstanceRole::Client => &self.settings.client_specs,
        };
        let request = client
            .run_instances()
            .image_id(image_id)
            .instance_type(instance_type.as_str().into())
            .key_name(testbed_id)
            .min_count(1)
            .max_count(1)
            .security_groups(&self.settings.testbed_id)
            .block_device_mappings(storage)
            .tag_specifications(tags);

        let response = request.send().await?;
        let instance = &response
            .instances()
            .first()
            .expect("AWS instances list should contain instances");

        Ok(self.make_instance(region, instance))
    }

    async fn delete_instance(&self, instance: Instance) -> CloudProviderResult<()> {
        let client = self.clients.get(&instance.region).ok_or_else(|| {
            CloudProviderError::Request(format!("Undefined region {:?}", instance.region))
        })?;

        client
            .terminate_instances()
            .set_instance_ids(Some(vec![instance.id.clone()]))
            .send()
            .await?;

        Ok(())
    }

    async fn register_ssh_public_key(&self, public_key: String) -> CloudProviderResult<()> {
        for client in self.clients.values() {
            let request = client
                .import_key_pair()
                .key_name(&self.settings.testbed_id)
                .public_key_material(Blob::new::<String>(public_key.clone()));

            let response = request.send().await;
            Self::check_but_ignore_duplicates(response)?;
        }
        Ok(())
    }

    async fn instance_setup_commands(&self) -> CloudProviderResult<Vec<String>> {
        if self.check_nvme_support().await? {
            Ok(self.nvme_mount_command())
        } else {
            Ok(Vec::new())
        }
    }
    #[cfg(test)]
    fn instances(&self) -> Vec<Instance> {
        // Only used under testing by the TestClient, unreachable cause no test
        // should use AwsClient.
        unreachable!()
    }
}