diff --git a/lnvps_db/src/lib.rs b/lnvps_db/src/lib.rs index 162b59e..9c3cd29 100644 --- a/lnvps_db/src/lib.rs +++ b/lnvps_db/src/lib.rs @@ -83,6 +83,9 @@ pub trait LNVpsDb: Sync + Send { /// List all VM's async fn list_vms(&self) -> Result>; + /// List all VM's on a given host + async fn list_vms_on_host(&self, host_id: u64) -> Result>; + /// List expired VM's async fn list_expired_vms(&self) -> Result>; diff --git a/lnvps_db/src/mysql.rs b/lnvps_db/src/mysql.rs index 0502a37..6701d97 100644 --- a/lnvps_db/src/mysql.rs +++ b/lnvps_db/src/mysql.rs @@ -223,6 +223,14 @@ impl LNVpsDb for LNVpsDbMysql { .map_err(Error::new) } + async fn list_vms_on_host(&self, host_id: u64) -> Result> { + sqlx::query_as("select * from vm where deleted = 0 and host_id = ?") + .bind(host_id) + .fetch_all(&self.db) + .await + .map_err(Error::new) + } + async fn list_expired_vms(&self) -> Result> { sqlx::query_as("select * from vm where expires > current_timestamp() and deleted = 0") .fetch_all(&self.db) diff --git a/src/mocks.rs b/src/mocks.rs index b9fc839..2e3e7aa 100644 --- a/src/mocks.rs +++ b/src/mocks.rs @@ -326,6 +326,15 @@ impl LNVpsDb for MockDb { Ok(vms.values().filter(|v| !v.deleted).cloned().collect()) } + async fn list_vms_on_host(&self, host_id: u64) -> anyhow::Result> { + let vms = self.vms.lock().await; + Ok(vms + .values() + .filter(|v| !v.deleted && v.host_id == host_id) + .cloned() + .collect()) + } + async fn list_expired_vms(&self) -> anyhow::Result> { let vms = self.vms.lock().await; Ok(vms diff --git a/src/provisioner/capacity.rs b/src/provisioner/capacity.rs new file mode 100644 index 0000000..f9b6b6e --- /dev/null +++ b/src/provisioner/capacity.rs @@ -0,0 +1,105 @@ +use anyhow::Result; +use lnvps_db::{LNVpsDb, VmHost, VmHostDisk, VmTemplate}; +use std::collections::HashMap; +use std::sync::Arc; + +/// Simple capacity reporting per node +#[derive(Clone)] +pub struct HostCapacity { + db: Arc, +} + +impl HostCapacity { + pub fn new(db: Arc) -> Self { + Self { db } + } + + pub async fn get_available_capacity(&self, host: &VmHost) -> Result { + let vms = self.db.list_vms_on_host(host.id).await?; + let storage = self.db.list_host_disks(host.id).await?; + let templates = self.db.list_vm_templates().await?; + + // a mapping between vm_id and template + let vm_template: HashMap = vms + .iter() + .filter_map(|v| { + templates + .iter() + .find(|t| t.id == v.template_id) + .and_then(|t| Some((v.id, t))) + }) + .collect(); + + let storage_disks: Vec = storage + .iter() + .map(|s| { + let usage = vm_template + .iter() + .filter(|(k, v)| v.id == s.id) + .fold(0, |acc, (k, v)| acc + v.disk_size); + DiskCapacity { + disk: s.clone(), + usage, + } + }) + .collect(); + + let cpu_consumed = vm_template.values().fold(0, |acc, vm| acc + vm.cpu); + let memory_consumed = vm_template.values().fold(0, |acc, vm| acc + vm.memory); + + Ok(AvailableCapacity { + cpu: host.cpu.saturating_sub(cpu_consumed), + memory: host.memory.saturating_sub(memory_consumed), + disks: storage_disks, + }) + } +} + +#[derive(Debug, Clone)] +pub struct AvailableCapacity { + /// Number of CPU cores available + pub cpu: u16, + /// Number of bytes of memory available + pub memory: u64, + /// List of disks on the host and its available space + pub disks: Vec, +} + +#[derive(Debug, Clone)] +pub struct DiskCapacity { + /// Disk ID + pub disk: VmHostDisk, + /// Space consumed by VMs + pub usage: u64, +} + +impl DiskCapacity { + pub fn available_capacity(&self) -> u64 { + self.disk.size.saturating_sub(self.usage) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mocks::MockDb; + + #[tokio::test] + async fn empty_available_capacity() -> Result<()> { + let db = Arc::new(MockDb::default()); + + let hc = HostCapacity::new(db.clone()); + let host = db.get_host(1).await?; + let cap = hc.get_available_capacity(&host).await?; + let disks = db.list_host_disks(1).await?; + /// check all resources are available + assert_eq!(cap.cpu, host.cpu); + assert_eq!(cap.memory, host.memory); + assert_eq!(cap.disks.len(), disks.len()); + for disk in cap.disks { + assert_eq!(0, disk.usage); + } + + Ok(()) + } +} diff --git a/src/provisioner/mod.rs b/src/provisioner/mod.rs index 13d3004..f15d062 100644 --- a/src/provisioner/mod.rs +++ b/src/provisioner/mod.rs @@ -1,5 +1,7 @@ +mod capacity; mod lnvps; mod network; +pub use capacity::*; pub use lnvps::*; pub use network::*;