From 72a2ab98158e956b95ed66be622e51294f9fd412 Mon Sep 17 00:00:00 2001
From: kieran <kieran@harkin.me>
Date: Fri, 29 Nov 2024 22:56:57 +0000
Subject: [PATCH] feat: add more dists

---
 lnvps_db/src/model.rs    |  6 ++++++
 src/api.rs               | 10 ++++++++--
 src/provisioner/lnvps.rs |  7 ++++---
 src/status.rs            |  5 +++--
 src/worker.rs            | 18 ++++++++++++++++++
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/lnvps_db/src/model.rs b/lnvps_db/src/model.rs
index 7774853..31e5689 100644
--- a/lnvps_db/src/model.rs
+++ b/lnvps_db/src/model.rs
@@ -113,6 +113,12 @@ pub enum OsDistribution {
     #[default]
     Ubuntu = 0,
     Debian = 1,
+    CentOS = 2,
+    Fedora = 3,
+    FreeBSD = 4,
+    OpenSUSE = 5,
+    ArchLinux = 6,
+    RedHatEnterprise = 7,
 }
 
 /// OS Images are templates which are used as a basis for
diff --git a/src/api.rs b/src/api.rs
index 556d890..7705f94 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -83,7 +83,10 @@ async fn v1_list_vms(
         }
 
         let state = vm_state.get_state(vm.id).await;
-        ret.push(ApiVmStatus { vm, status: state });
+        ret.push(ApiVmStatus {
+            vm,
+            status: state.unwrap_or_default(),
+        });
     }
 
     ApiData::ok(ret)
@@ -108,7 +111,10 @@ async fn v1_get_vm(
         t.hydrate_up(db.inner()).await?;
     }
     let state = vm_state.get_state(vm.id).await;
-    ApiData::ok(ApiVmStatus { vm, status: state })
+    ApiData::ok(ApiVmStatus {
+        vm,
+        status: state.unwrap_or_default(),
+    })
 }
 
 #[get("/api/v1/image")]
diff --git a/src/provisioner/lnvps.rs b/src/provisioner/lnvps.rs
index 166c6d4..6902ae5 100644
--- a/src/provisioner/lnvps.rs
+++ b/src/provisioner/lnvps.rs
@@ -246,7 +246,7 @@ impl Provisioner for LNVpsProvisioner {
         'ranges: for range in ip_ranges {
             let range_cidr: IpNetwork = range.cidr.parse()?;
             let ips = self.db.list_vm_ip_assignments_in_range(range.id).await?;
-            let ips: HashSet<IpAddr> = ips.iter().map_while(|i| i.ip.parse().ok()).collect();
+            let ips: HashSet<IpNetwork> = ips.iter().map_while(|i| i.ip.parse().ok()).collect();
 
             // pick an IP at random
             let cidr: Vec<IpAddr> = {
@@ -255,13 +255,14 @@ impl Provisioner for LNVpsProvisioner {
             };
 
             for ip in cidr {
-                if !ips.contains(&ip) {
+                let ip_net = IpNetwork::new(ip, range_cidr.prefix())?;
+                if !ips.contains(&ip_net) {
                     info!("Attempting to allocate IP for {vm_id} to {ip}");
                     let mut assignment = VmIpAssignment {
                         id: 0,
                         vm_id,
                         ip_range_id: range.id,
-                        ip: IpNetwork::new(ip, range_cidr.prefix())?.to_string(),
+                        ip: ip_net.to_string(),
                         ..Default::default()
                     };
                     let id = self.db.insert_vm_ip_assignment(&assignment).await?;
diff --git a/src/status.rs b/src/status.rs
index 12ef585..897e6e1 100644
--- a/src/status.rs
+++ b/src/status.rs
@@ -16,6 +16,7 @@ pub enum VmRunningState {
 
 #[derive(Clone, Serialize, Default)]
 pub struct VmState {
+    pub timestamp: u64,
     pub state: VmRunningState,
     pub cpu_usage: f32,
     pub mem_usage: f32,
@@ -51,8 +52,8 @@ impl VmStateCache {
         Ok(())
     }
 
-    pub async fn get_state(&self, id: u64) -> VmState {
+    pub async fn get_state(&self, id: u64) -> Option<VmState> {
         let guard = self.state.read().await;
-        guard.get(&id).cloned().unwrap_or_default()
+        guard.get(&id).cloned()
     }
 }
diff --git a/src/worker.rs b/src/worker.rs
index 83b9cc6..eea25a8 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -55,6 +55,7 @@ impl Worker {
         // TODO: remove assumption
         let db_id = (s.vm_id - 100) as u64;
         let state = VmState {
+            timestamp: Utc::now().timestamp() as u64,
             state: match s.status {
                 VmStatus::Stopped => VmRunningState::Stopped,
                 VmStatus::Running => VmRunningState::Running,
@@ -118,6 +119,23 @@ impl Worker {
                 }
             }
         }
+
+        let db_vms = self.db.list_vms().await?;
+        for vm in db_vms {
+            let state = if let Some(s) = self.vm_state_cache.get_state(vm.id).await {
+                if s.timestamp > Utc::now().timestamp() as u64 - 120 {
+                    Some(s)
+                } else {
+                    None
+                }
+            } else {
+                None
+            };
+
+            if state.is_none() && vm.expires > Utc::now() {
+                self.check_vm(vm.id).await?;
+            }
+        }
         Ok(())
     }