From 38c47e9e4ef2d0e16a04cf9900e2784832b1731f Mon Sep 17 00:00:00 2001 From: alemidev Date: Fri, 30 Sep 2022 23:10:18 +0200 Subject: [PATCH] feat: track cpu and disk (jank but works) --- src/main.rs | 175 +++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 132 insertions(+), 43 deletions(-) diff --git a/src/main.rs b/src/main.rs index 1e9c86e..3858d90 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,20 +1,63 @@ -#[macro_use] extern crate rocket; +#[macro_use] +extern crate rocket; use std::collections::HashMap; +use std::time::Duration; -use rocket::State; -use rocket::serde::{Serialize, json::Json}; -use rocket::tokio::sync::Mutex; -use systemstat::{System, Platform}; +use rocket::{ + response::status::Unauthorized, + serde::{json::Json, Serialize}, + tokio::{sync::Mutex, time::sleep}, + State, +}; +use systemstat::{CPULoad, DelayedMeasurement, Platform, System, BTreeMap, BlockDeviceStats}; + +fn to_403_err(e: T) -> Unauthorized +where + T: std::error::Error, +{ + Unauthorized(Some(e.to_string())) +} #[derive(Serialize)] struct SystemInfoView { session: String, mem: f32, cpu: f32, - tx: u64, - rx: u64, - disk: u64, + tx: f32, + rx: f32, + r: f32, + w: f32, +} + +struct DiskDeltaTracker { + r: u64, + w: u64, +} + +impl Default for DiskDeltaTracker { + fn default() -> Self { + DiskDeltaTracker { r: 0, w: 0 } + } +} + +impl DiskDeltaTracker { + pub fn delta_all(&mut self, disks: BTreeMap) -> (u64, u64) { + let mut r_tot : u64 = 0; + let mut w_tot : u64 = 0; + for drive in disks.values() { + r_tot += drive.read_sectors as u64; + w_tot += drive.write_sectors as u64; + } + return self.delta(r_tot, w_tot); + } + + pub fn delta(&mut self, r: u64, w: u64) -> (u64, u64) { + let out = (r - self.r, w - self.w); + self.r = r; + self.w = w; + return out; + } } struct NetDeltaTracker { @@ -37,44 +80,87 @@ impl NetDeltaTracker { } } -struct NetDeltaState { - pub store: Mutex>, +struct CpuDeltaTracker { + load: DelayedMeasurement, +} + +struct StateDeltaStore { + pub store_disk: Mutex>, + pub store_net: Mutex>, + pub store_cpu: Mutex>, } #[get("/")] -async fn stats_delta(session: String, state: &State) -> Json { +async fn stats_delta( + session: String, + state: &State, +) -> Result, Unauthorized> { let sys = System::new(); let (delta_tx, delta_rx); let (net_tx, net_rx) = match sys.network_stats(&"enp5s0".to_string()) { Ok(stats) => (stats.tx_bytes.as_u64(), stats.rx_bytes.as_u64()), - Err(x) => (0, 0), + Err(_e) => (0, 0), }; { - let mut store = state.inner().store.lock().await; - if !store.contains_key(&session) { - store.insert(session.clone(), NetDeltaTracker::default()); + let mut store_net = state.inner().store_net.lock().await; + if !store_net.contains_key(&session) { + store_net.insert(session.clone(), NetDeltaTracker::default()); } - let net_tracker = store.get_mut(&session).unwrap(); + let net_tracker = store_net.get_mut(&session).unwrap(); (delta_tx, delta_rx) = net_tracker.delta(net_tx, net_rx); } - Json( - SystemInfoView { - session, - mem: match sys.memory() { - Ok(mem) => { 1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32)) }, - Err(e) => { -1.0 }, - }, - cpu: 0.0, - tx: delta_tx, - rx: delta_rx, - disk: 0, + let cpu; + + { + let mut store_cpu = state.inner().store_cpu.lock().await; + if !store_cpu.contains_key(&session) { + let cpu = sys.cpu_load_aggregate().map_err(to_403_err)?; + store_cpu.insert(session.clone(), CpuDeltaTracker { load: cpu }); + sleep(Duration::from_secs(1)).await; } - ) + + let cpu_tracker = store_cpu.get_mut(&session).unwrap(); + + let cpu_stats = cpu_tracker.load.done().map_err(to_403_err)?; + cpu = 1.0 - cpu_stats.idle; + + let new_cpu = sys.cpu_load_aggregate().map_err(to_403_err)?; + store_cpu.insert(session.clone(), CpuDeltaTracker { load: new_cpu }); + } + + let mem = match sys.memory() { + Ok(mem) => 100.0 * (1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32))), + Err(_e) => -1.0, + }; + + let (disk_r, disk_w); + + { + let mut store_disk = state.inner().store_disk.lock().await; + if !store_disk.contains_key(&session) { + store_disk.insert(session.clone(), DiskDeltaTracker::default()); + } + + let net_tracker = store_disk.get_mut(&session).unwrap(); + + let disks = sys.block_device_statistics().map_err(to_403_err)?; + (disk_r, disk_w) = net_tracker.delta_all(disks); + } + + Ok(Json(SystemInfoView { + session, + mem, + cpu, + tx: delta_tx as f32 / (1024.0 * 1024.0), + rx: delta_rx as f32 / (1024.0 * 1024.0), + r: disk_r as f32 / (1024.0 * 2.0), + w: disk_w as f32 / (1024.0 * 2.0), + })) } #[get("/")] @@ -83,27 +169,30 @@ async fn stats_total() -> Json { let (net_tx, net_rx) = match sys.network_stats(&"enp5s0".to_string()) { Ok(stats) => (stats.tx_bytes.as_u64(), stats.rx_bytes.as_u64()), - Err(x) => (0, 0), + Err(_e) => (0, 0), }; - Json( - SystemInfoView { - session: "".to_string(), - mem: match sys.memory() { - Ok(mem) => { 1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32)) }, - Err(e) => { -1.0 }, - }, - cpu: 0.0, - tx: net_tx, - rx: net_rx, - disk: 0, - } - ) + Json(SystemInfoView { + session: "".to_string(), + mem: match sys.memory() { + Ok(mem) => 1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32)), + Err(_e) => -1.0, + }, + cpu: 0.0, + tx: net_tx as f32 / (1024.0 * 1024.0), + rx: net_rx as f32 / (1024.0 * 1024.0), + r: 0.0, + w: 0.0, + }) } #[launch] fn rocket() -> _ { rocket::build() - .manage(NetDeltaState { store: Mutex::new(HashMap::new()) }) + .manage(StateDeltaStore { + store_disk: Mutex::new(HashMap::new()), + store_net: Mutex::new(HashMap::new()), + store_cpu: Mutex::new(HashMap::new()), + }) .mount("/", routes![stats_delta, stats_total]) }