Making a decent wrapper for logical processors logic for Windows

heim-cpu/src/sys/windows/count.rs

@@ -1,11 +1,6 @@
-use std::alloc;
-use std::io;
-use std::mem;
-use std::ptr;
-
-use winapi::shared::minwindef;
-use winapi::um::{sysinfoapi, winbase, winnt};
+use winapi::um::{winbase, winnt};
 
+use super::wrappers::count::LogicalProcessors;
 use heim_common::prelude::*;
 
 pub fn logical_count() -> impl Future<Output = Result<u64>> {
@@ -24,59 +19,19 @@ pub fn logical_count() -> impl Future<Output = Result<u64>> {
 // around the `GetLogicalProcessorInformationEx` with `Iterator` interface
 // and proper deallocation on `Drop`
 pub fn physical_count() -> impl Future<Output = Result<Option<u64>>> {
-    let mut buffer_size = 0;
-
-    let result = unsafe {
-        sysinfoapi::GetLogicalProcessorInformationEx(
-            winnt::RelationAll,
-            ptr::null_mut(),
-            &mut buffer_size,
-        )
-    };
-    debug_assert_eq!(result, minwindef::FALSE);
-
-    let layout =
-        alloc::Layout::from_size_align(buffer_size as usize, mem::align_of::<u8>()).unwrap();
-
-    let buffer = unsafe { alloc::alloc(layout) };
-
-    let result = unsafe {
-        sysinfoapi::GetLogicalProcessorInformationEx(
-            winnt::RelationAll,
-            buffer as *mut _,
-            &mut buffer_size,
-        )
-    };
-
-    if result == minwindef::FALSE {
-        unsafe {
-            alloc::dealloc(buffer, layout);
-        }
-
-        future::err(io::Error::last_os_error().into())
-    } else {
-        let mut offset: usize = 0;
-        let mut count: u64 = 0;
-        loop {
-            if offset >= buffer_size as usize {
-                break;
-            }
-
-            unsafe {
-                let processor = buffer
-                    .add(offset)
-                    .cast::<winnt::SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>();
-                offset += (*processor).Size as usize;
-                if (*processor).Relationship == winnt::RelationProcessorCore {
-                    count += 1;
-                }
+    match LogicalProcessors::get() {
+        Ok(processors) => {
+            let count = processors
+                .iter()
+                .filter(|p| p.Relationship == winnt::RelationProcessorCore)
+                .count();
+
+            if count > 0 {
+                future::ok(Some(count as u64))
+            } else {
+                future::ok(None)
             }
         }
-
-        unsafe {
-            alloc::dealloc(buffer, layout);
-        }
-
-        future::ok(Some(count))
+        Err(e) => future::err(e.into()),
     }
 }

heim-cpu/src/sys/windows/mod.rs

@@ -4,6 +4,7 @@ mod stats;
 mod times;
 
 mod bindings;
+mod wrappers;
 
 pub use self::count::*;
 pub use self::freq::*;

heim-cpu/src/sys/windows/wrappers/count.rs

@@ -0,0 +1,94 @@
+use std::io;
+use std::ptr;
+
+use winapi::shared::{minwindef, winerror};
+use winapi::um::{sysinfoapi, winnt};
+
+/// This struct contains information about logical processors
+/// received from the `GetLogicalProcessorInformationEx` function.
+///
+/// Major problem is that each "Processor" struct in it
+/// has the variable size (described by the `Size` field),
+/// so we are storing data as a plain bytes vector
+/// and unsafely iterating on it later.
+#[derive(Debug)]
+pub struct LogicalProcessors {
+    buffer: Vec<u8>,
+}
+
+impl LogicalProcessors {
+    pub fn get() -> io::Result<Self> {
+        let mut buffer = vec![];
+        let mut buffer_size = 0u32;
+
+        let result = unsafe {
+            sysinfoapi::GetLogicalProcessorInformationEx(
+                winnt::RelationAll,
+                ptr::null_mut(),
+                &mut buffer_size,
+            )
+        };
+        debug_assert_eq!(result, minwindef::FALSE);
+
+        loop {
+            // Allocating enough memory to fill the buffer now
+            buffer.reserve(buffer_size as usize - buffer.capacity());
+
+            let result = unsafe {
+                sysinfoapi::GetLogicalProcessorInformationEx(
+                    winnt::RelationAll,
+                    buffer.as_mut_ptr() as *mut _,
+                    &mut buffer_size,
+                )
+            };
+
+            if result == minwindef::FALSE {
+                let e = io::Error::last_os_error();
+                match e.raw_os_error() {
+                    // Slight chance that there is now more CPU cores
+                    // and we need more memory?
+                    Some(value) if value == winerror::ERROR_INSUFFICIENT_BUFFER as i32 => continue,
+                    _ => return Err(e),
+                }
+            } else {
+                unsafe {
+                    buffer.set_len(buffer_size as usize);
+                }
+                break;
+            }
+        }
+
+        Ok(Self { buffer })
+    }
+
+    pub fn iter(&self) -> LogicalProcessorsIter<'_> {
+        LogicalProcessorsIter {
+            data: &self.buffer,
+            offset: 0,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct LogicalProcessorsIter<'p> {
+    data: &'p [u8],
+    offset: usize,
+}
+
+impl<'p> Iterator for LogicalProcessorsIter<'p> {
+    type Item = &'p winnt::SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.offset >= self.data.len() {
+            return None;
+        }
+
+        let core = unsafe {
+            let ptr = self.data.as_ptr().add(self.offset)
+                as winnt::PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX;
+            self.offset += (*ptr).Size as usize;
+            &*ptr
+        };
+        Some(core)
+    }
+}

heim-cpu/src/sys/windows/wrappers/mod.rs

@@ -0,0 +1 @@
+pub mod count;