565 lines
18 KiB
Rust
565 lines
18 KiB
Rust
//use chrono::Utc;
|
|
use nvml_wrapper::Nvml;
|
|
use reqwest::{Client, StatusCode};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::{error::Error, fs, process::Command, time::Duration};
|
|
use sysinfo::{CpuExt, DiskExt, System, SystemExt};
|
|
use tokio::time::{interval, sleep};
|
|
|
|
// Data structures matching the C# DTOs
|
|
#[derive(Serialize, Debug)]
|
|
struct RegistrationDto {
|
|
#[serde(rename = "id")]
|
|
id: i32,
|
|
#[serde(rename = "ipAddress")]
|
|
ip_address: String,
|
|
#[serde(rename = "cpuType")]
|
|
cpu_type: String,
|
|
#[serde(rename = "cpuCores")]
|
|
cpu_cores: i32,
|
|
#[serde(rename = "gpuType")]
|
|
gpu_type: String,
|
|
#[serde(rename = "ramSize")]
|
|
ram_size: f64,
|
|
}
|
|
|
|
#[derive(Serialize, Debug)]
|
|
struct MetricDto {
|
|
#[serde(rename = "serverId")]
|
|
server_id: i32,
|
|
#[serde(rename = "ipAddress")]
|
|
ip_address: String,
|
|
#[serde(rename = "cpu_Load")]
|
|
cpu_load: f64,
|
|
#[serde(rename = "cpu_Temp")]
|
|
cpu_temp: f64,
|
|
#[serde(rename = "gpu_Load")]
|
|
gpu_load: f64,
|
|
#[serde(rename = "gpu_Temp")]
|
|
gpu_temp: f64,
|
|
#[serde(rename = "gpu_Vram_Size")]
|
|
gpu_vram_size: f64,
|
|
#[serde(rename = "gpu_Vram_Usage")]
|
|
gpu_vram_usage: f64,
|
|
#[serde(rename = "ram_Load")]
|
|
ram_load: f64,
|
|
#[serde(rename = "ram_Size")]
|
|
ram_size: f64,
|
|
#[serde(rename = "disk_Size")]
|
|
disk_size: f64,
|
|
#[serde(rename = "disk_Usage")]
|
|
disk_usage: f64,
|
|
#[serde(rename = "disk_Temp")]
|
|
disk_temp: f64,
|
|
#[serde(rename = "net_In")]
|
|
net_in: f64,
|
|
#[serde(rename = "net_Out")]
|
|
net_out: f64,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct IdResponse {
|
|
id: i32,
|
|
#[serde(rename = "ipAddress")]
|
|
ip_address: String,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct HeartbeatPayload {
|
|
#[serde(rename = "IpAddress")]
|
|
ip_address: String,
|
|
}
|
|
|
|
struct HardwareInfo {
|
|
cpu_type: String,
|
|
cpu_cores: i32,
|
|
gpu_type: String,
|
|
ram_size: f64,
|
|
ip_address: String,
|
|
}
|
|
|
|
impl HardwareInfo {
|
|
async fn collect() -> Result<Self, Box<dyn Error>> {
|
|
let mut sys = System::new();
|
|
sys.refresh_cpu();
|
|
sys.refresh_memory();
|
|
|
|
let cpus = sys.cpus();
|
|
let cpu_type = cpus
|
|
.get(0)
|
|
.map(|c| c.brand().to_string())
|
|
.unwrap_or("Unknown CPU".to_string());
|
|
let cpu_cores = cpus.len() as i32;
|
|
let ram_gb = (sys.total_memory() as f64) / 1024.0 / 1024.0;
|
|
let gpu_type = Self::detect_gpu_name();
|
|
let ip_address = local_ip_address::local_ip()?.to_string();
|
|
|
|
Ok(Self {
|
|
cpu_type,
|
|
cpu_cores,
|
|
gpu_type,
|
|
ram_size: ram_gb,
|
|
ip_address,
|
|
})
|
|
}
|
|
|
|
fn detect_gpu_name() -> String {
|
|
Self::try_nvml_gpu_name()
|
|
.or_else(Self::fallback_gpu_name)
|
|
.unwrap_or_else(|| "Unknown GPU".to_string())
|
|
}
|
|
|
|
fn try_nvml_gpu_name() -> Option<String> {
|
|
let nvml = Nvml::init().ok()?;
|
|
let device = nvml.device_by_index(0).ok()?;
|
|
device.name().ok().map(|s| s.to_string())
|
|
}
|
|
|
|
fn fallback_gpu_name() -> Option<String> {
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
let output = std::process::Command::new("lshw")
|
|
.args(&["-C", "display"])
|
|
.output()
|
|
.ok()?;
|
|
Some(
|
|
String::from_utf8_lossy(&output.stdout)
|
|
.lines()
|
|
.find(|l| l.contains("product:"))
|
|
.map(|l| l.trim().replace("product:", "").trim().to_string())
|
|
.unwrap_or("Unknown GPU".to_string()),
|
|
)
|
|
}
|
|
|
|
#[cfg(target_os = "windows")]
|
|
{
|
|
let output = std::process::Command::new("wmic")
|
|
.args(&["path", "win32_VideoController", "get", "name"])
|
|
.output()
|
|
.ok()?;
|
|
Some(
|
|
String::from_utf8_lossy(&output.stdout)
|
|
.lines()
|
|
.nth(1)
|
|
.map(|s| s.trim().to_string())
|
|
.unwrap_or("Unknown GPU".to_string()),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn get_server_id_by_ip(base_url: &str, ip: &str) -> Result<(i32, String), Box<dyn Error>> {
|
|
let client = Client::builder()
|
|
.danger_accept_invalid_certs(true)
|
|
.build()?;
|
|
|
|
let url = format!("{}/monitoring/server-id-by-ip?ipAddress={}", base_url, ip);
|
|
|
|
loop {
|
|
println!("Attempting to fetch server ID for IP {}...", ip);
|
|
match client.get(&url).send().await {
|
|
Ok(resp) if resp.status().is_success() => {
|
|
let text = resp.text().await?;
|
|
println!("Raw response: {}", text); // Debug output
|
|
|
|
let id_resp: IdResponse = serde_json::from_str(&text).map_err(|e| {
|
|
println!("Failed to parse response: {}", e);
|
|
e
|
|
})?;
|
|
|
|
println!(
|
|
"✅ Received ID {} for IP {}",
|
|
id_resp.id, id_resp.ip_address
|
|
);
|
|
return Ok((id_resp.id, id_resp.ip_address));
|
|
}
|
|
Ok(resp) if resp.status() == StatusCode::NOT_FOUND => {
|
|
println!(
|
|
"❌ Server with IP {} not found in database (will retry in 10 seconds)",
|
|
ip
|
|
);
|
|
sleep(Duration::from_secs(10)).await;
|
|
}
|
|
Ok(resp) => {
|
|
println!(
|
|
"⚠️ Server responded with status: {} - {}",
|
|
resp.status(),
|
|
resp.text().await?
|
|
);
|
|
sleep(Duration::from_secs(10)).await;
|
|
}
|
|
Err(err) => {
|
|
println!("⚠️ Request failed: {} (will retry in 10 seconds)", err);
|
|
sleep(Duration::from_secs(10)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn register_with_server(base_url: &str) -> Result<(i32, String), Box<dyn Error>> {
|
|
// First get local IP
|
|
let ip = local_ip_address::local_ip()?.to_string();
|
|
println!("Local IP address detected: {}", ip);
|
|
|
|
// Get server ID from backend (this will retry until successful)
|
|
let (server_id, registered_ip) = get_server_id_by_ip(base_url, &ip).await?;
|
|
|
|
// Create HTTP client for registration
|
|
let client = Client::builder()
|
|
.danger_accept_invalid_certs(true)
|
|
.build()?;
|
|
|
|
// Collect hardware info
|
|
let hardware = HardwareInfo::collect().await?;
|
|
|
|
// Prepare registration data
|
|
let registration = RegistrationDto {
|
|
id: server_id,
|
|
ip_address: registered_ip.clone(),
|
|
cpu_type: hardware.cpu_type,
|
|
cpu_cores: hardware.cpu_cores,
|
|
gpu_type: hardware.gpu_type,
|
|
ram_size: hardware.ram_size,
|
|
};
|
|
|
|
// Try to register (will retry on failure)
|
|
loop {
|
|
println!("Attempting to register with server...");
|
|
let url = format!("{}/monitoring/register-agent-by-id", base_url);
|
|
match client.post(&url).json(®istration).send().await {
|
|
Ok(resp) if resp.status().is_success() => {
|
|
println!("✅ Successfully registered with server.");
|
|
return Ok((server_id, registered_ip));
|
|
}
|
|
Ok(resp) => {
|
|
let status = resp.status();
|
|
let text = resp.text().await.unwrap_or_default();
|
|
println!(
|
|
"⚠️ Registration failed ({}): {} (will retry in 10 seconds)",
|
|
status, text
|
|
);
|
|
}
|
|
Err(err) => {
|
|
println!("⚠️ Registration error: {} (will retry in 10 seconds)", err);
|
|
}
|
|
}
|
|
sleep(Duration::from_secs(10)).await;
|
|
}
|
|
}
|
|
|
|
async fn heartbeat_loop(base_url: &str, ip: &str) -> Result<(), Box<dyn Error>> {
|
|
let client = Client::builder()
|
|
.danger_accept_invalid_certs(true)
|
|
.build()?;
|
|
let url = format!("{}/heartbeat/receive", base_url);
|
|
|
|
loop {
|
|
let payload = HeartbeatPayload {
|
|
ip_address: ip.to_string(),
|
|
};
|
|
|
|
match client.post(&url).json(&payload).send().await {
|
|
Ok(res) if res.status().is_success() => {
|
|
println!("Heartbeat sent successfully.");
|
|
}
|
|
Ok(res) => eprintln!("Server responded with status: {}", res.status()),
|
|
Err(e) => eprintln!("Heartbeat error: {}", e),
|
|
}
|
|
|
|
sleep(Duration::from_secs(20)).await;
|
|
}
|
|
}
|
|
|
|
struct MetricsCollector {
|
|
sys: System,
|
|
nvml: Option<Nvml>,
|
|
server_id: i32,
|
|
ip_address: String,
|
|
}
|
|
|
|
impl MetricsCollector {
|
|
fn new(server_id: i32, ip_address: String) -> Self {
|
|
Self {
|
|
sys: System::new(),
|
|
nvml: Nvml::init().ok(),
|
|
server_id,
|
|
ip_address,
|
|
}
|
|
}
|
|
|
|
async fn collect_and_send_loop(&mut self, base_url: &str) -> Result<(), Box<dyn Error>> {
|
|
let client = Client::new();
|
|
let url = format!("{}/monitoring/metric", base_url);
|
|
let mut interval = interval(Duration::from_secs(20));
|
|
|
|
loop {
|
|
interval.tick().await;
|
|
let metric = self.collect_metrics();
|
|
println!("Collected metrics: {:?}", metric);
|
|
|
|
match client.post(&url).json(&metric).send().await {
|
|
Ok(res) => println!(
|
|
"✅ Sent metrics for server {} | Status: {}",
|
|
metric.server_id,
|
|
res.status()
|
|
),
|
|
Err(err) => eprintln!("❌ Failed to send metrics: {}", err),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn collect_metrics(&mut self) -> MetricDto {
|
|
self.sys.refresh_all();
|
|
|
|
// CPU
|
|
let cpu_load = self.sys.global_cpu_info().cpu_usage() as f64;
|
|
let cpu_temp = get_cpu_temp().unwrap_or(0.0) as f64;
|
|
|
|
// RAM
|
|
let total_memory = self.sys.total_memory();
|
|
let used_memory = self.sys.used_memory();
|
|
let ram_load = (used_memory as f64 / total_memory as f64) * 100.0;
|
|
let ram_size = (total_memory as f64) / 1024.0 / 1024.0;
|
|
|
|
// Disk
|
|
let disk = self.sys.disks().first();
|
|
// In collect_metrics():
|
|
let (disk_size, disk_usage, disk_temp) = {
|
|
let mut total_size = 0u64;
|
|
let mut total_used = 0u64;
|
|
let mut temp = 0.0;
|
|
let mut count = 0;
|
|
|
|
for disk in self.sys.disks() {
|
|
total_size += disk.total_space();
|
|
total_used += disk.total_space() - disk.available_space();
|
|
count += 1;
|
|
}
|
|
|
|
// Disk temperature (Linux only)
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
if let Ok(dir) = fs::read_dir("/sys/block") {
|
|
for entry in dir.flatten() {
|
|
if let Some(disk_name) = entry.file_name().to_str() {
|
|
if disk_name.starts_with("sd") || disk_name.starts_with("nvme") {
|
|
let temp_path = format!(
|
|
"/sys/block/{}/device/hwmon/hwmon*/temp1_input",
|
|
disk_name
|
|
);
|
|
if let Ok(paths) = glob::glob(&temp_path) {
|
|
for path in paths.flatten() {
|
|
if let Ok(content) = fs::read_to_string(path) {
|
|
if let Ok(t) = content.trim().parse::<f32>() {
|
|
temp += t / 1000.0; // Convert millidegrees
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let size_gb = if count > 0 {
|
|
(total_size as f64) / 1024.0 / 1024.0 / 1024.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let usage = if total_size > 0 {
|
|
(total_used as f64 / total_size as f64) * 100.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let avg_temp = if count > 0 { temp / count as f64 } else { 0.0 };
|
|
|
|
(size_gb, usage, avg_temp)
|
|
};
|
|
// GPU (NVIDIA)
|
|
let (gpu_temp, gpu_load, vram_used, vram_total) = if let Some(nvml) = &self.nvml {
|
|
if let Ok(device) = nvml.device_by_index(0) {
|
|
let temp = device
|
|
.temperature(nvml_wrapper::enum_wrappers::device::TemperatureSensor::Gpu)
|
|
.unwrap_or(0) as f64;
|
|
let load = device
|
|
.utilization_rates()
|
|
.map(|u| u.gpu as f64)
|
|
.unwrap_or(0.0);
|
|
let mem = device.memory_info().ok();
|
|
let used = mem
|
|
.clone()
|
|
.map(|m| (m.used as f64) / 1024.0 / 1024.0 / 1024.0)
|
|
.unwrap_or(0.0); // GB
|
|
let total = mem
|
|
.map(|m| (m.total as f64) / 1024.0 / 1024.0 / 1024.0)
|
|
.unwrap_or(0.0); // GB
|
|
(temp, load, used, total)
|
|
} else {
|
|
(0.0, 0.0, 0.0, 0.0)
|
|
}
|
|
} else {
|
|
(0.0, 0.0, 0.0, 0.0)
|
|
};
|
|
|
|
// Network (convert bytes to bits)
|
|
let (net_in, net_out) = get_network_traffic().unwrap_or((0, 0));
|
|
let net_in_bits = (net_in as f64) * 8.0;
|
|
let net_out_bits = (net_out as f64) * 8.0;
|
|
|
|
MetricDto {
|
|
server_id: self.server_id,
|
|
ip_address: self.ip_address.clone(),
|
|
cpu_load,
|
|
cpu_temp,
|
|
gpu_load,
|
|
gpu_temp,
|
|
gpu_vram_size: vram_total,
|
|
gpu_vram_usage: if vram_total > 0.0 {
|
|
(vram_used / vram_total) * 100.0
|
|
} else {
|
|
0.0
|
|
},
|
|
ram_load,
|
|
ram_size,
|
|
disk_size,
|
|
disk_usage: disk_usage,
|
|
disk_temp: 0.0, // not supported
|
|
net_in: net_in_bits,
|
|
net_out: net_out_bits,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn get_cpu_temp() -> Option<f32> {
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
// Linux: sensors command or sysfs
|
|
if let Ok(output) = Command::new("sensors").output() {
|
|
let stdout = String::from_utf8_lossy(&output.stdout);
|
|
for line in stdout.lines() {
|
|
if line.to_lowercase().contains("package id")
|
|
|| line.to_lowercase().contains("cpu temp")
|
|
{
|
|
if let Some(temp_str) = line.split_whitespace().find(|s| s.contains("°C")) {
|
|
let number: String = temp_str
|
|
.chars()
|
|
.filter(|c| c.is_ascii_digit() || *c == '.')
|
|
.collect();
|
|
return number.parse::<f32>().ok();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fallback to sysfs (common path for Intel/AMD)
|
|
if let Ok(content) = fs::read_to_string("/sys/class/thermal/thermal_zone0/temp") {
|
|
if let Ok(temp) = content.trim().parse::<f32>() {
|
|
return Some(temp / 1000.0); // Convert millidegrees to degrees
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(target_os = "windows")]
|
|
{
|
|
// Windows: WMI query
|
|
let output = Command::new("wmic")
|
|
.args(&["cpu", "get", "Temperature", "/Value"])
|
|
.output()
|
|
.ok()?;
|
|
|
|
let stdout = String::from_utf8_lossy(&output.stdout);
|
|
for line in stdout.lines() {
|
|
if line.starts_with("Temperature=") {
|
|
if let Ok(temp) = line.replace("Temperature=", "").trim().parse::<f32>() {
|
|
return Some(temp); // Returns in Celsius
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn get_network_traffic() -> Option<(u64, u64)> {
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
let content = fs::read_to_string("/proc/net/dev").ok()?;
|
|
let mut rx_total = 0u64;
|
|
let mut tx_total = 0u64;
|
|
|
|
for line in content.lines().skip(2) {
|
|
let parts: Vec<&str> = line.split_whitespace().collect();
|
|
if parts.len() < 10 || parts[0].ends_with(":") {
|
|
continue;
|
|
}
|
|
if parts[0].contains("lo:") {
|
|
continue;
|
|
}
|
|
rx_total += parts[1].parse::<u64>().unwrap_or(0);
|
|
tx_total += parts[9].parse::<u64>().unwrap_or(0);
|
|
}
|
|
Some((rx_total, tx_total))
|
|
}
|
|
|
|
#[cfg(target_os = "windows")]
|
|
{
|
|
use std::process::Stdio;
|
|
let output = Command::new("netstat")
|
|
.args(&["-e"])
|
|
.stdout(Stdio::piped())
|
|
.output()
|
|
.ok()?;
|
|
|
|
let stdout = String::from_utf8_lossy(&output.stdout);
|
|
let mut lines = stdout.lines();
|
|
|
|
// Find the line with statistics
|
|
while let Some(line) = lines.next() {
|
|
if line.contains("Bytes") {
|
|
if let Some(stats_line) = lines.next() {
|
|
let parts: Vec<&str> = stats_line.split_whitespace().collect();
|
|
if parts.len() >= 2 {
|
|
let rx = parts[0].parse::<u64>().unwrap_or(0);
|
|
let tx = parts[1].parse::<u64>().unwrap_or(0);
|
|
return Some((rx, tx));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn Error>> {
|
|
let server_base_url = "http://localhost:5000";
|
|
|
|
// Registration phase
|
|
println!("Starting registration process...");
|
|
let (server_id, ip_address) = register_with_server(server_base_url).await?;
|
|
|
|
// Start heartbeat in background
|
|
let heartbeat_handle = tokio::spawn({
|
|
let ip = ip_address.clone();
|
|
async move {
|
|
if let Err(e) = heartbeat_loop(server_base_url, &ip).await {
|
|
eprintln!("Heartbeat loop failed: {}", e);
|
|
}
|
|
}
|
|
});
|
|
|
|
// Start metrics collection
|
|
println!("Starting metrics collection...");
|
|
let mut metrics_collector = MetricsCollector::new(server_id, ip_address);
|
|
metrics_collector
|
|
.collect_and_send_loop(server_base_url)
|
|
.await?;
|
|
|
|
heartbeat_handle.await?;
|
|
Ok(())
|
|
}
|