tikv_jemalloc_ctl/
stats_print.rs

1//! Bulk statistics output.
2
3use libc::{c_char, c_void};
4use std::any::Any;
5use std::ffi::CStr;
6use std::io::{self, Write};
7use std::panic::{self, AssertUnwindSafe};
8
9/// Statistics configuration.
10///
11/// All options default to `false`.
12#[derive(Copy, Clone, Default)]
13pub struct Options {
14    /// If set, the output will be JSON-formatted.
15    ///
16    /// This corresponds to the `J` character.
17    pub json_format: bool,
18
19    /// If set, information that never changes during execution will be skipped.
20    ///
21    /// This corresponds to the `g` character.
22    pub skip_constants: bool,
23
24    /// If set, merged information about arenas will be skipped.
25    ///
26    /// This corresponds to the `m` character.
27    pub skip_merged_arenas: bool,
28
29    /// If set, information about individual arenas will be skipped.
30    ///
31    /// This corresponds to the `a` character.
32    pub skip_per_arena: bool,
33
34    /// If set, information about individual size classes for bins will be skipped.
35    ///
36    /// This corresponds to the `b` character.
37    pub skip_bin_size_classes: bool,
38
39    /// If set, information about individual size classes for large objects will be skipped.
40    ///
41    /// This corresponds to the `l` character.
42    pub skip_large_size_classes: bool,
43
44    /// If set, mutex statistics will be skipped.
45    ///
46    /// This corresponds to the `x` character.
47    pub skip_mutex_statistics: bool,
48
49    _p: (),
50}
51
52struct State<W> {
53    writer: W,
54    error: io::Result<()>,
55    panic: Result<(), Box<dyn Any + Send>>,
56}
57
58extern "C" fn callback<W>(opaque: *mut c_void, buf: *const c_char)
59where
60    W: Write,
61{
62    unsafe {
63        let state = &mut *(opaque as *mut State<W>);
64        if state.error.is_err() || state.panic.is_err() {
65            return;
66        }
67
68        let buf = CStr::from_ptr(buf);
69        match panic::catch_unwind(AssertUnwindSafe(|| {
70            state.writer.write_all(buf.to_bytes())
71        })) {
72            Ok(Ok(_)) => {}
73            Ok(Err(e)) => state.error = Err(e),
74            Err(e) => state.panic = Err(e),
75        }
76    }
77}
78
79/// Writes allocator statistics.
80///
81/// The information is the same that can be retrieved by the individual lookup methods in this
82/// crate, but all done at once.
83#[cfg_attr(feature = "cargo-clippy", allow(clippy::cast_possible_wrap))]
84pub fn stats_print<W>(writer: W, options: Options) -> io::Result<()>
85where
86    W: Write,
87{
88    unsafe {
89        let mut state = State {
90            writer,
91            error: Ok(()),
92            panic: Ok(()),
93        };
94        let mut opts = [0; 8];
95        let mut i = 0;
96        if options.json_format {
97            opts[i] = b'J' as c_char;
98            i += 1;
99        }
100        if options.skip_constants {
101            opts[i] = b'g' as c_char;
102            i += 1;
103        }
104        if options.skip_merged_arenas {
105            opts[i] = b'm' as c_char;
106            i += 1;
107        }
108        if options.skip_per_arena {
109            opts[i] = b'a' as c_char;
110            i += 1;
111        }
112        if options.skip_bin_size_classes {
113            opts[i] = b'b' as c_char;
114            i += 1;
115        }
116        if options.skip_large_size_classes {
117            opts[i] = b'l' as c_char;
118            i += 1;
119        }
120        if options.skip_mutex_statistics {
121            opts[i] = b'x' as c_char;
122            i += 1;
123        }
124        opts[i] = 0;
125
126        tikv_jemalloc_sys::malloc_stats_print(
127            Some(callback::<W>),
128            &mut state as *mut _ as *mut c_void,
129            opts.as_ptr(),
130        );
131        if let Err(e) = state.panic {
132            panic::resume_unwind(e);
133        }
134        state.error
135    }
136}
137
138#[cfg(test)]
139mod test {
140    use super::*;
141
142    #[test]
143    fn basic() {
144        let mut buf = vec![];
145        stats_print(&mut buf, Options::default()).unwrap();
146        println!("{}", String::from_utf8(buf).unwrap());
147    }
148
149    #[test]
150    fn all_options() {
151        let mut buf = vec![];
152        let options = Options {
153            json_format: true,
154            skip_constants: true,
155            skip_merged_arenas: true,
156            skip_per_arena: true,
157            skip_bin_size_classes: true,
158            skip_large_size_classes: true,
159            skip_mutex_statistics: true,
160            _p: (),
161        };
162        stats_print(&mut buf, options).unwrap();
163        println!("{}", String::from_utf8(buf).unwrap());
164    }
165}