caco3_web/jemalloc/
init.rs

1use std::env;
2use std::fmt::Display;
3use std::fmt::Write;
4use std::os::unix::prelude::CommandExt;
5use std::process::Command;
6
7use serde::{Deserialize, Serialize};
8
9/// Jemalloc configuration.
10#[derive(Clone, Debug, Deserialize, Serialize, Default)]
11pub struct Jemalloc {
12    #[serde(default)]
13    pub background_thread: bool,
14    #[serde(default)]
15    pub max_background_threads: Option<u32>,
16    #[serde(default)]
17    pub number_of_arenas: Option<u32>,
18    #[serde(default)]
19    pub extra_conf: Option<String>
20}
21
22pub const POSSIBLE_MALLOC_CONF_ENVIRONMENT_VARIABLES: &[&str] =
23    &["MALLOC_CONF", "_RJEM_MALLOC_CONF"];
24
25/// Return `true` if jemalloc background managed threads is supported.
26pub const fn is_background_thread_supported() -> bool {
27    // See https://github.com/tikv/jemallocator/blob/main/jemalloc-sys/src/env.rs
28    if cfg!(target_env = "musl") {
29        return false;
30    }
31    // Background thread on MacOS is not supported.
32    // See https://github.com/jemalloc/jemalloc/issues/843
33    if cfg!(target_os = "macos") {
34        return false;
35    }
36    true
37}
38
39/// Re-execute current process to apply jemalloc configuration.
40pub fn apply_config(config: &Jemalloc, f: impl FnOnce(&str)) -> ! {
41    // Some configuration of jemalloc need to be configured before main program is started.
42    // But at this point, main program has been started, how do we solve this?
43    //
44    // We replace current process with itself but with properly jemalloc configuration.
45    let malloc_conf = config.to_config();
46
47    let mut args = env::args_os();
48    let program = args.next().expect("Process name");
49    let mut cmd = Command::new(program);
50    cmd.args(args);
51    for name in POSSIBLE_MALLOC_CONF_ENVIRONMENT_VARIABLES {
52        cmd.env(name, &malloc_conf);
53    }
54    f(&malloc_conf);
55    let err = cmd.exec();
56    panic!("jemalloc: exec error: {:?}", err);
57}
58
59/// Returns `true` if jemalloc is configured.
60pub fn is_configured() -> bool {
61    POSSIBLE_MALLOC_CONF_ENVIRONMENT_VARIABLES
62        .iter()
63        .any(|name| env::var_os(name).is_some())
64}
65
66impl Jemalloc {
67    pub fn to_config(&self) -> String {
68        let mut config = String::with_capacity(64);
69        // Abort program if invalid jemalloc configurations are found.
70        config.push_str("abort_conf:true");
71
72        let mut write_config = |key: &str, value: &dyn Display| {
73            write!(&mut config, ",{}:{}", key, value)
74                .expect("a Display implementation returned an error unexpectedly");
75        };
76        if self.background_thread {
77            // Do nothing, this is intended.
78            // background thread should be enabled at runtime to avoid deadlock.
79        }
80        if let Some(v) = self.max_background_threads {
81            write_config("max_background_threads", &v);
82        }
83        if let Some(v) = self.number_of_arenas {
84            write_config("narenas", &v);
85        }
86        if let Some(extra_conf) = self.extra_conf.as_deref() {
87            config.push(',');
88            config.push_str(extra_conf);
89        }
90        config
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn jemalloc_to_config() {
100        let val = Jemalloc {
101            background_thread: false,
102            max_background_threads: None,
103            number_of_arenas: None,
104            extra_conf: None,
105        };
106        assert_eq!(val.to_config(), "abort_conf:true");
107
108        let val = Jemalloc {
109            background_thread: false,
110            max_background_threads: None,
111            number_of_arenas: None,
112            extra_conf: Some("tcache:false".to_owned()),
113        };
114        assert_eq!(val.to_config(), "abort_conf:true,tcache:false");
115
116        let val = Jemalloc {
117            background_thread: false,
118            max_background_threads: None,
119            number_of_arenas: Some(16),
120            extra_conf: None,
121        };
122        assert_eq!(val.to_config(), "abort_conf:true,narenas:16");
123
124        let val = Jemalloc {
125            background_thread: true,
126            max_background_threads: None,
127            number_of_arenas: None,
128            extra_conf: None,
129        };
130        assert_eq!(val.to_config(), "abort_conf:true");
131
132        let val = Jemalloc {
133            background_thread: false,
134            max_background_threads: Some(4),
135            number_of_arenas: None,
136            extra_conf: None,
137        };
138        assert_eq!(val.to_config(), "abort_conf:true,max_background_threads:4");
139
140        let val = Jemalloc {
141            background_thread: true,
142            max_background_threads: Some(8),
143            number_of_arenas: Some(64),
144            extra_conf: None,
145        };
146        assert_eq!(
147            val.to_config(),
148            "abort_conf:true,max_background_threads:8,narenas:64"
149        );
150    }
151}
caco3_web/jemalloc/init.rs

caco3_web/jemalloc/
init.rs
