code optimizations to reduce cpu usage of agent on all platforms and additional unit test.

This commit is contained in:
jasonwitty 2025-08-28 16:03:05 -07:00
parent 7592709a43
commit df2308e6e9
4 changed files with 214 additions and 29 deletions

40
Cargo.lock generated
View File

@ -158,7 +158,7 @@ dependencies = [
"sha1", "sha1",
"sync_wrapper", "sync_wrapper",
"tokio", "tokio",
"tokio-tungstenite", "tokio-tungstenite 0.24.0",
"tower 0.5.2", "tower 0.5.2",
"tower-layer", "tower-layer",
"tower-service", "tower-service",
@ -2181,13 +2181,13 @@ dependencies = [
"sysinfo", "sysinfo",
"tempfile", "tempfile",
"tokio", "tokio",
"tokio-tungstenite", "tokio-tungstenite 0.24.0",
"url", "url",
] ]
[[package]] [[package]]
name = "socktop_agent" name = "socktop_agent"
version = "1.40.66" version = "1.40.67"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"assert_cmd", "assert_cmd",
@ -2210,6 +2210,7 @@ dependencies = [
"tempfile", "tempfile",
"time", "time",
"tokio", "tokio",
"tokio-tungstenite 0.21.0",
"tonic-build", "tonic-build",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
@ -2463,6 +2464,18 @@ dependencies = [
"tokio", "tokio",
] ]
[[package]]
name = "tokio-tungstenite"
version = "0.21.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c83b561d025642014097b66e6c1bb422783339e0909e4429cde4749d1990bc38"
dependencies = [
"futures-util",
"log",
"tokio",
"tungstenite 0.21.0",
]
[[package]] [[package]]
name = "tokio-tungstenite" name = "tokio-tungstenite"
version = "0.24.0" version = "0.24.0"
@ -2475,7 +2488,7 @@ dependencies = [
"rustls-pki-types", "rustls-pki-types",
"tokio", "tokio",
"tokio-rustls 0.26.2", "tokio-rustls 0.26.2",
"tungstenite", "tungstenite 0.24.0",
] ]
[[package]] [[package]]
@ -2608,6 +2621,25 @@ dependencies = [
"tracing-log", "tracing-log",
] ]
[[package]]
name = "tungstenite"
version = "0.21.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ef1a641ea34f399a848dea702823bbecfb4c486f911735368f1f137cb8257e1"
dependencies = [
"byteorder",
"bytes",
"data-encoding",
"http",
"httparse",
"log",
"rand",
"sha1",
"thiserror 1.0.69",
"url",
"utf-8",
]
[[package]] [[package]]
name = "tungstenite" name = "tungstenite"
version = "0.24.0" version = "0.24.0"

View File

@ -1,6 +1,6 @@
[package] [package]
name = "socktop_agent" name = "socktop_agent"
version = "1.40.66" version = "1.40.67"
authors = ["Jason Witty <jasonpwitty+socktop@proton.me>"] authors = ["Jason Witty <jasonpwitty+socktop@proton.me>"]
description = "Remote system monitor over WebSocket, TUI like top" description = "Remote system monitor over WebSocket, TUI like top"
edition = "2021" edition = "2021"
@ -36,3 +36,4 @@ protoc-bin-vendored = "3"
[dev-dependencies] [dev-dependencies]
assert_cmd = "2.0" assert_cmd = "2.0"
tempfile = "3.10" tempfile = "3.10"
tokio-tungstenite = "0.21"

View File

@ -49,6 +49,15 @@ struct GpuCache {
} }
static GPUC: OnceCell<Mutex<GpuCache>> = OnceCell::new(); static GPUC: OnceCell<Mutex<GpuCache>> = OnceCell::new();
// Static caches for unchanging data
static HOSTNAME: OnceCell<String> = OnceCell::new();
struct NetworkNameCache {
names: Vec<String>,
infos: Vec<NetworkInfo>,
}
static NETWORK_CACHE: OnceCell<Mutex<NetworkNameCache>> = OnceCell::new();
static CPU_VEC: OnceCell<Mutex<Vec<f32>>> = OnceCell::new();
fn cached_temp() -> Option<f32> { fn cached_temp() -> Option<f32> {
if !temp_enabled() { if !temp_enabled() {
return None; return None;
@ -121,9 +130,19 @@ pub async fn collect_fast_metrics(state: &AppState) -> Metrics {
warn!("sysinfo selective refresh panicked: {e:?}"); warn!("sysinfo selective refresh panicked: {e:?}");
} }
let hostname = state.hostname.clone(); // Get or initialize hostname once
let hostname = HOSTNAME.get_or_init(|| state.hostname.clone()).clone();
// Reuse CPU vector to avoid allocation
let cpu_total = sys.global_cpu_usage(); let cpu_total = sys.global_cpu_usage();
let cpu_per_core: Vec<f32> = sys.cpus().iter().map(|c| c.cpu_usage()).collect(); let cpu_per_core = {
let vec_lock = CPU_VEC.get_or_init(|| Mutex::new(Vec::with_capacity(32)));
let mut vec = vec_lock.lock().unwrap();
vec.clear();
vec.extend(sys.cpus().iter().map(|c| c.cpu_usage()));
vec.clone() // Still need to clone but the allocation is reused
};
let mem_total = sys.total_memory(); let mem_total = sys.total_memory();
let mem_used = mem_total.saturating_sub(sys.available_memory()); let mem_used = mem_total.saturating_sub(sys.available_memory());
let swap_total = sys.total_swap(); let swap_total = sys.total_swap();
@ -156,17 +175,38 @@ pub async fn collect_fast_metrics(state: &AppState) -> Metrics {
None None
}; };
// Networks // Networks with reusable name cache
let networks: Vec<NetworkInfo> = { let networks = {
let mut nets = state.networks.lock().await; let mut nets = state.networks.lock().await;
nets.refresh(false); nets.refresh(false);
nets.iter()
.map(|(name, data)| NetworkInfo { // Get or initialize network cache
name: name.to_string(), let cache = NETWORK_CACHE.get_or_init(|| {
Mutex::new(NetworkNameCache {
names: Vec::new(),
infos: Vec::with_capacity(4), // Most systems have few network interfaces
})
});
let mut cache = cache.lock().unwrap();
// Collect current network names
let current_names: Vec<_> = nets.keys().map(|name| name.to_string()).collect();
// Update cached network names if they changed
if cache.names != current_names {
cache.names = current_names;
}
// Reuse NetworkInfo objects
cache.infos.clear();
for (name, data) in nets.iter() {
cache.infos.push(NetworkInfo {
name: name.to_string(), // We'll still clone but avoid Vec reallocation
received: data.total_received(), received: data.total_received(),
transmitted: data.total_transmitted(), transmitted: data.total_transmitted(),
}) });
.collect() }
cache.infos.clone()
}; };
// GPUs: if we already determined none exist, short-circuit (no repeated probing) // GPUs: if we already determined none exist, short-circuit (no repeated probing)

View File

@ -7,13 +7,33 @@ use axum::{
}; };
use flate2::{write::GzEncoder, Compression}; use flate2::{write::GzEncoder, Compression};
use futures_util::StreamExt; use futures_util::StreamExt;
use once_cell::sync::OnceCell;
use std::collections::HashMap; use std::collections::HashMap;
use std::io::Write; use std::io::Write;
use tokio::sync::Mutex;
use crate::metrics::{collect_disks, collect_fast_metrics, collect_processes_all}; use crate::metrics::{collect_disks, collect_fast_metrics, collect_processes_all};
use crate::proto::pb; use crate::proto::pb;
use crate::state::AppState; use crate::state::AppState;
// Compression threshold based on typical payload size
const COMPRESSION_THRESHOLD: usize = 768;
// Reusable buffer for compression to avoid allocations
struct CompressionCache {
processes_vec: Vec<pb::Process>,
}
impl CompressionCache {
fn new() -> Self {
Self {
processes_vec: Vec::with_capacity(512), // Typical process count
}
}
}
static COMPRESSION_CACHE: OnceCell<Mutex<CompressionCache>> = OnceCell::new();
pub async fn ws_handler( pub async fn ws_handler(
ws: WebSocketUpgrade, ws: WebSocketUpgrade,
State(state): State<AppState>, State(state): State<AppState>,
@ -46,39 +66,51 @@ async fn handle_socket(mut socket: WebSocket, state: AppState) {
} }
Message::Text(ref text) if text == "get_processes" => { Message::Text(ref text) if text == "get_processes" => {
let payload = collect_processes_all(&state).await; let payload = collect_processes_all(&state).await;
// Map to protobuf message // Map to protobuf message
let rows: Vec<pb::Process> = payload // Get cached buffers
.top_processes let cache = COMPRESSION_CACHE.get_or_init(|| Mutex::new(CompressionCache::new()));
.into_iter() let mut cache = cache.lock().await;
.map(|p| pb::Process {
// Reuse process vector to build the list
cache.processes_vec.clear();
cache
.processes_vec
.extend(payload.top_processes.into_iter().map(|p| pb::Process {
pid: p.pid, pid: p.pid,
name: p.name, name: p.name,
cpu_usage: p.cpu_usage, cpu_usage: p.cpu_usage,
mem_bytes: p.mem_bytes, mem_bytes: p.mem_bytes,
}) }));
.collect();
let pb = pb::Processes { let pb = pb::Processes {
process_count: payload.process_count as u64, process_count: payload.process_count as u64,
rows, rows: std::mem::take(&mut cache.processes_vec),
}; };
let mut buf = Vec::with_capacity(8 * 1024); let mut buf = Vec::with_capacity(8 * 1024);
if prost::Message::encode(&pb, &mut buf).is_err() { if prost::Message::encode(&pb, &mut buf).is_err() {
let _ = socket.send(Message::Close(None)).await; let _ = socket.send(Message::Close(None)).await;
} else { } else {
// compress if large // compress if large
if buf.len() <= 768 { if buf.len() <= COMPRESSION_THRESHOLD {
let _ = socket.send(Message::Binary(buf)).await; let _ = socket.send(Message::Binary(buf)).await;
} else { } else {
let mut enc = GzEncoder::new(Vec::new(), Compression::fast()); // Create a new encoder for each message to ensure proper gzip headers
if enc.write_all(&buf).is_ok() { let mut encoder =
let bin = enc.finish().unwrap_or(buf); GzEncoder::new(Vec::with_capacity(buf.len()), Compression::fast());
let _ = socket.send(Message::Binary(bin)).await; match encoder.write_all(&buf).and_then(|_| encoder.finish()) {
} else { Ok(compressed) => {
let _ = socket.send(Message::Binary(compressed)).await;
}
Err(_) => {
let _ = socket.send(Message::Binary(buf)).await; let _ = socket.send(Message::Binary(buf)).await;
} }
} }
} }
} }
drop(cache); // Explicit drop to release mutex early
}
Message::Close(_) => break, Message::Close(_) => break,
_ => {} _ => {}
} }
@ -91,7 +123,7 @@ async fn handle_socket(mut socket: WebSocket, state: AppState) {
// Small, cheap gzip for larger payloads; send text for small. // Small, cheap gzip for larger payloads; send text for small.
async fn send_json<T: serde::Serialize>(ws: &mut WebSocket, value: &T) -> Result<(), axum::Error> { async fn send_json<T: serde::Serialize>(ws: &mut WebSocket, value: &T) -> Result<(), axum::Error> {
let json = serde_json::to_string(value).expect("serialize"); let json = serde_json::to_string(value).expect("serialize");
if json.len() <= 768 { if json.len() <= COMPRESSION_THRESHOLD {
return ws.send(Message::Text(json)).await; return ws.send(Message::Text(json)).await;
} }
let mut enc = GzEncoder::new(Vec::new(), Compression::fast()); let mut enc = GzEncoder::new(Vec::new(), Compression::fast());
@ -99,3 +131,83 @@ async fn send_json<T: serde::Serialize>(ws: &mut WebSocket, value: &T) -> Result
let bin = enc.finish().unwrap_or_else(|_| json.into_bytes()); let bin = enc.finish().unwrap_or_else(|_| json.into_bytes());
ws.send(Message::Binary(bin)).await ws.send(Message::Binary(bin)).await
} }
#[cfg(test)]
mod tests {
use super::*;
use prost::Message as ProstMessage;
use sysinfo::System;
#[tokio::test]
async fn test_process_list_not_empty() {
// Initialize system data first to ensure we have processes
let mut sys = System::new_all();
sys.refresh_all();
// Create state and put the refreshed system in it
let state = AppState::new();
{
let mut sys_lock = state.sys.lock().await;
*sys_lock = sys;
}
// Get processes directly using the collection function
let processes = collect_processes_all(&state).await;
// Convert to protobuf message format
let cache = COMPRESSION_CACHE.get_or_init(|| Mutex::new(CompressionCache::new()));
let mut cache = cache.lock().await;
// Reuse process vector to build the list
cache.processes_vec.clear();
cache
.processes_vec
.extend(processes.top_processes.into_iter().map(|p| pb::Process {
pid: p.pid,
name: p.name,
cpu_usage: p.cpu_usage,
mem_bytes: p.mem_bytes,
}));
// Create the protobuf message
let pb = pb::Processes {
process_count: processes.process_count as u64,
rows: cache.processes_vec.clone(),
};
// Test protobuf encoding/decoding
let mut buf = Vec::new();
prost::Message::encode(&pb, &mut buf).expect("Failed to encode protobuf");
let decoded = pb::Processes::decode(buf.as_slice()).expect("Failed to decode protobuf");
// Print debug info
println!("Process count: {}", pb.process_count);
println!("Process vector length: {}", pb.rows.len());
println!("Encoded size: {} bytes", buf.len());
println!("Decoded process count: {}", decoded.rows.len());
// Print first few processes if available
for (i, process) in pb.rows.iter().take(5).enumerate() {
println!(
"Process {}: {} (PID: {}) CPU: {:.1}% MEM: {} bytes",
i + 1,
process.name,
process.pid,
process.cpu_usage,
process.mem_bytes
);
}
// Validate
assert!(!pb.rows.is_empty(), "Process list should not be empty");
assert!(
pb.process_count > 0,
"Process count should be greater than 0"
);
assert_eq!(
pb.process_count as usize,
pb.rows.len(),
"Process count mismatch with actual rows"
);
}
}