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
11#[cfg(feature = "jemalloc_conf")]
12use const_str::concat_bytes;
13use jevmalloc::{
14	Jemalloc,
15	global::hook::{ALLOC, ALLOC_ZEROED},
16	stats::print as print_stats,
17};
18pub use jevmalloc::{arenas::trim, background_thread_enable};
19use libc::{STDOUT_FILENO, c_void, write};
20
21use crate::{arrayvec::ArrayVec, utils::BoolExt};
22
23/// Line buffer for one allocation-trace record. A record of three integers at
24/// their maximum widths occupies 74 bytes.
25type TraceLine = ArrayVec<u8, 128>;
26
27/// Provides the process-wide jemalloc startup configuration.
28///
29/// Jemalloc reads this symbol during allocator initialization, which can occur
30/// before `main`. The NUL-terminated options tune the cache and decay
31/// thresholds. CPU-affine arenas, metadata huge pages and background thread
32/// settings are requested only on the platforms whose jemalloc provides them.
33#[cfg(feature = "jemalloc_conf")]
34#[used]
35#[unsafe(no_mangle)]
36pub static malloc_conf: &[u8] = MALLOC_CONF;
37
38/// Provides the jemalloc startup configuration under the symbol name a
39/// prefixed build reads.
40///
41/// Jemalloc renames its public symbols when configured with a prefix, which
42/// `jevmalloc-sys` does for musl, Apple, Android, DragonFly and OpenBSD, and
43/// which a substituted `JEMALLOC_OVERRIDE` library may do on any target.
44/// Defining both names lets the linked allocator take whichever one it declares
45/// and leaves the other unreferenced.
46#[cfg(feature = "jemalloc_conf")]
47#[used]
48#[unsafe(export_name = "_rjem_malloc_conf")]
49pub static MALLOC_CONF_PREFIXED: &[u8] = MALLOC_CONF;
50
51#[cfg(feature = "jemalloc_conf")]
52const MALLOC_CONF: &[u8] = concat_bytes!(
53	"tcache:true",
54	MALLOC_CONF_PERCPU_ARENA,
55	MALLOC_CONF_METADATA_THP,
56	MALLOC_CONF_BACKGROUND,
57	",lg_extent_max_active_fit:4",
58	",oversize_threshold:2097152",
59	",tcache_max:8192",
60	",dirty_decay_ms:16000",
61	",muzzy_decay_ms:144000",
62	//MALLOC_CONF_PROF,
63	0
64);
65
66/// Assigns arenas by CPU where jemalloc can tell which CPU a thread runs on.
67///
68/// Jemalloc needs `sched_getcpu`, or its own Apple and Windows paths. Without
69/// them a release build falls back to its default arena count and prints a
70/// notice, and a debug build aborts during allocator initialization.
71#[cfg(feature = "jemalloc_conf")]
72const MALLOC_CONF_PERCPU_ARENA: &str = if cfg!(any(
73	target_os = "linux",
74	target_os = "android",
75	target_os = "freebsd",
76	target_os = "dragonfly",
77	target_vendor = "apple",
78	target_os = "windows",
79)) {
80	",percpu_arena:percpu"
81} else {
82	""
83};
84
85/// Asks jemalloc to back allocator metadata with huge pages where it can.
86///
87/// Jemalloc needs `MADV_HUGEPAGE` (which its build ignores on 32-bit ARM) or
88/// `memcntl`. Without either, a debug build aborts during allocator
89/// initialization, and the key's only remaining effect in release is
90/// whether the tcache stacks come from the never-purging base allocator,
91/// since base blocks keep the 2 MiB alignment and rounding of
92/// `BASE_BLOCK_MIN_ALIGN` either way.
93#[cfg(feature = "jemalloc_conf")]
94const MALLOC_CONF_METADATA_THP: &str = if cfg!(any(
95	all(any(target_os = "linux", target_os = "android"), not(target_arch = "arm")),
96	target_os = "illumos",
97	target_os = "solaris",
98)) {
99	",metadata_thp:always"
100} else {
101	""
102};
103
104// Apple's jemalloc has no background threads and prints a notice when asked.
105#[cfg(all(feature = "jemalloc_conf", not(target_vendor = "apple")))]
106const MALLOC_CONF_BACKGROUND: &str = ",background_thread:false,max_background_threads:-1";
107#[cfg(all(feature = "jemalloc_conf", target_vendor = "apple"))]
108const MALLOC_CONF_BACKGROUND: &str = "";
109
110#[cfg(all(
111	feature = "jemalloc_conf",
112	feature = "jemalloc_prof",
113	target_arch = "x86_64",
114))]
115const _MALLOC_CONF_PROF: &str = ",prof_active:false";
116#[cfg(all(
117	feature = "jemalloc_conf",
118	any(not(feature = "jemalloc_prof"), not(target_arch = "x86_64")),
119))]
120const _MALLOC_CONF_PROF: &str = "";
121
122#[global_allocator]
123static JEMALLOC: Jemalloc = Jemalloc;
124
125static GLOBAL_ALLOCS: AtomicU64 = AtomicU64::new(0);
126static COUNT_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
127static TRACE_GLOBAL_ALLOCS: AtomicBool = AtomicBool::new(false);
128
129/// Registers the allocation-observer callbacks during process startup.
130///
131/// Normal and zeroed allocations made after registration feed the same counting
132/// and tracing instrumentation. The allocator reads these slots only when
133/// `jevmalloc` is built with its `global_hooks` feature.
134#[crate::ctor(unsafe)]
135fn _static_initialization() {
136	// SAFETY: Mutable static globals in jemalloc crate; must be initialized
137	// properly and uniquely.
138	unsafe { ALLOC = Some(global_alloc_hook) };
139
140	// SAFETY: As above.
141	unsafe { ALLOC_ZEROED = Some(global_alloc_zeroed_hook) };
142}
143
144fn global_alloc_hook(layout: Layout) {
145	catch_unwind(move || handle_global_alloc(layout))
146		.map_err(|_| abort())
147		.ok();
148}
149
150fn global_alloc_zeroed_hook(layout: Layout) {
151	catch_unwind(move || handle_global_alloc(layout))
152		.map_err(|_| abort())
153		.ok();
154}
155
156fn handle_global_alloc(layout: Layout) {
157	let do_count = COUNT_GLOBAL_ALLOCS.load(Ordering::Relaxed);
158	let count = GLOBAL_ALLOCS.fetch_add(do_count.into(), Ordering::Relaxed);
159
160	if TRACE_GLOBAL_ALLOCS.load(Ordering::Relaxed) {
161		let mut buf = TraceLine::new();
162
163		writeln!(&mut buf, "{count} align={} size={}", layout.align(), layout.size())
164			.expect("writeln! to buffer failed");
165
166		// SAFETY: Valid ptr and len from buf for writing to stdout.
167		unsafe { write(STDOUT_FILENO, buf.as_ptr().cast::<c_void>(), buf.len()) }
168			.ge(&0)
169			.into_result()
170			.expect("write(2) error");
171	}
172}
173
174/// Returns the process allocation count observed by the allocator hook.
175///
176/// The counter uses relaxed ordering and advances only when internal allocation
177/// counting is enabled. It is intended for allocation measurements rather than
178/// synchronized accounting.
179#[inline]
180#[must_use]
181pub fn global_alloc_count() -> u64 { GLOBAL_ALLOCS.load(Ordering::Relaxed) }
182
183/// Collects jemalloc's UTF-8 statistics report with the supplied print options.
184///
185/// Returns `None` if jemalloc produces no report, the report exceeds 1 MiB, or
186/// the report contains invalid UTF-8.
187#[must_use]
188pub fn memory_stats(opts: &str) -> Option<String> {
189	const MAX_LENGTH: usize = 1_048_576;
190
191	let mut stats = vec![0; MAX_LENGTH];
192	let length = print_stats(opts, &mut stats).ok()?.len();
193	if length == 0 {
194		return None;
195	}
196
197	stats.truncate(length);
198	String::from_utf8(stats).ok()
199}
200
201/// Exposes jemalloc state controls associated with the calling thread.
202///
203/// These functions resolve the thread's own arena before applying the
204/// operation. They return jevmalloc's control errors unchanged.
205pub mod this_thread {
206	pub use jevmalloc::thread::this::{decay, set_muzzy_decay};
207}
tuwunel_core/alloc/je.rs

tuwunel_core/alloc/
je.rs
