tuwunel_core/alloc/
je.rs

1//! jemalloc allocator
2
3use std::{
4	alloc::Layout,
5	io::Write,
6	panic::catch_unwind,
7	process::abort,
8	sync::atomic::{AtomicBool, AtomicU64, Ordering},
9};
10
11use jevmalloc::{
12	Jemalloc,
13	global::hook::{ALLOC, ALLOC_ZEROED},
14	stats::print as print_stats,
15};
16pub use jevmalloc::{arenas::trim, background_thread_enable};
17use libc::{STDOUT_FILENO, c_void, write};
18
19use crate::{arrayvec::ArrayVec, utils::BoolExt};
20
21/// Line buffer for one allocation-trace record. A record of three integers at
22/// their maximum widths occupies 74 bytes.
23type TraceLine = ArrayVec<u8, 128>;
24
25/// Provides the process-wide jemalloc startup configuration.
26///
27/// Jemalloc reads this unmangled symbol during allocator initialization, which
28/// can occur before `main`. The NUL-terminated options enable CPU-affine
29/// arenas, background purging, metadata huge pages, and tuned cache and decay
30/// thresholds.
31#[cfg(feature = "jemalloc_conf")]
32#[used]
33#[unsafe(no_mangle)]
34pub static malloc_conf: &[u8] = const_str::concat_bytes!(
35	"tcache:true",
36	",percpu_arena:percpu",
37	",metadata_thp:always",
38	",background_thread:true",
39	",max_background_threads:-1",
40	",lg_extent_max_active_fit:4",
41	",oversize_threshold:2097152",
42	",tcache_max:524288",
43	",dirty_decay_ms:16000",
44	",muzzy_decay_ms:144000",
45	//MALLOC_CONF_PROF,
46	0
47);
48
49#[cfg(all(
50	feature = "jemalloc_conf",
51	feature = "jemalloc_prof",
52	target_arch = "x86_64",
53))]
54const _MALLOC_CONF_PROF: &str = ",prof_active:false";
55#[cfg(all(
56	feature = "jemalloc_conf",
57	any(not(feature = "jemalloc_prof"), not(target_arch = "x86_64")),
58))]
59const _MALLOC_CONF_PROF: &str = "";
60
61#[global_allocator]
62static JEMALLOC: Jemalloc = Jemalloc;
63
64static GLOBAL_ALLOCS: AtomicU64 = AtomicU64::new(0);
65static COUNT_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
66static TRACE_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
67
68/// Registers the allocation-observer callbacks during process startup.
69///
70/// Normal and zeroed allocations made after registration feed the same counting
71/// and tracing instrumentation. The allocator reads these slots only when
72/// `jevmalloc` is built with its `global_hooks` feature.
73#[crate::ctor(unsafe)]
74fn _static_initialization() {
75	// SAFETY: Mutable static globals in jemalloc crate; must be initialized
76	// properly and uniquely.
77	unsafe { ALLOC = Some(global_alloc_hook) };
78
79	// SAFETY: As above.
80	unsafe { ALLOC_ZEROED = Some(global_alloc_zeroed_hook) };
81}
82
83fn global_alloc_hook(layout: Layout) {
84	catch_unwind(move || handle_global_alloc(layout))
85		.map_err(|_| abort())
86		.ok();
87}
88
89fn global_alloc_zeroed_hook(layout: Layout) {
90	catch_unwind(move || handle_global_alloc(layout))
91		.map_err(|_| abort())
92		.ok();
93}
94
95fn handle_global_alloc(layout: Layout) {
96	let do_count = COUNT_GLOBAL_ALLOCS.load(Ordering::Relaxed);
97	let count = GLOBAL_ALLOCS.fetch_add(do_count.into(), Ordering::Relaxed);
98
99	if TRACE_GLOBAL_ALLOCS.load(Ordering::Relaxed) {
100		let mut buf = TraceLine::new();
101
102		writeln!(&mut buf, "{count} align={} size={}", layout.align(), layout.size())
103			.expect("writeln! to buffer failed");
104
105		// SAFETY: Valid ptr and len from buf for writing to stdout.
106		unsafe { write(STDOUT_FILENO, buf.as_ptr().cast::<c_void>(), buf.len()) }
107			.ge(&0)
108			.into_result()
109			.expect("write(2) error");
110	}
111}
112
113/// Returns the process allocation count observed by the allocator hook.
114///
115/// The counter uses relaxed ordering and advances only when internal allocation
116/// counting is enabled. It is intended for allocation measurements rather than
117/// synchronized accounting.
118#[inline]
119#[must_use]
120pub fn global_alloc_count() -> u64 { GLOBAL_ALLOCS.load(Ordering::Relaxed) }
121
122/// Collects jemalloc's UTF-8 statistics report with the supplied print options.
123///
124/// Returns `None` if jemalloc produces no report, the report exceeds 1 MiB, or
125/// the report contains invalid UTF-8.
126#[must_use]
127pub fn memory_stats(opts: &str) -> Option<String> {
128	const MAX_LENGTH: usize = 1_048_576;
129
130	let mut stats = vec![0; MAX_LENGTH];
131	let length = print_stats(opts, &mut stats).ok()?.len();
132	if length == 0 {
133		return None;
134	}
135
136	stats.truncate(length);
137	String::from_utf8(stats).ok()
138}
139
140/// Exposes jemalloc state controls associated with the calling thread.
141///
142/// These functions resolve the thread's own arena before applying the
143/// operation. They return jevmalloc's control errors unchanged.
144pub mod this_thread {
145	pub use jevmalloc::thread::this::{decay, set_muzzy_decay};
146}
tuwunel_core/alloc/je.rs

tuwunel_core/alloc/
je.rs
