use libc::{c_char, c_void};
use std::any::Any;
use std::ffi::CStr;
use std::io::{self, Write};
use std::panic::{self, AssertUnwindSafe};
#[derive(Copy, Clone, Default)]
pub struct Options {
pub json_format: bool,
pub skip_constants: bool,
pub skip_merged_arenas: bool,
pub skip_per_arena: bool,
pub skip_bin_size_classes: bool,
pub skip_large_size_classes: bool,
pub skip_mutex_statistics: bool,
_p: (),
}
struct State<W> {
writer: W,
error: io::Result<()>,
panic: Result<(), Box<dyn Any + Send>>,
}
extern "C" fn callback<W>(opaque: *mut c_void, buf: *const c_char)
where
W: Write,
{
unsafe {
let state = &mut *(opaque as *mut State<W>);
if state.error.is_err() || state.panic.is_err() {
return;
}
let buf = CStr::from_ptr(buf);
match panic::catch_unwind(AssertUnwindSafe(|| {
state.writer.write_all(buf.to_bytes())
})) {
Ok(Ok(_)) => {}
Ok(Err(e)) => state.error = Err(e),
Err(e) => state.panic = Err(e),
}
}
}
#[cfg_attr(feature = "cargo-clippy", allow(clippy::cast_possible_wrap))]
pub fn stats_print<W>(writer: W, options: Options) -> io::Result<()>
where
W: Write,
{
unsafe {
let mut state = State {
writer,
error: Ok(()),
panic: Ok(()),
};
let mut opts = [0; 8];
let mut i = 0;
if options.json_format {
opts[i] = b'J' as c_char;
i += 1;
}
if options.skip_constants {
opts[i] = b'g' as c_char;
i += 1;
}
if options.skip_merged_arenas {
opts[i] = b'm' as c_char;
i += 1;
}
if options.skip_per_arena {
opts[i] = b'a' as c_char;
i += 1;
}
if options.skip_bin_size_classes {
opts[i] = b'b' as c_char;
i += 1;
}
if options.skip_large_size_classes {
opts[i] = b'l' as c_char;
i += 1;
}
if options.skip_mutex_statistics {
opts[i] = b'x' as c_char;
i += 1;
}
opts[i] = 0;
tikv_jemalloc_sys::malloc_stats_print(
Some(callback::<W>),
&mut state as *mut _ as *mut c_void,
opts.as_ptr(),
);
if let Err(e) = state.panic {
panic::resume_unwind(e);
}
state.error
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn basic() {
let mut buf = vec![];
stats_print(&mut buf, Options::default()).unwrap();
println!("{}", String::from_utf8(buf).unwrap());
}
#[test]
fn all_options() {
let mut buf = vec![];
let options = Options {
json_format: true,
skip_constants: true,
skip_merged_arenas: true,
skip_per_arena: true,
skip_bin_size_classes: true,
skip_large_size_classes: true,
skip_mutex_statistics: true,
_p: (),
};
stats_print(&mut buf, options).unwrap();
println!("{}", String::from_utf8(buf).unwrap());
}
}