Skip to main content

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}