feat: track cpu and disk (jank but works)

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<T>(e: T) -> Unauthorized<String>
+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<String, BlockDeviceStats>) -> (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<HashMap<String, NetDeltaTracker>>,
+struct CpuDeltaTracker {
+	load: DelayedMeasurement<CPULoad>,
+}
+
+struct StateDeltaStore {
+	pub store_disk: Mutex<HashMap<String, DiskDeltaTracker>>,
+	pub store_net: Mutex<HashMap<String, NetDeltaTracker>>,
+	pub store_cpu: Mutex<HashMap<String, CpuDeltaTracker>>,
 }
 
 #[get("/<session>")]
-async fn stats_delta(session: String, state: &State<NetDeltaState>) -> Json<SystemInfoView> {
+async fn stats_delta(
+	session: String,
+	state: &State<StateDeltaStore>,
+) -> Result<Json<SystemInfoView>, Unauthorized<String>> {
 	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<SystemInfoView> {
 
 	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 {
+	Json(SystemInfoView {
 		session: "<TOTAL>".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 },
+			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,
-		}
-	)
+		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])
 }