poll when instance running

Cargo.lock

@@ -1888,6 +1888,7 @@ dependencies = [
  "serde",
  "serde_derive",
  "serde_json",
+ "tokio",
  "toml",
 ]
 

lib/helpers/Cargo.toml

@@ -9,6 +9,7 @@ license = "GPL-3.0-or-later"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+tokio = "1.25"
 types = { package = "monitor_types", path = "../types" }
 periphery_client = { path = "../periphery_client" }
 async_timing_util = "0.1.14"

lib/helpers/src/aws/mod.rs

@@ -1,6 +1,15 @@
+use std::time::Duration;
+
 use anyhow::{anyhow, Context};
-use aws_sdk_ec2::{model::InstanceType, Client, Region};
-use periphery_client::PeripheryClient;
+use aws_sdk_ec2::model::{
+    BlockDeviceMapping, EbsBlockDevice, InstanceStateChange, InstanceStateName, InstanceStatus,
+};
+pub use aws_sdk_ec2::{
+    model::InstanceType,
+    output::{DescribeInstanceStatusOutput, TerminateInstancesOutput},
+    Client, Region,
+};
+use types::Server;
 
 /// must provide creds in env or with ~/.aws/credentials
 pub async fn create_ec2_client(
@@ -20,34 +29,44 @@ pub async fn create_ec2_client(
 
 pub struct Ec2Instance {
     pub id: String,
-    pub periphery_address: String,
+    pub server: Server,
 }
 
-const POLL_RATE_SECS: i32 = 1;
+const POLL_RATE_SECS: u64 = 2;
 const MAX_POLL_TRIES: usize = 30;
 
-/// should poll the periphery
+/// this will only resolve after the instance is running
+/// should still poll the periphery agent after creation
 pub async fn create_instance_with_ami(
-    client: Client,
+    client: &Client,
     ami_id: &str,
     instance_type: &str,
     security_groups_ids: &Vec<String>,
+    volume_size_gb: i32,
 ) -> anyhow::Result<Ec2Instance> {
     let instance_type = InstanceType::from(instance_type);
     if let InstanceType::Unknown(t) = instance_type {
         return Err(anyhow!("unknown instance type {t:?}"));
     }
+    let block_device_mapping = BlockDeviceMapping::builder()
+        .set_device_name(String::from("/dev/sda1").into())
+        .set_ebs(
+            EbsBlockDevice::builder()
+                .volume_size(volume_size_gb)
+                .build()
+                .into(),
+        )
+        .build();
     let mut req = client
         .run_instances()
         .image_id(ami_id)
         .instance_type(instance_type)
+        .block_device_mappings(block_device_mapping)
         .min_count(1)
         .max_count(1);
-
     for sg_id in security_groups_ids {
         req = req.security_group_ids(sg_id);
     }
-
     let res = req
         .send()
         .await
@@ -64,10 +83,72 @@ pub async fn create_instance_with_ami(
     let ip = instance
         .private_ip_address()
         .ok_or(anyhow!("instance does not have private ip"))?;
-    // client.describe_instances().
-    let periphery_address = format!("http://{ip}:8000");
-    Ok(Ec2Instance {
-        id: instance_id,
-        periphery_address,
-    })
+    let server = Server {
+        address: format!("http://{ip}:8000"),
+        ..Default::default()
+    };
+    for _ in 0..MAX_POLL_TRIES {
+        let state_name = get_ec2_instance_state_name(&client, &instance_id).await?;
+        if state_name == Some(InstanceStateName::Running) {
+            return Ok(Ec2Instance {
+                id: instance_id,
+                server,
+            });
+        }
+        tokio::time::sleep(Duration::from_secs(POLL_RATE_SECS)).await;
+    }
+    Err(anyhow!("instance not running after polling"))
+}
+
+pub async fn get_ec2_instance_status(
+    client: &Client,
+    instance_id: &str,
+) -> anyhow::Result<Option<InstanceStatus>> {
+    let status = client
+        .describe_instance_status()
+        .instance_ids(instance_id)
+        .send()
+        .await
+        .context("failed to get instance status from aws")?
+        .instance_statuses()
+        .ok_or(anyhow!("instance statuses is None"))?
+        .get(0)
+        .map(|s| s.to_owned());
+    Ok(status)
+}
+
+pub async fn get_ec2_instance_state_name(
+    client: &Client,
+    instance_id: &str,
+) -> anyhow::Result<Option<InstanceStateName>> {
+    let status = get_ec2_instance_status(client, instance_id).await?;
+    if status.is_none() {
+        return Ok(None);
+    }
+    let state = status
+        .unwrap()
+        .instance_state()
+        .ok_or(anyhow!("instance state is None"))?
+        .name()
+        .ok_or(anyhow!("instance state name is None"))?
+        .to_owned();
+    Ok(Some(state))
+}
+
+pub async fn terminate_ec2_instance(
+    client: &Client,
+    instance_id: &str,
+) -> anyhow::Result<InstanceStateChange> {
+    let res = client
+        .terminate_instances()
+        .instance_ids(instance_id)
+        .send()
+        .await
+        .context("failed to terminate instance from aws")?
+        .terminating_instances()
+        .ok_or(anyhow!("terminating instances is None"))?
+        .get(0)
+        .ok_or(anyhow!("terminating instances is empty"))?
+        .to_owned();
+    Ok(res)
 }