use core::alloc::{GlobalAlloc, Layout};
use core::ffi::c_void;
use std::ffi::CString;
use std::path::PathBuf;
pub mod sys;
pub struct MiMalloc;
unsafe impl GlobalAlloc for MiMalloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
sys::mi_malloc_aligned(layout.size(), layout.align()).cast()
}
unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
sys::mi_free(ptr.cast::<c_void>());
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
sys::mi_realloc_aligned(ptr.cast::<c_void>(), new_size, layout.align()).cast()
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
sys::mi_zalloc_aligned(layout.size(), layout.align()).cast()
}
}
pub unsafe fn unwrapped_malloc(size: usize, alignment: usize) -> *mut u8 {
unsafe { sys::mi_unwrapped_malloc(size, alignment).cast() }
}
pub unsafe fn unwrapped_free(p: *mut u8) {
unsafe { sys::mi_unwrapped_free(p.cast()) }
}
pub unsafe fn unwrapped_realloc(p: *mut u8, new_size: usize, alignment: usize) -> *mut u8 {
unsafe { sys::mi_unwrapped_realloc(p.cast(), new_size, alignment).cast() }
}
pub fn enable_heap_profiling() -> bool {
prof::start(0)
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum ProfConfigMode {
#[default]
Fallback,
Override,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum DumpFormat {
#[default]
Text,
Proto,
}
#[non_exhaustive]
#[derive(Debug, Clone, Default)]
pub struct ProfConfig {
pub mode: ProfConfigMode,
pub sample_interval: Option<usize>,
pub max_profiler_bytes: Option<usize>,
pub seed: Option<u64>,
pub accum: bool,
pub max_stack_depth: Option<usize>,
pub dump_at_exit: Option<PathBuf>,
pub dump_format: DumpFormat,
}
pub fn enable_heap_profiling_with(config: &ProfConfig) -> bool {
let dump_at_exit_c: Option<CString> = match &config.dump_at_exit {
Some(path) => match path.to_str().and_then(|s| CString::new(s).ok()) {
Some(c) => Some(c),
None => return false,
},
None => None,
};
let mut raw: sys::mi_prof_config_t = unsafe { core::mem::zeroed() };
raw.size = core::mem::size_of::<sys::mi_prof_config_t>();
raw.version = sys::MI_PROF_CONFIG_VERSION;
raw.mode = match config.mode {
ProfConfigMode::Fallback => sys::MI_PROF_CONFIG_FALLBACK,
ProfConfigMode::Override => sys::MI_PROF_CONFIG_OVERRIDE,
};
raw.sample_interval = config.sample_interval.unwrap_or(0);
raw.max_profiler_bytes = config.max_profiler_bytes.unwrap_or(0);
raw.seed = config.seed.unwrap_or(0);
raw.accum = config.accum;
raw.max_stack_depth = config.max_stack_depth.unwrap_or(0);
raw.dump_at_exit = dump_at_exit_c
.as_ref()
.map_or(core::ptr::null(), |c| c.as_ptr());
raw.dump_format = match config.dump_format {
DumpFormat::Text => sys::MI_PROF_FORMAT_TEXT,
DumpFormat::Proto => sys::MI_PROF_FORMAT_PROTO,
};
unsafe { sys::mi_prof_start_ex(&raw) }
}
pub mod prof {
use core::ffi::{c_char, c_void};
use std::ffi::{CStr, CString};
use std::io;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::path::Path;
use crate::sys;
pub fn start(sample_rate: usize) -> bool {
unsafe { sys::mi_prof_start(sample_rate) }
}
#[doc(hidden)]
pub fn start_seeded(sample_rate: usize, seed: u64) -> bool {
unsafe { sys::mi_prof_start_seeded(sample_rate, seed) }
}
pub fn stop() {
unsafe { sys::mi_prof_stop() }
}
pub fn is_enabled() -> bool {
unsafe { sys::mi_prof_is_enabled() }
}
pub fn reset() {
unsafe { sys::mi_prof_reset() }
}
pub fn dump_file(path: &Path) -> io::Result<()> {
let path = path.to_str().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "profile path is not UTF-8")
})?;
let path = CString::new(path).map_err(|_| {
io::Error::new(io::ErrorKind::InvalidInput, "profile path contains NUL")
})?;
if unsafe { sys::mi_prof_dump(path.as_ptr()) } {
Ok(())
} else {
Err(io::Error::last_os_error())
}
}
unsafe extern "C" fn write_cb(arg: *mut c_void, buf: *const c_char, len: usize) {
let out = &mut *(arg as *mut Vec<u8>);
out.extend_from_slice(core::slice::from_raw_parts(buf.cast::<u8>(), len));
}
pub fn dump_to_vec() -> Vec<u8> {
let mut out = Vec::new();
let ok =
unsafe { sys::mi_prof_dump_writer(Some(write_cb), (&mut out as *mut Vec<u8>).cast()) };
if ok {
out
} else {
Vec::new()
}
}
pub fn dump_proto_to_vec() -> Vec<u8> {
let mut out = Vec::new();
let ok = unsafe {
sys::mi_prof_dump_proto_writer(Some(write_cb), (&mut out as *mut Vec<u8>).cast())
};
if ok {
out
} else {
Vec::new()
}
}
pub fn dump_proto_file(path: &Path) -> io::Result<()> {
let path = path.to_str().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "profile path is not UTF-8")
})?;
let path = CString::new(path).map_err(|_| {
io::Error::new(io::ErrorKind::InvalidInput, "profile path contains NUL")
})?;
if unsafe { sys::mi_prof_dump_proto(path.as_ptr()) } {
Ok(())
} else {
Err(io::Error::last_os_error())
}
}
#[derive(Debug, Clone, Default)]
pub struct ProfStats {
pub enabled: bool,
pub accum: bool,
pub sample_rate: usize,
pub live_samples: usize,
pub live_bytes: usize,
pub accum_samples: usize,
pub accum_bytes: usize,
pub unique_stacks: usize,
pub arena_committed: usize,
pub stack_table_overflows: usize,
pub dropped_samples: usize,
pub heap: HeapStats,
}
#[derive(Debug, Clone, Default)]
pub struct HeapStats {
pub committed: usize,
pub reserved: usize,
pub malloc_requested: usize,
pub pages: usize,
pub pages_abandoned: usize,
pub heaps: usize,
pub theaps: usize,
pub purged: usize,
pub detailed: bool,
}
pub fn stats() -> ProfStats {
let mut raw: sys::mi_prof_stats_t = unsafe { core::mem::zeroed() };
raw.size = core::mem::size_of::<sys::mi_prof_stats_t>();
raw.version = sys::MI_PROF_STAT_VERSION;
if unsafe { sys::mi_prof_stats_get(&mut raw) } {
ProfStats {
enabled: raw.enabled,
accum: raw.accum,
sample_rate: raw.sample_rate,
live_samples: raw.live_samples,
live_bytes: raw.live_bytes,
accum_samples: raw.accum_samples,
accum_bytes: raw.accum_bytes,
unique_stacks: raw.unique_stacks,
arena_committed: raw.arena_committed,
stack_table_overflows: raw.stack_table_overflows,
dropped_samples: raw.dropped_samples,
heap: HeapStats {
committed: raw.heap_committed,
reserved: raw.heap_reserved,
malloc_requested: raw.heap_malloc_requested,
pages: raw.heap_pages,
pages_abandoned: raw.heap_pages_abandoned,
heaps: raw.heap_count,
theaps: raw.theap_count,
purged: raw.heap_purged,
detailed: raw.heap_stats_detailed,
},
}
} else {
ProfStats::default()
}
}
#[derive(Debug, Clone)]
pub struct Sample {
pub stack: Vec<usize>,
pub live_objects: usize,
pub live_bytes: usize,
pub accum_objects: usize,
pub accum_bytes: usize,
}
impl Sample {
pub fn estimated_bytes(&self, sample_rate: usize) -> u64 {
if self.live_objects == 0 || self.live_bytes == 0 {
return 0;
}
if sample_rate <= 1 {
return self.live_bytes as u64;
}
let avg = self.live_bytes as f64 / self.live_objects as f64;
let scale = 1.0 / (1.0 - (-avg / sample_rate as f64).exp());
(self.live_bytes as f64 * scale) as u64
}
}
struct SnapshotGuard(*mut sys::mi_prof_snapshot_t);
impl Drop for SnapshotGuard {
fn drop(&mut self) {
unsafe { sys::mi_prof_snapshot_free(self.0) }
}
}
unsafe extern "C" fn collect_visitor(
info: *const sys::mi_prof_sample_info_t,
arg: *mut c_void,
) -> bool {
let result = catch_unwind(AssertUnwindSafe(|| unsafe {
let out = &mut *(arg as *mut Vec<Sample>);
let info = &*info;
let stack = (0..info.depth)
.map(|i| *info.stack.add(i) as usize)
.collect();
out.push(Sample {
stack,
live_objects: info.live_objects,
live_bytes: info.live_bytes,
accum_objects: info.accum_objects,
accum_bytes: info.accum_bytes,
});
}));
result.is_ok()
}
pub fn samples() -> Vec<Sample> {
let snap = unsafe { sys::mi_prof_snapshot_new() };
if snap.is_null() {
return Vec::new();
}
let guard = SnapshotGuard(snap);
let mut out: Vec<Sample> = Vec::new();
unsafe {
sys::mi_prof_snapshot_visit(
guard.0,
collect_visitor,
(&mut out as *mut Vec<Sample>).cast(),
);
}
out
}
#[derive(Debug, Clone)]
pub struct ModuleInfo {
pub path: String,
pub base: usize,
pub size: usize,
}
unsafe extern "C" fn modules_visitor(
info: *const sys::mi_prof_module_info_t,
arg: *mut c_void,
) -> bool {
let result = catch_unwind(AssertUnwindSafe(|| unsafe {
let out = &mut *(arg as *mut Vec<ModuleInfo>);
let info = &*info;
let path = CStr::from_ptr(info.path).to_string_lossy().into_owned();
out.push(ModuleInfo {
path,
base: info.base,
size: info.size,
});
}));
result.is_ok()
}
pub fn modules() -> Vec<ModuleInfo> {
let mut out: Vec<ModuleInfo> = Vec::new();
unsafe {
sys::mi_prof_modules_visit(
modules_visitor,
(&mut out as *mut Vec<ModuleInfo>).cast(),
);
}
out
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
static PROF_TEST_LOCK: Mutex<()> = Mutex::new(());
fn reset_profiler() {
if prof::is_enabled() {
prof::stop();
}
}
#[test]
fn enable_heap_profiling_with_default_config_starts_profiler() {
let _guard = PROF_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
reset_profiler();
let config = ProfConfig::default();
assert!(enable_heap_profiling_with(&config));
assert!(prof::is_enabled());
prof::stop();
}
#[test]
fn enable_heap_profiling_with_override_mode_sets_sample_interval() {
let _guard = PROF_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
reset_profiler();
let config = ProfConfig {
mode: ProfConfigMode::Override,
sample_interval: Some(4096),
..Default::default()
};
assert!(enable_heap_profiling_with(&config));
assert!(prof::is_enabled());
assert_eq!(prof::stats().sample_rate, 4096);
prof::stop();
}
#[test]
fn unwrapped_malloc_write_realloc_grow_verify_free() {
unsafe {
let size = 64usize;
let p = unwrapped_malloc(size, 0);
assert!(!p.is_null());
for i in 0..size {
*p.add(i) = (i % 256) as u8;
}
let new_size = 256usize;
let p2 = unwrapped_realloc(p, new_size, 0);
assert!(!p2.is_null());
for i in 0..size {
assert_eq!(*p2.add(i), (i % 256) as u8);
}
unwrapped_free(p2);
}
}
#[test]
fn unwrapped_free_null_is_noop() {
unsafe {
unwrapped_free(core::ptr::null_mut());
}
}
#[test]
fn unwrapped_malloc_rejects_non_power_of_two_alignment() {
unsafe {
let p = unwrapped_malloc(16, 3);
assert!(p.is_null());
}
}
#[test]
fn unwrapped_realloc_with_null_ptr_behaves_like_malloc() {
unsafe {
let p = unwrapped_realloc(core::ptr::null_mut(), 32, 0);
assert!(!p.is_null());
unwrapped_free(p);
}
}
#[test]
fn unwrapped_realloc_with_zero_size_frees_and_returns_null() {
unsafe {
let p = unwrapped_malloc(32, 0);
assert!(!p.is_null());
let p2 = unwrapped_realloc(p, 0, 0);
assert!(p2.is_null());
}
}
}