1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

//! Bulk statistics output.

use libc::{c_char, c_void};
use std::any::Any;
use std::ffi::CStr;
use std::io::{self, Write};
use std::panic::{self, AssertUnwindSafe};

/// Statistics configuration.
///
/// All options default to `false`.
#[derive(Copy, Clone, Default)]
pub struct Options {
    /// If set, the output will be JSON-formatted.
    ///
    /// This corresponds to the `J` character.
    pub json_format: bool,

    /// If set, information that never changes during execution will be skipped.
    ///
    /// This corresponds to the `g` character.
    pub skip_constants: bool,

    /// If set, merged information about arenas will be skipped.
    ///
    /// This corresponds to the `m` character.
    pub skip_merged_arenas: bool,

    /// If set, information about individual arenas will be skipped.
    ///
    /// This corresponds to the `a` character.
    pub skip_per_arena: bool,

    /// If set, information about individual size classes for bins will be skipped.
    ///
    /// This corresponds to the `b` character.
    pub skip_bin_size_classes: bool,

    /// If set, information about individual size classes for large objects will be skipped.
    ///
    /// This corresponds to the `l` character.
    pub skip_large_size_classes: bool,

    /// If set, mutex statistics will be skipped.
    ///
    /// This corresponds to the `x` character.
    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),
        }
    }
}

/// Writes allocator statistics.
///
/// The information is the same that can be retrieved by the individual lookup methods in this
/// crate, but all done at once.
#[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());
    }
}