finish macOS bindings for I/O

ClementTsang · ClementTsang · Apr 10, 2023 · Mar 9, 2023 · Mar 11, 2023 · Mar 14, 2023
commit ad1d8f9953d5e51ddea061cc3b7a2b9ea42aeb42
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -110,6 +110,7 @@ libc = "0.2.141"
 procfs = { version = "0.15.1", default-features = false }
 
 [target.'cfg(target_os = "macos")'.dependencies]
+core-foundation = "0.9.3"
 mach2 = "0.4.1"
 
 [target.'cfg(target_os = "windows")'.dependencies]

diff --git a/src/app/data_harvester/disks/system/linux/io_counters.rs b/src/app/data_harvester/disks/system/linux/io_counters.rs
@@ -25,14 +25,8 @@ const DISK_SECTOR_SIZE: u64 = 512;
 #[derive(Debug, Default)]
 pub struct IoCounters {
     name: String,
-    read_count: u64,
-    read_merged_count: u64,
     read_bytes: u64,
-    read_time_secs: u64,
-    write_count: u64,
-    write_merged_count: u64,
     write_bytes: u64,
-    write_time_secs: u64,
 }
 
 impl IoCounters {
@@ -71,26 +65,20 @@ impl FromStr for IoCounters {
 
         let name = next_part(&mut parts)?.to_string();
 
-        let read_count = next_part(&mut parts)?.parse()?;
-        let read_merged_count = next_part(&mut parts)?.parse()?;
+        let _read_count = next_part(&mut parts)?.parse()?;
+        let _read_merged_count = next_part(&mut parts)?.parse()?;
         let read_bytes = next_part(&mut parts)?.parse::<u64>()? * DISK_SECTOR_SIZE;
-        let read_time_secs = next_part(&mut parts)?.parse()?;
+        let _read_time_secs = next_part(&mut parts)?.parse()?;
 
-        let write_count = next_part(&mut parts)?.parse()?;
-        let write_merged_count = next_part(&mut parts)?.parse()?;
+        let _write_count = next_part(&mut parts)?.parse()?;
+        let _write_merged_count = next_part(&mut parts)?.parse()?;
         let write_bytes = next_part(&mut parts)?.parse::<u64>()? * DISK_SECTOR_SIZE;
-        let write_time_secs = next_part(&mut parts)?.parse()?;
+        let _write_time_secs = next_part(&mut parts)?.parse()?;
 
         Ok(IoCounters {
             name,
-            read_count,
-            read_merged_count,
             read_bytes,
-            read_time_secs,
-            write_count,
-            write_merged_count,
             write_bytes,
-            write_time_secs,
         })
     }
 }

diff --git a/src/app/data_harvester/disks/system/macos/io_counters.rs b/src/app/data_harvester/disks/system/macos/io_counters.rs
@@ -2,16 +2,15 @@
 
 use std::ffi::OsStr;
 
+use anyhow::bail;
+
+use super::io_kit::{self, get_dict, get_disks, get_i64, get_string};
+
 #[derive(Debug, Default)]
 pub struct IoCounters {
     name: String,
-
-    read_count: u64,
     read_bytes: u64,
-    read_time_secs: u64,
-    write_count: u64,
     write_bytes: u64,
-    write_time_secs: u64,
 }
 
 impl IoCounters {
@@ -28,7 +27,33 @@ impl IoCounters {
     }
 }
 
+fn get_device_io(device: io_kit::IoObject) -> anyhow::Result<IoCounters> {
+    let parent = device.service_parent()?;
+
+    if parent.conforms_to_block_storage_driver() {
+        let disk_props = device.properties()?;
+        let parent_props = parent.properties()?;
+
+        let name = get_string(&disk_props, "BSD Name")?;
+        let stats = get_dict(&parent_props, "Statistics")?;
+
+        let read_bytes = get_i64(&stats, "Bytes (Read)")? as u64;
+        let write_bytes = get_i64(&stats, "Bytes (Write)")? as u64;
+
+        // let read_count = stats.get_i64("Operations (Read)")? as u64;
+        // let write_count = stats.get_i64("Operations (Write)")? as u64;
+
+        Ok(IoCounters {
+            name,
+            read_bytes,
+            write_bytes,
+        })
+    } else {
+        bail!("{parent:?}, the parent of {device:?} does not conform to IOBlockStorageDriver")
+    }
+}
+
 /// Returns an iterator of disk I/O stats. Pulls data through IOKit.
 pub fn io_stats() -> anyhow::Result<Vec<anyhow::Result<IoCounters>>> {
-    Ok(vec![])
+    Ok(get_disks()?.map(get_device_io).collect())
 }
diff --git a/src/app/data_harvester/disks/system/macos/io_kit.rs b/src/app/data_harvester/disks/system/macos/io_kit.rs
@@ -0,0 +1,10 @@
+mod bindings;
+
+mod io_iterator;
+pub use io_iterator::*;
+
+mod io_object;
+pub use io_object::*;
+
+mod io_disks;
+pub use io_disks::get_disks;
diff --git a/src/app/data_harvester/disks/system/macos/io_kit/bindings.rs b/src/app/data_harvester/disks/system/macos/io_kit/bindings.rs
@@ -0,0 +1,59 @@
+//! C FFI bindings for [IOKit](https://developer.apple.com/documentation/iokit/).
+//!
+//! Based on [heim](https://github.com/heim-rs/heim/blob/master/heim-common/src/sys/macos/iokit/io_master_port.rs)
+//! and [sysinfo's implementation](https://github.com/GuillaumeGomez/sysinfo/blob/master/src/apple/macos/ffi.rs).
+//!
+//! Ideally, we can remove this if sysinfo ever gains disk I/O capabilities.
+
+use core_foundation::base::{mach_port_t, CFAllocatorRef};
+use core_foundation::dictionary::CFMutableDictionaryRef;
+
+use libc::c_char;
+use mach2::port::MACH_PORT_NULL;
+use mach2::{boolean::boolean_t, kern_return::kern_return_t};
+
+#[allow(non_camel_case_types)]
+pub type io_object_t = mach_port_t;
+
+#[allow(non_camel_case_types)]
+pub type io_iterator_t = io_object_t;
+#[allow(non_camel_case_types)]
+pub type io_registry_entry_t = io_object_t;
+
+pub type IOOptionBits = u32;
+
+/// See https://github.com/1kc/librazermacos/pull/27#issuecomment-1042368531.
+#[allow(non_upper_case_globals)]
+pub const kIOMasterPortDefault: mach_port_t = MACH_PORT_NULL;
+
+#[allow(non_upper_case_globals)]
+pub const kIOServicePlane: &str = "IOService\0";
+
+#[allow(non_upper_case_globals)]
+pub const kIOMediaClass: &str = "IOMedia\0";
+
+// See [here](https://developer.apple.com/documentation/iokit) for more details.
+extern "C" {
+
+    pub fn IOServiceGetMatchingServices(
+        mainPort: mach_port_t, matching: CFMutableDictionaryRef, existing: *mut io_iterator_t,
+    ) -> kern_return_t;
+
+    pub fn IOServiceMatching(name: *const c_char) -> CFMutableDictionaryRef;
+
+    pub fn IOIteratorNext(iterator: io_iterator_t) -> io_object_t;
+
+    pub fn IOObjectRelease(obj: io_object_t) -> kern_return_t;
+
+    pub fn IORegistryEntryGetParentEntry(
+        entry: io_registry_entry_t, plane: *const libc::c_char, parent: *mut io_registry_entry_t,
+    ) -> kern_return_t;
+
+    pub fn IOObjectConformsTo(object: io_object_t, className: *const libc::c_char) -> boolean_t;
+
+    pub fn IORegistryEntryCreateCFProperties(
+        entry: io_registry_entry_t, properties: *mut CFMutableDictionaryRef,
+        allocator: CFAllocatorRef, options: IOOptionBits,
+    ) -> kern_return_t;
+
+}
diff --git a/src/app/data_harvester/disks/system/macos/io_kit/io_disks.rs b/src/app/data_harvester/disks/system/macos/io_kit/io_disks.rs
@@ -0,0 +1,24 @@
+use anyhow::bail;
+
+use mach2::kern_return;
+
+use super::{bindings::*, IoIterator};
+
+pub fn get_disks() -> anyhow::Result<IoIterator> {
+    let mut media_iter: io_iterator_t = 0;
+
+    // SAFETY: This is a safe syscall,
+    let result = unsafe {
+        IOServiceGetMatchingServices(
+            kIOMasterPortDefault,
+            IOServiceMatching(kIOMediaClass.as_ptr().cast()),
+            &mut media_iter,
+        )
+    };
+
+    if result == kern_return::KERN_SUCCESS {
+        Ok(media_iter.into())
+    } else {
+        bail!("IOServiceGetMatchingServices failed, error code {result}");
+    }
+}
diff --git a/src/app/data_harvester/disks/system/macos/io_kit/io_iterator.rs b/src/app/data_harvester/disks/system/macos/io_kit/io_iterator.rs
@@ -0,0 +1,51 @@
+//! Based on [heim's](https://github.com/heim-rs/heim/blob/master/heim-common/src/sys/macos/iokit/io_iterator.rs).
+//! implementation.
+
+use std::ops::{Deref, DerefMut};
+
+use mach2::kern_return;
+
+use super::{bindings::*, io_object::IoObject};
+
+/// Safe wrapper around the IOKit `io_iterator_t` type.
+#[derive(Debug)]
+pub struct IoIterator(io_iterator_t);
+
+impl From<io_iterator_t> for IoIterator {
+    fn from(iter: io_iterator_t) -> IoIterator {
+        IoIterator(iter)
+    }
+}
+
+impl Deref for IoIterator {
+    type Target = io_iterator_t;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for IoIterator {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl Iterator for IoIterator {
+    type Item = IoObject;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Basically, we just stop when we hit 0.
+        match unsafe { IOIteratorNext(self.0) } {
+            0 => None,
+            io_object => Some(IoObject::from(io_object)),
+        }
+    }
+}
+
+impl Drop for IoIterator {
+    fn drop(&mut self) {
+        let result = unsafe { IOObjectRelease(self.0) };
+        assert_eq!(result, kern_return::KERN_SUCCESS);
+    }
+}