vendor/sysinfo/src/windows/cpu.rs

   1 // Take a look at the license at the top of the repository in the LICENSE file.
   2
   3 use crate::sys::tools::KeyHandler;
   4 use crate::{CpuExt, CpuRefreshKind, LoadAvg};
   5
   6 use std::collections::HashMap;
   7 use std::io::Error;
   8 use std::mem;
   9 use std::ops::DerefMut;
  10 use std::ptr::null_mut;
  11 use std::sync::Mutex;
  12
  13 use ntapi::ntpoapi::PROCESSOR_POWER_INFORMATION;
  14
  15 use winapi::shared::minwindef::FALSE;
  16 use winapi::shared::winerror::{ERROR_INSUFFICIENT_BUFFER, ERROR_SUCCESS};
  17 use winapi::um::handleapi::CloseHandle;
  18 use winapi::um::pdh::{
  19     PdhAddEnglishCounterA, PdhAddEnglishCounterW, PdhCloseQuery, PdhCollectQueryData,
  20     PdhCollectQueryDataEx, PdhGetFormattedCounterValue, PdhOpenQueryA, PdhRemoveCounter,
  21     PDH_FMT_COUNTERVALUE, PDH_FMT_DOUBLE, PDH_HCOUNTER, PDH_HQUERY,
  22 };
  23 use winapi::um::powerbase::CallNtPowerInformation;
  24 use winapi::um::synchapi::CreateEventA;
  25 use winapi::um::sysinfoapi::GetLogicalProcessorInformationEx;
  26 use winapi::um::sysinfoapi::SYSTEM_INFO;
  27 use winapi::um::winbase::{RegisterWaitForSingleObject, INFINITE};
  28 use winapi::um::winnt::{
  29     ProcessorInformation, RelationAll, RelationProcessorCore, BOOLEAN, HANDLE,
  30     PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX, PVOID, WT_EXECUTEDEFAULT,
  31 };
  32
  33 // This formula comes from Linux's include/linux/sched/loadavg.h
  34 // https://github.com/torvalds/linux/blob/345671ea0f9258f410eb057b9ced9cefbbe5dc78/include/linux/sched/loadavg.h#L20-L23
  35 #[allow(clippy::excessive_precision)]
  36 const LOADAVG_FACTOR_1F: f64 = 0.9200444146293232478931553241;
  37 #[allow(clippy::excessive_precision)]
  38 const LOADAVG_FACTOR_5F: f64 = 0.9834714538216174894737477501;
  39 #[allow(clippy::excessive_precision)]
  40 const LOADAVG_FACTOR_15F: f64 = 0.9944598480048967508795473394;
  41 // The time interval in seconds between taking load counts, same as Linux
  42 const SAMPLING_INTERVAL: usize = 5;
  43
  44 // maybe use a read/write lock instead?
  45 static LOAD_AVG: once_cell::sync::Lazy<Mutex<Option<LoadAvg>>> =
  46     once_cell::sync::Lazy::new(|| unsafe { init_load_avg() });
  47
  48 pub(crate) fn get_load_average() -> LoadAvg {
  49     if let Ok(avg) = LOAD_AVG.lock() {
  50         if let Some(avg) = &*avg {
  51             return avg.clone();
  52         }
  53     }
  54     LoadAvg::default()
  55 }
  56
  57 unsafe extern "system" fn load_avg_callback(counter: PVOID, _: BOOLEAN) {
  58     let mut display_value = mem::MaybeUninit::<PDH_FMT_COUNTERVALUE>::uninit();
  59
  60     if PdhGetFormattedCounterValue(
  61         counter as _,
  62         PDH_FMT_DOUBLE,
  63         null_mut(),
  64         display_value.as_mut_ptr(),
  65     ) != ERROR_SUCCESS as _
  66     {
  67         return;
  68     }
  69     let display_value = display_value.assume_init();
  70     if let Ok(mut avg) = LOAD_AVG.lock() {
  71         if let Some(avg) = avg.deref_mut() {
  72             let current_load = display_value.u.doubleValue();
  73
  74             avg.one = avg.one * LOADAVG_FACTOR_1F + current_load * (1.0 - LOADAVG_FACTOR_1F);
  75             avg.five = avg.five * LOADAVG_FACTOR_5F + current_load * (1.0 - LOADAVG_FACTOR_5F);
  76             avg.fifteen =
  77                 avg.fifteen * LOADAVG_FACTOR_15F + current_load * (1.0 - LOADAVG_FACTOR_15F);
  78         }
  79     }
  80 }
  81
  82 unsafe fn init_load_avg() -> Mutex<Option<LoadAvg>> {
  83     // You can see the original implementation here: https://github.com/giampaolo/psutil
  84     let mut query = null_mut();
  85
  86     if PdhOpenQueryA(null_mut(), 0, &mut query) != ERROR_SUCCESS as _ {
  87         sysinfo_debug!("init_load_avg: PdhOpenQueryA failed");
  88         return Mutex::new(None);
  89     }
  90
  91     let mut counter: PDH_HCOUNTER = mem::zeroed();
  92     if PdhAddEnglishCounterA(
  93         query,
  94         b"\\System\\Cpu Queue Length\0".as_ptr() as _,
  95         0,
  96         &mut counter,
  97     ) != ERROR_SUCCESS as _
  98     {
  99         PdhCloseQuery(query);
 100         sysinfo_debug!("init_load_avg: failed to get CPU queue length");
 101         return Mutex::new(None);
 102     }
 103
 104     let event = CreateEventA(null_mut(), FALSE, FALSE, b"LoadUpdateEvent\0".as_ptr() as _);
 105     if event.is_null() {
 106         PdhCloseQuery(query);
 107         sysinfo_debug!("init_load_avg: failed to create event `LoadUpdateEvent`");
 108         return Mutex::new(None);
 109     }
 110
 111     if PdhCollectQueryDataEx(query, SAMPLING_INTERVAL as _, event) != ERROR_SUCCESS as _ {
 112         PdhCloseQuery(query);
 113         sysinfo_debug!("init_load_avg: PdhCollectQueryDataEx failed");
 114         return Mutex::new(None);
 115     }
 116
 117     let mut wait_handle = null_mut();
 118     if RegisterWaitForSingleObject(
 119         &mut wait_handle,
 120         event,
 121         Some(load_avg_callback),
 122         counter as _,
 123         INFINITE,
 124         WT_EXECUTEDEFAULT,
 125     ) == 0
 126     {
 127         PdhRemoveCounter(counter);
 128         PdhCloseQuery(query);
 129         sysinfo_debug!("init_load_avg: RegisterWaitForSingleObject failed");
 130         Mutex::new(None)
 131     } else {
 132         Mutex::new(Some(LoadAvg::default()))
 133     }
 134 }
 135
 136 struct InternalQuery {
 137     query: PDH_HQUERY,
 138     event: HANDLE,
 139     data: HashMap<String, PDH_HCOUNTER>,
 140 }
 141
 142 unsafe impl Send for InternalQuery {}
 143 unsafe impl Sync for InternalQuery {}
 144
 145 impl Drop for InternalQuery {
 146     fn drop(&mut self) {
 147         unsafe {
 148             for (_, counter) in self.data.iter() {
 149                 PdhRemoveCounter(*counter);
 150             }
 151
 152             if !self.event.is_null() {
 153                 CloseHandle(self.event);
 154             }
 155
 156             if !self.query.is_null() {
 157                 PdhCloseQuery(self.query);
 158             }
 159         }
 160     }
 161 }
 162
 163 pub(crate) struct Query {
 164     internal: InternalQuery,
 165 }
 166
 167 impl Query {
 168     pub fn new() -> Option<Query> {
 169         let mut query = null_mut();
 170         unsafe {
 171             if PdhOpenQueryA(null_mut(), 0, &mut query) == ERROR_SUCCESS as i32 {
 172                 let q = InternalQuery {
 173                     query,
 174                     event: null_mut(),
 175                     data: HashMap::new(),
 176                 };
 177                 Some(Query { internal: q })
 178             } else {
 179                 sysinfo_debug!("Query::new: PdhOpenQueryA failed");
 180                 None
 181             }
 182         }
 183     }
 184
 185     #[allow(clippy::ptr_arg)]
 186     pub fn get(&self, name: &String) -> Option<f32> {
 187         if let Some(counter) = self.internal.data.get(name) {
 188             unsafe {
 189                 let mut display_value = mem::MaybeUninit::<PDH_FMT_COUNTERVALUE>::uninit();
 190                 let counter: PDH_HCOUNTER = *counter;
 191
 192                 let ret = PdhGetFormattedCounterValue(
 193                     counter,
 194                     PDH_FMT_DOUBLE,
 195                     null_mut(),
 196                     display_value.as_mut_ptr(),
 197                 ) as u32;
 198                 let display_value = display_value.assume_init();
 199                 return if ret == ERROR_SUCCESS as _ {
 200                     let data = *display_value.u.doubleValue();
 201                     Some(data as f32)
 202                 } else {
 203                     sysinfo_debug!("Query::get: PdhGetFormattedCounterValue failed");
 204                     Some(0.)
 205                 };
 206             }
 207         }
 208         None
 209     }
 210
 211     #[allow(clippy::ptr_arg)]
 212     pub fn add_english_counter(&mut self, name: &String, getter: Vec<u16>) -> bool {
 213         if self.internal.data.contains_key(name) {
 214             sysinfo_debug!("Query::add_english_counter: doesn't have key `{:?}`", name);
 215             return false;
 216         }
 217         unsafe {
 218             let mut counter: PDH_HCOUNTER = std::mem::zeroed();
 219             let ret = PdhAddEnglishCounterW(self.internal.query, getter.as_ptr(), 0, &mut counter);
 220             if ret == ERROR_SUCCESS as _ {
 221                 self.internal.data.insert(name.clone(), counter);
 222             } else {
 223                 sysinfo_debug!(
 224                     "Query::add_english_counter: failed to add counter '{}': {:x}...",
 225                     name,
 226                     ret,
 227                 );
 228                 return false;
 229             }
 230         }
 231         true
 232     }
 233
 234     pub fn refresh(&self) {
 235         unsafe {
 236             if PdhCollectQueryData(self.internal.query) != ERROR_SUCCESS as _ {
 237                 sysinfo_debug!("failed to refresh CPU data");
 238             }
 239         }
 240     }
 241 }
 242
 243 pub(crate) struct CpusWrapper {
 244     global: Cpu,
 245     cpus: Vec<Cpu>,
 246     got_cpu_frequency: bool,
 247 }
 248
 249 impl CpusWrapper {
 250     pub fn new() -> Self {
 251         Self {
 252             global: Cpu::new_with_values("Total CPU".to_owned(), String::new(), String::new(), 0),
 253             cpus: Vec::new(),
 254             got_cpu_frequency: false,
 255         }
 256     }
 257
 258     pub fn global_cpu(&self) -> &Cpu {
 259         &self.global
 260     }
 261
 262     pub fn global_cpu_mut(&mut self) -> &mut Cpu {
 263         &mut self.global
 264     }
 265
 266     pub fn cpus(&self) -> &[Cpu] {
 267         &self.cpus
 268     }
 269
 270     fn init_if_needed(&mut self, refresh_kind: CpuRefreshKind) {
 271         if self.cpus.is_empty() {
 272             let (cpus, vendor_id, brand) = super::tools::init_cpus(refresh_kind);
 273             self.cpus = cpus;
 274             self.global.vendor_id = vendor_id;
 275             self.global.brand = brand;
 276             self.got_cpu_frequency = refresh_kind.frequency();
 277         }
 278     }
 279
 280     pub fn len(&mut self) -> usize {
 281         self.init_if_needed(CpuRefreshKind::new());
 282         self.cpus.len()
 283     }
 284
 285     pub fn iter_mut(&mut self, refresh_kind: CpuRefreshKind) -> impl Iterator<Item = &mut Cpu> {
 286         self.init_if_needed(refresh_kind);
 287         self.cpus.iter_mut()
 288     }
 289
 290     pub fn get_frequencies(&mut self) {
 291         if self.got_cpu_frequency {
 292             return;
 293         }
 294         let frequencies = get_frequencies(self.cpus.len());
 295
 296         for (cpu, frequency) in self.cpus.iter_mut().zip(frequencies) {
 297             cpu.set_frequency(frequency);
 298         }
 299         self.global
 300             .set_frequency(self.cpus.get(0).map(|cpu| cpu.frequency()).unwrap_or(0));
 301         self.got_cpu_frequency = true;
 302     }
 303 }
 304
 305 #[doc = include_str!("../../md_doc/cpu.md")]
 306 pub struct Cpu {
 307     name: String,
 308     cpu_usage: f32,
 309     key_used: Option<KeyHandler>,
 310     vendor_id: String,
 311     brand: String,
 312     frequency: u64,
 313 }
 314
 315 impl CpuExt for Cpu {
 316     fn cpu_usage(&self) -> f32 {
 317         self.cpu_usage
 318     }
 319
 320     fn name(&self) -> &str {
 321         &self.name
 322     }
 323
 324     fn frequency(&self) -> u64 {
 325         self.frequency
 326     }
 327
 328     fn vendor_id(&self) -> &str {
 329         &self.vendor_id
 330     }
 331
 332     fn brand(&self) -> &str {
 333         &self.brand
 334     }
 335 }
 336
 337 impl Cpu {
 338     pub(crate) fn new_with_values(
 339         name: String,
 340         vendor_id: String,
 341         brand: String,
 342         frequency: u64,
 343     ) -> Cpu {
 344         Cpu {
 345             name,
 346             cpu_usage: 0f32,
 347             key_used: None,
 348             vendor_id,
 349             brand,
 350             frequency,
 351         }
 352     }
 353
 354     pub(crate) fn set_cpu_usage(&mut self, value: f32) {
 355         self.cpu_usage = value;
 356     }
 357
 358     pub(crate) fn set_frequency(&mut self, value: u64) {
 359         self.frequency = value;
 360     }
 361 }
 362
 363 fn get_vendor_id_not_great(info: &SYSTEM_INFO) -> String {
 364     use winapi::um::winnt;
 365     // https://docs.microsoft.com/fr-fr/windows/win32/api/sysinfoapi/ns-sysinfoapi-system_info
 366     unsafe {
 367         match info.u.s().wProcessorArchitecture {
 368             winnt::PROCESSOR_ARCHITECTURE_INTEL => "Intel x86",
 369             winnt::PROCESSOR_ARCHITECTURE_MIPS => "MIPS",
 370             winnt::PROCESSOR_ARCHITECTURE_ALPHA => "RISC Alpha",
 371             winnt::PROCESSOR_ARCHITECTURE_PPC => "PPC",
 372             winnt::PROCESSOR_ARCHITECTURE_SHX => "SHX",
 373             winnt::PROCESSOR_ARCHITECTURE_ARM => "ARM",
 374             winnt::PROCESSOR_ARCHITECTURE_IA64 => "Intel Itanium-based x64",
 375             winnt::PROCESSOR_ARCHITECTURE_ALPHA64 => "RISC Alpha x64",
 376             winnt::PROCESSOR_ARCHITECTURE_MSIL => "MSIL",
 377             winnt::PROCESSOR_ARCHITECTURE_AMD64 => "(Intel or AMD) x64",
 378             winnt::PROCESSOR_ARCHITECTURE_IA32_ON_WIN64 => "Intel Itanium-based x86",
 379             winnt::PROCESSOR_ARCHITECTURE_NEUTRAL => "unknown",
 380             winnt::PROCESSOR_ARCHITECTURE_ARM64 => "ARM x64",
 381             winnt::PROCESSOR_ARCHITECTURE_ARM32_ON_WIN64 => "ARM",
 382             winnt::PROCESSOR_ARCHITECTURE_IA32_ON_ARM64 => "Intel Itanium-based x86",
 383             _ => "unknown",
 384         }
 385         .to_owned()
 386     }
 387 }
 388
 389 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
 390 pub(crate) fn get_vendor_id_and_brand(info: &SYSTEM_INFO) -> (String, String) {
 391     #[cfg(target_arch = "x86")]
 392     use std::arch::x86::__cpuid;
 393     #[cfg(target_arch = "x86_64")]
 394     use std::arch::x86_64::__cpuid;
 395
 396     unsafe fn add_u32(v: &mut Vec<u8>, i: u32) {
 397         let i = &i as *const u32 as *const u8;
 398         v.push(*i);
 399         v.push(*i.offset(1));
 400         v.push(*i.offset(2));
 401         v.push(*i.offset(3));
 402     }
 403
 404     unsafe {
 405         // First, we try to get the complete name.
 406         let res = __cpuid(0x80000000);
 407         let n_ex_ids = res.eax;
 408         let brand = if n_ex_ids >= 0x80000004 {
 409             let mut extdata = Vec::with_capacity(5);
 410
 411             for i in 0x80000000..=n_ex_ids {
 412                 extdata.push(__cpuid(i));
 413             }
 414
 415             // 4 * u32 * nb_entries
 416             let mut out = Vec::with_capacity(4 * std::mem::size_of::<u32>() * 3);
 417             for data in extdata.iter().take(5).skip(2) {
 418                 add_u32(&mut out, data.eax);
 419                 add_u32(&mut out, data.ebx);
 420                 add_u32(&mut out, data.ecx);
 421                 add_u32(&mut out, data.edx);
 422             }
 423             let mut pos = 0;
 424             for e in out.iter() {
 425                 if *e == 0 {
 426                     break;
 427                 }
 428                 pos += 1;
 429             }
 430             match std::str::from_utf8(&out[..pos]) {
 431                 Ok(s) => s.to_owned(),
 432                 _ => String::new(),
 433             }
 434         } else {
 435             String::new()
 436         };
 437
 438         // Failed to get full name, let's retry for the short version!
 439         let res = __cpuid(0);
 440         let mut x = Vec::with_capacity(3 * std::mem::size_of::<u32>());
 441         add_u32(&mut x, res.ebx);
 442         add_u32(&mut x, res.edx);
 443         add_u32(&mut x, res.ecx);
 444         let mut pos = 0;
 445         for e in x.iter() {
 446             if *e == 0 {
 447                 break;
 448             }
 449             pos += 1;
 450         }
 451         let vendor_id = match std::str::from_utf8(&x[..pos]) {
 452             Ok(s) => s.to_owned(),
 453             Err(_) => get_vendor_id_not_great(info),
 454         };
 455         (vendor_id, brand)
 456     }
 457 }
 458
 459 #[cfg(all(not(target_arch = "x86_64"), not(target_arch = "x86")))]
 460 pub(crate) fn get_vendor_id_and_brand(info: &SYSTEM_INFO) -> (String, String) {
 461     (get_vendor_id_not_great(info), String::new())
 462 }
 463
 464 pub(crate) fn get_key_used(p: &mut Cpu) -> &mut Option<KeyHandler> {
 465     &mut p.key_used
 466 }
 467
 468 // From https://stackoverflow.com/a/43813138:
 469 //
 470 // If your PC has 64 or fewer logical cpus installed, the above code will work fine. However,
 471 // if your PC has more than 64 logical cpus installed, use GetActiveCpuCount() or
 472 // GetLogicalCpuInformation() to determine the total number of logical cpus installed.
 473 pub(crate) fn get_frequencies(nb_cpus: usize) -> Vec<u64> {
 474     let size = nb_cpus * mem::size_of::<PROCESSOR_POWER_INFORMATION>();
 475     let mut infos: Vec<PROCESSOR_POWER_INFORMATION> = Vec::with_capacity(nb_cpus);
 476
 477     unsafe {
 478         if CallNtPowerInformation(
 479             ProcessorInformation,
 480             null_mut(),
 481             0,
 482             infos.as_mut_ptr() as _,
 483             size as _,
 484         ) == 0
 485         {
 486             infos.set_len(nb_cpus);
 487             // infos.Number
 488             return infos
 489                 .into_iter()
 490                 .map(|i| i.CurrentMhz as u64)
 491                 .collect::<Vec<_>>();
 492         }
 493     }
 494     sysinfo_debug!("get_frequencies: CallNtPowerInformation failed");
 495     vec![0; nb_cpus]
 496 }
 497
 498 pub(crate) fn get_physical_core_count() -> Option<usize> {
 499     // we cannot use the number of cpus here to pre calculate the buf size
 500     // GetLogicalCpuInformationEx with RelationProcessorCore passed to it not only returns
 501     // the logical cores but also numa nodes
 502     //
 503     // GetLogicalCpuInformationEx: https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getlogicalprocessorinformationex
 504
 505     let mut needed_size = 0;
 506     unsafe {
 507         GetLogicalProcessorInformationEx(RelationAll, null_mut(), &mut needed_size);
 508
 509         let mut buf: Vec<u8> = Vec::with_capacity(needed_size as _);
 510
 511         loop {
 512             if GetLogicalProcessorInformationEx(
 513                 RelationAll,
 514                 buf.as_mut_ptr() as *mut _,
 515                 &mut needed_size,
 516             ) == FALSE
 517             {
 518                 let e = Error::last_os_error();
 519                 // For some reasons, the function might return a size not big enough...
 520                 match e.raw_os_error() {
 521                     Some(value) if value == ERROR_INSUFFICIENT_BUFFER as _ => {}
 522                     _ => {
 523                         sysinfo_debug!(
 524                             "get_physical_core_count: GetLogicalCpuInformationEx failed"
 525                         );
 526                         return None;
 527                     }
 528                 }
 529             } else {
 530                 break;
 531             }
 532             buf.reserve(needed_size as usize - buf.capacity());
 533         }
 534
 535         buf.set_len(needed_size as _);
 536
 537         let mut i = 0;
 538         let raw_buf = buf.as_ptr();
 539         let mut count = 0;
 540         while i < buf.len() {
 541             let p = &*(raw_buf.add(i) as PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX);
 542             i += p.Size as usize;
 543             if p.Relationship == RelationProcessorCore {
 544                 // Only count the physical cores.
 545                 count += 1;
 546             }
 547         }
 548         Some(count)
 549     }
 550 }