windows stuff

src/win/cpu.rs

@@ -0,0 +1,59 @@
+use crate::data::CpuTime;
+use crate::{Error, Result, SensorSource};
+use winapi::shared::minwindef;
+use winapi::um::{processthreadsapi, winbase, winnt};
+
+pub struct CpuTimeSource {
+    cpu_count: f32,
+}
+
+impl CpuTimeSource {
+    pub fn new() -> Result<CpuTimeSource> {
+        Ok(CpuTimeSource {
+            cpu_count: dbg!(cpu_count()?),
+        })
+    }
+}
+
+impl SensorSource for CpuTimeSource {
+    type Data = CpuTime;
+
+    fn read(&mut self) -> Result<Self::Data> {
+        let mut user = minwindef::FILETIME::default();
+        let mut kernel = minwindef::FILETIME::default();
+        let mut idle = minwindef::FILETIME::default();
+
+        let result =
+            unsafe { processthreadsapi::GetSystemTimes(&mut idle, &mut kernel, &mut user) };
+
+        if result == 0 {
+            Err(Error::last_os_error("GetSystemTimes"))
+        } else {
+            let user = time_to_float(user);
+            let idle = time_to_float(idle);
+            // Same as `psutil` subtracting idle time
+            // and leaving only busy kernel time
+            let system = time_to_float(kernel) - idle;
+
+            Ok(CpuTime((user + system) / self.cpu_count))
+        }
+    }
+}
+
+fn time_to_float(time: minwindef::FILETIME) -> f32 {
+    const HI_T: f64 = 429.496_729_6;
+    const LO_T: f64 = 1e-7;
+
+    let low = LO_T * f64::from(time.dwLowDateTime);
+    HI_T.mul_add(f64::from(time.dwHighDateTime), low) as f32
+}
+
+fn cpu_count() -> Result<f32> {
+    let result = unsafe { winbase::GetActiveProcessorCount(winnt::ALL_PROCESSOR_GROUPS) };
+
+    if result > 0 {
+        Ok(result as f32)
+    } else {
+        Err(Error::last_os_error("GetActiveProcessorCount"))
+    }
+}

src/win/mod.rs

@@ -0,0 +1,101 @@
+mod cpu;
+mod disk;
+mod reg;
+mod wmi;
+
+use self::cpu::CpuTimeSource;
+use crate::data::{DiskUsage, GpuMemory, GpuUsage, Memory, NetStats};
+use crate::win::wmi::WmiSensor;
+use crate::Result;
+use crate::{hostname, SensorData, SensorSource};
+use once_cell::sync::Lazy;
+use os_thread_local::ThreadLocal;
+use std::borrow::Cow;
+use std::sync::Mutex;
+use sysinfo::{ComponentExt, DiskExt, NetworkExt, System, SystemExt};
+
+pub struct Sensors {
+    pub hostname: String,
+    pub system: Mutex<System>,
+    cpu: Mutex<CpuTimeSource>,
+    gpu_mem_total: u64,
+}
+
+static WMI: Lazy<ThreadLocal<WmiSensor>> =
+    Lazy::new(|| ThreadLocal::new(|| WmiSensor::new().expect("failed to init wmi")));
+
+impl Sensors {
+    pub fn new() -> Result<Sensors> {
+        let mut system = System::new_all();
+        system.refresh_all();
+        println!("{:?}", system);
+        for component in system.components() {
+            println!("{} :{}°C", component.label(), component.temperature());
+        }
+
+        let gpu_mem_total = reg::total_gpu_memory()?;
+
+        Ok(Sensors {
+            hostname: hostname()?,
+            system: Mutex::new(system),
+            cpu: Mutex::new(CpuTimeSource::new()?),
+            gpu_mem_total,
+        })
+    }
+}
+
+pub fn get_metrics(sensors: &Sensors) -> Result<String> {
+    let mut system = sensors.system.lock().unwrap();
+    system.refresh_disks();
+    system.refresh_networks();
+    system.refresh_memory();
+
+    let hostname = &sensors.hostname;
+    let mut result = String::with_capacity(256);
+
+    let memory = Memory {
+        total: system.total_memory(),
+        available: system.available_memory(),
+        free: system.free_memory(),
+    };
+    memory.write(&mut result, &hostname);
+    for disk in system.disks() {
+        let space = DiskUsage {
+            name: disk.name().to_string_lossy().into(),
+            size: disk.total_space(),
+            free: disk.available_space(),
+        };
+        space.write(&mut result, &hostname);
+    }
+    for (interface, net) in system.networks() {
+        let usage = NetStats {
+            interface: interface.into(),
+            bytes_received: net.total_received(),
+            bytes_sent: net.total_transmitted(),
+        };
+        usage.write(&mut result, &hostname);
+    }
+    let cpu = sensors.cpu.lock().unwrap().read()?;
+    cpu.write(&mut result, &hostname);
+
+    let gpu_mem_used = WMI.with(|wmi| wmi.gpu_mem())?;
+    let gpu_mem = GpuMemory {
+        total: sensors.gpu_mem_total,
+        free: sensors.gpu_mem_total - gpu_mem_used,
+    };
+    gpu_mem.write(&mut result, &hostname);
+
+    let gpu_engines = WMI.with(|wmi| wmi.gpu_usage())?;
+    for (name, usage) in gpu_engines.into_iter() {
+        let gpu_usage = GpuUsage {
+            system: Cow::Owned(name),
+            usage,
+        };
+        gpu_usage.write(&mut result, &hostname);
+    }
+    if let Some(disk_usage) = WMI.with(|wmi| wmi.disk_usage())? {
+        disk_usage.write(&mut result, &hostname);
+    }
+
+    Ok(result)
+}

src/win/reg.rs

@@ -0,0 +1,24 @@
+use crate::{Error, Result};
+use serde::Deserialize;
+use winreg::enums::*;
+use winreg::RegKey;
+
+#[derive(Debug, Deserialize)]
+struct GpuInfo {
+    #[serde(rename = "HardwareInformation.qwMemorySize")]
+    memory_size: Option<u64>,
+}
+
+pub fn total_gpu_memory() -> Result<u64> {
+    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+    let mut mem = 0;
+    for i in 0..3 {
+        if let Ok(gpu_key) = hklm.open_subkey(
+            format!("SYSTEM\\ControlSet001\\Control\\Class\\{{4d36e968-e325-11ce-bfc1-08002be10318}}\\{:04}", i),
+        ) {
+            let info: GpuInfo = gpu_key.decode().map_err(|e| Error::Reg(e.to_string()))?;
+            mem += info.memory_size.unwrap_or_default();
+        }
+    }
+    Ok(mem)
+}

src/win/wmi.rs

@@ -0,0 +1,82 @@
+use crate::data::DiskStats;
+use crate::Result;
+use serde::Deserialize;
+use std::collections::HashMap;
+use wmi::{COMLibrary, WMIConnection};
+
+pub struct WmiSensor {
+    wmi_con: WMIConnection,
+}
+
+impl WmiSensor {
+    pub fn new() -> Result<Self> {
+        let com_con = COMLibrary::new()?;
+        let wmi_con = WMIConnection::new(com_con.into())?;
+
+        Ok(WmiSensor { wmi_con })
+    }
+
+    pub fn gpu_mem(&self) -> Result<u64> {
+        #[derive(Deserialize, Debug)]
+        #[allow(non_camel_case_types)]
+        struct Win32_PerfFormattedData_GPUPerformanceCounters_GPUAdapterMemory {
+            #[serde(rename = "DedicatedUsage")]
+            dedicated_usage: u64,
+        }
+
+        let results: Vec<Win32_PerfFormattedData_GPUPerformanceCounters_GPUAdapterMemory> =
+            self.wmi_con.query()?;
+        Ok(results.iter().map(|result| result.dedicated_usage).sum())
+    }
+
+    pub fn gpu_usage(&self) -> Result<HashMap<String, u32>> {
+        #[derive(Deserialize, Debug)]
+        #[allow(non_camel_case_types)]
+        struct Win32_PerfFormattedData_GPUPerformanceCounters_GPUEngine {
+            #[serde(rename = "Name")]
+            name: String,
+            #[serde(rename = "UtilizationPercentage")]
+            usage: u32,
+        }
+
+        let results: Vec<Win32_PerfFormattedData_GPUPerformanceCounters_GPUEngine> =
+            self.wmi_con.query()?;
+
+        let mut data = HashMap::default();
+
+        for result in results {
+            if let Some(eng_type) = result.name.split("_engtype_").skip(1).next() {
+                let entry = data.entry(eng_type.to_string()).or_default();
+                *entry += result.usage;
+            }
+        }
+
+        Ok(data)
+    }
+
+    pub fn disk_usage(&self) -> Result<Option<DiskStats>> {
+        #[derive(Deserialize, Debug)]
+        #[allow(non_camel_case_types)]
+        struct Win32_PerfRawData_Counters_FileSystemDiskActivity {
+            #[serde(rename = "Name")]
+            name: String,
+            #[serde(rename = "FileSystemBytesRead")]
+            read: u64,
+            #[serde(rename = "FileSystemBytesWritten")]
+            written: u64,
+        }
+
+        let results: Vec<Win32_PerfRawData_Counters_FileSystemDiskActivity> =
+            self.wmi_con.query()?;
+        for result in results {
+            if result.name == "_Total" {
+                return Ok(Some(DiskStats {
+                    interface: "Total".to_string(),
+                    bytes_sent: result.written,
+                    bytes_received: result.read,
+                }));
+            }
+        }
+        Ok(None)
+    }
+}