use crate::EnumerationAnomaly;
use crate::observation::Source;
use crate::records::RecordWalk;
use std::io;
use crate::records::SIZING_ATTEMPTS;
use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
use windows_sys::Win32::System::SystemInformation::{
CpuSetInformation, GetSystemCpuSetInformation, SYSTEM_CPU_SET_INFORMATION,
};
use windows_sys::Win32::System::Threading::GetCurrentProcess;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CpuSet {
pub id: u32,
pub group: u16,
pub logical_processor_index: u8,
pub core_index: u8,
pub last_level_cache_index: u8,
pub numa_node_index: u8,
pub efficiency_class: u8,
pub parked: bool,
pub allocated: bool,
pub allocated_to_target_process: bool,
pub real_time: bool,
pub scheduling_class: u8,
pub allocation_tag: u64,
}
mod flags {
pub(super) const PARKED: u8 = 1 << 0;
pub(super) const ALLOCATED: u8 = 1 << 1;
pub(super) const ALLOCATED_TO_TARGET_PROCESS: u8 = 1 << 2;
pub(super) const REAL_TIME: u8 = 1 << 3;
}
pub(crate) fn enumerate() -> io::Result<(Vec<CpuSet>, Option<EnumerationAnomaly>)> {
let mut grew = None;
for _ in 0..SIZING_ATTEMPTS {
let mut length: u32 = 0;
let probe = unsafe {
GetSystemCpuSetInformation(
std::ptr::null_mut(),
0,
&raw mut length,
GetCurrentProcess(),
0,
)
};
if probe != 0 {
return Ok((Vec::new(), None));
}
let error = io::Error::last_os_error();
if error.raw_os_error() != Some(ERROR_INSUFFICIENT_BUFFER as i32) {
return Err(error);
}
if length == 0 {
return Ok((Vec::new(), None));
}
let mut storage = vec![0_u64; (length as usize).div_ceil(8)];
let buffer = storage.as_mut_ptr().cast::<u8>();
let mut actual_length = length;
let ok = unsafe {
GetSystemCpuSetInformation(
buffer.cast(),
length,
&raw mut actual_length,
GetCurrentProcess(),
0,
)
};
if ok != 0 {
return Ok(unsafe { decode(buffer.cast_const(), actual_length) });
}
let error = io::Error::last_os_error();
if error.raw_os_error() != Some(ERROR_INSUFFICIENT_BUFFER as i32) {
return Err(error);
}
grew = Some(error);
}
Err(grew.expect("SIZING_ATTEMPTS is a non-zero constant, so the loop ran at least once"))
}
const SIZE_OFFSET: usize = core::mem::offset_of!(SYSTEM_CPU_SET_INFORMATION, Size);
const TYPE_OFFSET: usize = core::mem::offset_of!(SYSTEM_CPU_SET_INFORMATION, Type);
const UNION_OFFSET: usize = core::mem::offset_of!(SYSTEM_CPU_SET_INFORMATION, Anonymous);
unsafe fn decode(base: *const u8, length: u32) -> (Vec<CpuSet>, Option<EnumerationAnomaly>) {
use windows_sys::Win32::System::SystemInformation::{
SYSTEM_CPU_SET_INFORMATION_0, SYSTEM_CPU_SET_INFORMATION_0_0,
};
const CPUSET_OFFSET: usize = core::mem::offset_of!(SYSTEM_CPU_SET_INFORMATION_0, CpuSet);
macro_rules! field {
($name:ident) => {
UNION_OFFSET
+ CPUSET_OFFSET
+ core::mem::offset_of!(SYSTEM_CPU_SET_INFORMATION_0_0, $name)
};
}
let mut walk = unsafe {
RecordWalk::new(
base,
length,
SIZE_OFFSET,
size_of::<SYSTEM_CPU_SET_INFORMATION>(),
Source::CpuSets,
)
};
let mut records = Vec::new();
for record in &mut walk {
let decoded = unsafe {
(|| {
if record.read::<i32>(TYPE_OFFSET)? != CpuSetInformation {
return None;
}
let all_flags = record.read::<u8>(field!(Anonymous1))?;
Some(CpuSet {
id: record.read(field!(Id))?,
group: record.read(field!(Group))?,
logical_processor_index: record.read(field!(LogicalProcessorIndex))?,
core_index: record.read(field!(CoreIndex))?,
last_level_cache_index: record.read(field!(LastLevelCacheIndex))?,
numa_node_index: record.read(field!(NumaNodeIndex))?,
efficiency_class: record.read(field!(EfficiencyClass))?,
parked: all_flags & flags::PARKED != 0,
allocated: all_flags & flags::ALLOCATED != 0,
allocated_to_target_process: all_flags & flags::ALLOCATED_TO_TARGET_PROCESS
!= 0,
real_time: all_flags & flags::REAL_TIME != 0,
scheduling_class: record.read(field!(Anonymous2))?,
allocation_tag: record.read(field!(AllocationTag))?,
})
})()
};
if let Some(cpu_set) = decoded {
records.push(cpu_set);
}
}
(records, walk.anomaly())
}
#[cfg(test)]
mod tests;