heim_cpu/sys/windows/wrappers/
count.rs

1use std::io;
2use std::ptr;
3
4use winapi::shared::{minwindef, winerror};
5use winapi::um::{sysinfoapi, winnt};
6
7/// This struct contains information about logical processors
8/// received from the `GetLogicalProcessorInformationEx` function.
9///
10/// Major problem is that each "Processor" struct in it
11/// has the variable size (described by the `Size` field),
12/// so we are storing data as a plain bytes vector
13/// and unsafely iterating on it later.
14#[derive(Debug)]
15pub struct LogicalProcessors {
16    buffer: Vec<u8>,
17}
18
19impl LogicalProcessors {
20    // TODO: Ensure that allowing this lint is fine.
21    // It is triggered by casting `*mut u8` to the logical processor info struct pointer,
22    // which is `#[repr(C)]` and generally should be fine.
23    #[allow(clippy::cast_ptr_alignment)]
24    pub fn get() -> io::Result<Self> {
25        let mut buffer = vec![];
26        let mut buffer_size = 0u32;
27
28        let result = unsafe {
29            sysinfoapi::GetLogicalProcessorInformationEx(
30                winnt::RelationAll,
31                ptr::null_mut(),
32                &mut buffer_size,
33            )
34        };
35        debug_assert_eq!(result, minwindef::FALSE);
36
37        loop {
38            // Allocating enough memory to fill the buffer now
39            buffer.reserve(buffer_size as usize - buffer.capacity());
40
41            let result = unsafe {
42                sysinfoapi::GetLogicalProcessorInformationEx(
43                    winnt::RelationAll,
44                    buffer.as_mut_ptr() as *mut _,
45                    &mut buffer_size,
46                )
47            };
48
49            if result == minwindef::FALSE {
50                let e = io::Error::last_os_error();
51                match e.raw_os_error() {
52                    // Slight chance that there is now more CPU cores
53                    // and we need more memory?
54                    Some(value) if value == winerror::ERROR_INSUFFICIENT_BUFFER as i32 => continue,
55                    _ => return Err(e),
56                }
57            } else {
58                unsafe {
59                    buffer.set_len(buffer_size as usize);
60                }
61                break;
62            }
63        }
64
65        Ok(Self { buffer })
66    }
67
68    pub fn iter(&self) -> LogicalProcessorsIter<'_> {
69        LogicalProcessorsIter {
70            data: &self.buffer,
71            offset: 0,
72        }
73    }
74}
75
76#[derive(Debug)]
77pub struct LogicalProcessorsIter<'p> {
78    data: &'p [u8],
79    offset: usize,
80}
81
82impl<'p> Iterator for LogicalProcessorsIter<'p> {
83    type Item = &'p winnt::SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX;
84
85    #[allow(clippy::cast_ptr_alignment)]
86    fn next(&mut self) -> Option<Self::Item> {
87        if self.offset >= self.data.len() {
88            return None;
89        }
90
91        let core = unsafe {
92            let ptr = self.data.as_ptr().add(self.offset)
93                as winnt::PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX;
94            self.offset += (*ptr).Size as usize;
95            &*ptr
96        };
97        Some(core)
98    }
99}
heim_cpu/sys/windows/wrappers/count.rs

heim_cpu/sys/windows/wrappers/
count.rs
