stoolap_jemalloc/lib.rs
1// Copyright 2026 Stoolap Contributors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A jemalloc-style memory allocator in pure Rust, with a sampling heap
16//! profiler that writes pprof profiles.
17//!
18//! ```no_run
19//! use stoolap_jemalloc::Jemalloc;
20//!
21//! #[global_allocator]
22//! static GLOBAL: Jemalloc = Jemalloc::new();
23//!
24//! fn main() {
25//! stoolap_jemalloc::prof::activate();
26//! // ... work ...
27//! stoolap_jemalloc::prof::write_pprof("heap.pb").unwrap();
28//! }
29//! ```
30//!
31//! Then `go tool pprof -http=: heap.pb`.
32//!
33//! Design, after jemalloc:
34//! - size classes with 8-byte steps up to 128 bytes, then four per doubling
35//! - a thread cache per thread, flushed in batches and trimmed over time
36//! - 4 arenas per CPU; small classes come from slabs, larger ones from page
37//! runs inside 4 MiB chunks, and those above 1 MiB from their own mappings
38//! - freed pages go back to the OS after 5 to 10 seconds
39//! - heap profiling by sampling, on average once per 512 KiB allocated
40
41mod arena;
42pub mod background;
43mod base;
44mod chunk;
45#[cfg(all(unix, not(miri)))]
46mod fork;
47mod huge;
48mod lock;
49mod malloc;
50mod os;
51pub mod prof;
52mod size_class;
53mod stats;
54mod tcache;
55
56pub use stats::{Stats, stats};
57
58use core::alloc::{GlobalAlloc, Layout};
59
60/// The allocator; install it with `#[global_allocator]`:
61///
62/// ```no_run
63/// use stoolap_jemalloc::Jemalloc;
64///
65/// #[global_allocator]
66/// static GLOBAL: Jemalloc = Jemalloc::new();
67/// ```
68#[derive(Debug, Clone, Copy, Default)]
69pub struct Jemalloc {
70 profiling: bool,
71}
72
73impl Jemalloc {
74 /// The allocator, with heap profiling off until `prof::activate()`
75 #[must_use]
76 pub const fn new() -> Self {
77 Jemalloc { profiling: false }
78 }
79
80 /// The allocator with heap profiling on from the process's first
81 /// allocation, as jemalloc's `prof:true,prof_active:true` options
82 /// set it, without a call to `prof::activate()`.
83 /// `prof::deactivate()` still stops it.
84 #[must_use]
85 pub const fn with_profiling(self) -> Self {
86 Jemalloc { profiling: true }
87 }
88}
89
90unsafe impl GlobalAlloc for Jemalloc {
91 #[inline(always)]
92 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
93 malloc::alloc(layout.size(), layout.align(), self.profiling)
94 }
95
96 #[inline(always)]
97 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
98 malloc::dealloc(ptr, layout.size(), layout.align());
99 }
100
101 #[inline(always)]
102 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
103 malloc::alloc_zeroed(layout.size(), layout.align(), self.profiling)
104 }
105
106 #[inline]
107 unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
108 malloc::realloc(ptr, layout.size(), layout.align(), new_size, self.profiling)
109 }
110}
111
112/// Returns the calling thread's cached objects, then every free page and
113/// cached huge mapping, to the OS
114pub fn purge() {
115 unsafe {
116 let t = tcache::current();
117 if !t.is_null() {
118 tcache::flush_all(&mut *t);
119 }
120 arena::purge();
121 huge::purge();
122 }
123}