zeropool/pool.rs
1use std::sync::atomic::{AtomicU64, Ordering};
2
3use crate::allocator::{Allocator, HeapAllocator};
4use crate::size_class::{ClassTable, SizeClass};
5use crate::stats::{Counters, Stats, snapshot};
6use crate::tls::TlsState;
7
8/// Global counter for unique pool instance IDs.
9static NEXT_ID: AtomicU64 = AtomicU64::new(1);
10
11/// Shared state backing all [`ZeroPool`] handles.
12///
13/// Holds identity, class routing table, and runtime configuration.
14/// All behavior lives on [`ZeroPool`].
15#[derive(Debug)]
16pub(crate) struct State {
17 pub id: u64,
18 pub table: ClassTable,
19 pub tls_cache_size: usize,
20 pub min_buffer_size: usize,
21 pub pinned_memory: bool,
22 pub batch_size: usize,
23 pub track_stats: bool,
24 pub counters: Counters,
25 pub allocator: Box<dyn Allocator>,
26}
27
28/// A user-space byte allocator with size-class bucketing and thread-local caching.
29///
30/// # Architecture
31///
32/// ```text
33/// Thread 1 Thread 2 Thread N
34/// ┌────────────┐ ┌────────────┐ ┌────────────┐
35/// │ TLS Cache │ │ TLS Cache │ │ TLS Cache │ ← Lock-free
36/// │ [class 0] │ │ [class 0] │ │ [class 0] │ per-class
37/// │ [class 1] │ │ [class 1] │ │ [class 1] │ LIFO caches
38/// │ ... │ │ ... │ │ ... │
39/// └─────┬──────┘ └─────┬──────┘ └─────┬──────┘
40/// │ batch │ batch │ batch
41/// └──────────┬───────┴───────────────────┘
42/// │
43/// ┌───────▼────────┐
44/// │ Shared Pool │
45/// │ (lock-free) │
46/// │ │
47/// │ [4KB queue] │ ArrayQueue per class
48/// │ [16KB queue] │ CAS-based push/pop
49/// │ [64KB queue] │ No mutex needed
50/// │ [256KB queue] │
51/// │ [1MB queue] │
52/// │ [4MB queue] │
53/// │ [16MB queue] │
54/// │ [64MB queue] │
55/// └────────────────┘
56/// ```
57#[derive(Debug)]
58pub struct ZeroPool {
59 pub(crate) state: State,
60}
61
62impl ZeroPool {
63 /// Create a new allocator with system-aware defaults.
64 ///
65 /// Chain configuration methods to customize before use:
66 ///
67 /// ```
68 /// use zeropool::ZeroPool;
69 ///
70 /// // Defaults
71 /// let pool = ZeroPool::new();
72 ///
73 /// // Custom
74 /// let pool = ZeroPool::new()
75 /// .min_buffer_size(4096)
76 /// .tls_cache_size(8)
77 /// .max_buffers_per_class(64)
78 /// .batch_size(4)
79 /// .track_stats(true);
80 /// ```
81 pub fn new() -> Self {
82 use crate::config::{
83 DEFAULT_MIN_BUFFER_SIZE, cpu_count, default_batch_size, default_max_buffers_per_class,
84 default_tls_cache_size,
85 };
86 let cpus = cpu_count();
87 let tls = default_tls_cache_size(cpus);
88 Self {
89 state: State {
90 id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
91 table: ClassTable::new(default_max_buffers_per_class(cpus)),
92 tls_cache_size: tls,
93 min_buffer_size: DEFAULT_MIN_BUFFER_SIZE,
94 pinned_memory: false,
95 batch_size: default_batch_size(tls),
96 track_stats: false,
97 counters: Counters::new(),
98 allocator: Box::new(HeapAllocator),
99 },
100 }
101 }
102
103 /// Set a custom allocator for buffer creation.
104 ///
105 /// Default: [`HeapAllocator`] (standard `Vec::with_capacity`).
106 ///
107 /// ```
108 /// use zeropool::{Allocator, ZeroPool};
109 ///
110 /// struct MyAllocator;
111 /// impl Allocator for MyAllocator {
112 /// fn allocate(&self, capacity: usize) -> Vec<u8> {
113 /// vec![0; capacity]
114 /// }
115 /// }
116 ///
117 /// let pool = ZeroPool::new().allocator(MyAllocator);
118 /// ```
119 pub fn allocator(self, alloc: impl Allocator) -> Self {
120 self.rebuild(|s| s.allocator = Box::new(alloc))
121 }
122
123 /// Set the minimum buffer size to keep in the pool.
124 ///
125 /// Buffers smaller than this are discarded on dealloc.
126 /// Default: 4KB
127 pub fn min_buffer_size(self, size: usize) -> Self {
128 self.rebuild(|s| s.min_buffer_size = size)
129 }
130
131 /// Set the number of buffers kept in thread-local cache per size class.
132 ///
133 /// Higher values reduce shared pool access but increase per-thread memory.
134 /// Also recomputes batch size (half of TLS cache, min 2) unless
135 /// `.batch_size()` is called afterwards to override.
136 /// Default: 2–8 based on CPU count
137 pub fn tls_cache_size(self, size: usize) -> Self {
138 assert!(size > 0, "tls_cache_size must be > 0");
139 self.rebuild(|s| {
140 s.tls_cache_size = size;
141 s.batch_size = crate::config::default_batch_size(size);
142 })
143 }
144
145 /// Set the maximum number of buffers per size class in the shared pool.
146 ///
147 /// Default: 32–128 based on CPU count
148 pub fn max_buffers_per_class(self, count: usize) -> Self {
149 assert!(count > 0, "max_buffers_per_class must be > 0");
150 self.rebuild(|s| s.table = ClassTable::new(count))
151 }
152
153 /// Enable pinned memory (mlock) for allocated buffers.
154 ///
155 /// Locks buffers in RAM to prevent swapping.
156 /// Default: false
157 pub fn pinned_memory(self, enabled: bool) -> Self {
158 self.rebuild(|s| s.pinned_memory = enabled)
159 }
160
161 /// Set the batch size for TLS ↔ shared pool transfers.
162 ///
163 /// When a thread-local cache misses, this many buffers are moved at once
164 /// from the shared pool (magazine-style).
165 /// Default: half of TLS cache size (min 2)
166 pub fn batch_size(self, size: usize) -> Self {
167 self.rebuild(|s| s.batch_size = size)
168 }
169
170 /// Enable or disable runtime statistics tracking.
171 ///
172 /// Disabled by default because hot-path atomic counters are measurable
173 /// overhead in tight allocation loops. Enable this when you need
174 /// [`stats()`](Self::stats) to report allocation counters.
175 pub fn track_stats(self, enabled: bool) -> Self {
176 self.rebuild(|s| s.track_stats = enabled)
177 }
178
179 fn rebuild(mut self, f: impl FnOnce(&mut State)) -> Self {
180 f(&mut self.state);
181 self
182 }
183
184 /// Allocate a zero-initialized buffer of at least `size` bytes.
185 ///
186 /// Returns a [`Buf`](crate::Buf) that automatically deallocates back
187 /// to the pool on drop.
188 ///
189 /// # Performance
190 ///
191 /// 1. **Fastest**: TLS cache pop (lock-free, ~24ns)
192 /// 2. **Fast**: Batch refill from shared pool (lock-free CAS)
193 /// 3. **Cold**: Fresh allocation via the configured [`Allocator`]
194 ///
195 /// # Example
196 /// ```
197 /// use zeropool::ZeroPool;
198 ///
199 /// let pool = ZeroPool::new();
200 /// let mut buf = pool.alloc(1024);
201 /// buf[0] = 42;
202 /// ```
203 #[inline]
204 #[must_use]
205 pub fn alloc(&self, size: usize) -> crate::Buf<'_> {
206 if self.state.track_stats {
207 self.state.counters.gets.fetch_add(1, Ordering::Relaxed);
208 }
209
210 let Some((class_idx, class)) = self.state.table.route(size) else {
211 if self.state.track_stats {
212 self.state.counters.oversize.fetch_add(1, Ordering::Relaxed);
213 self.state.counters.allocations.fetch_add(1, Ordering::Relaxed);
214 }
215 let mut buf = self.state.allocator.allocate(size);
216 buf.truncate(size);
217 self.pin(&mut buf);
218 return crate::Buf::new(buf, self, u8::MAX);
219 };
220
221 // ── TLS fast path (lock-free) ──────────────────────────────
222 let tls_result = TlsState::with(|tls| {
223 if !tls.owns(self.state.id) {
224 tls.bind(self.state.id, self.state.tls_cache_size);
225 }
226
227 if let Some(buf) = tls.caches[class_idx].pop() {
228 return Some((buf, true));
229 }
230
231 tls.refill(class_idx, class, self.state.batch_size).map(|buf| (buf, false))
232 });
233
234 let ci = class_idx as u8;
235
236 if let Some((mut buf, from_tls)) = tls_result {
237 if from_tls {
238 if self.state.track_stats {
239 self.state.counters.tls_hits.fetch_add(1, Ordering::Relaxed);
240 }
241 } else if self.state.track_stats {
242 self.state.counters.shared_hits.fetch_add(1, Ordering::Relaxed);
243 }
244 SizeClass::resize_zeroed(&mut buf, size);
245 return crate::Buf::new(buf, self, ci);
246 }
247
248 // ── Cold path: fresh allocation ────────────────────────────
249 if self.state.track_stats {
250 self.state.counters.allocations.fetch_add(1, Ordering::Relaxed);
251 }
252 let mut buf = self.state.allocator.allocate(class.class_size);
253 buf.truncate(size);
254 self.pin(&mut buf);
255 crate::Buf::new(buf, self, ci)
256 }
257
258 /// Allocate a buffer without zeroing its contents.
259 ///
260 /// This is the high-performance allocation path for workloads that fully
261 /// overwrite the buffer before reading it. The returned [`BufUninit`](crate::BufUninit)
262 /// does not expose a readable byte slice in safe code.
263 #[inline]
264 #[must_use]
265 pub fn alloc_uninit(&self, size: usize) -> crate::BufUninit<'_> {
266 if self.state.track_stats {
267 self.state.counters.gets.fetch_add(1, Ordering::Relaxed);
268 }
269
270 let Some((class_idx, class)) = self.state.table.route(size) else {
271 if self.state.track_stats {
272 self.state.counters.oversize.fetch_add(1, Ordering::Relaxed);
273 self.state.counters.allocations.fetch_add(1, Ordering::Relaxed);
274 }
275 let mut buf = Vec::with_capacity(size);
276 SizeClass::resize_uninit(&mut buf, size);
277 self.pin(&mut buf);
278 return crate::BufUninit::new(buf, self, u8::MAX);
279 };
280
281 let tls_result = TlsState::with(|tls| {
282 if !tls.owns(self.state.id) {
283 tls.bind(self.state.id, self.state.tls_cache_size);
284 }
285
286 if let Some(buf) = tls.caches[class_idx].pop() {
287 return Some((buf, true));
288 }
289
290 tls.refill(class_idx, class, self.state.batch_size).map(|buf| (buf, false))
291 });
292
293 let ci = class_idx as u8;
294
295 if let Some((mut buf, from_tls)) = tls_result {
296 if from_tls {
297 if self.state.track_stats {
298 self.state.counters.tls_hits.fetch_add(1, Ordering::Relaxed);
299 }
300 } else if self.state.track_stats {
301 self.state.counters.shared_hits.fetch_add(1, Ordering::Relaxed);
302 }
303 SizeClass::resize_uninit(&mut buf, size);
304 return crate::BufUninit::new(buf, self, ci);
305 }
306
307 if self.state.track_stats {
308 self.state.counters.allocations.fetch_add(1, Ordering::Relaxed);
309 }
310 let mut buf = Vec::with_capacity(class.class_size);
311 SizeClass::resize_uninit(&mut buf, size);
312 self.pin(&mut buf);
313 crate::BufUninit::new(buf, self, ci)
314 }
315
316 /// Return a buffer to the pool for reuse.
317 ///
318 /// `class_hint` is the class index stored in [`Buf`](crate::Buf)
319 /// at allocation time (`u8::MAX` for oversize buffers that bypass pooling).
320 #[inline(always)]
321 pub(crate) fn dealloc(&self, mut buffer: Vec<u8>, class_hint: u8) {
322 if self.state.track_stats {
323 self.state.counters.puts.fetch_add(1, Ordering::Relaxed);
324 }
325 buffer.clear();
326
327 if class_hint == u8::MAX {
328 if self.state.track_stats {
329 self.state.counters.discards.fetch_add(1, Ordering::Relaxed);
330 }
331 return;
332 }
333
334 let cap = buffer.capacity();
335
336 if cap < self.state.min_buffer_size {
337 if self.state.track_stats {
338 self.state.counters.discards.fetch_add(1, Ordering::Relaxed);
339 }
340 return;
341 }
342
343 let class_idx = if cap >= ClassTable::boundary(class_hint as usize) {
344 class_hint as usize
345 } else {
346 let Some((idx, _)) = self.state.table.route_capacity(cap) else {
347 return;
348 };
349 idx
350 };
351
352 self.pin(&mut buffer);
353
354 // ── TLS fast path ──────────────────────────────────────────
355 let overflow = TlsState::with(|tls| {
356 if !tls.owns(self.state.id) {
357 tls.bind(self.state.id, self.state.tls_cache_size);
358 }
359
360 let class = &self.state.table[class_idx];
361
362 if tls.caches[class_idx].len() >= tls.limit {
363 tls.spill(class_idx, class, self.state.batch_size);
364 }
365
366 if tls.caches[class_idx].len() < tls.limit {
367 tls.caches[class_idx].push(buffer);
368 return None;
369 }
370
371 Some(buffer)
372 });
373
374 if let Some(buf) = overflow {
375 let _ = self.state.table[class_idx].push(buf);
376 }
377 }
378
379 /// Warm up the pool by pre-allocating buffers for the given size class.
380 ///
381 /// # Example
382 /// ```
383 /// use zeropool::ZeroPool;
384 ///
385 /// let pool = ZeroPool::new().min_buffer_size(0).track_stats(true);
386 /// pool.warm(16, 64 * 1024); // 16 × 64KB buffers
387 /// ```
388 pub fn warm(&self, count: usize, size: usize) {
389 let Some((_, class)) = self.state.table.route(size) else {
390 return;
391 };
392
393 for _ in 0..count {
394 let mut buf = self.state.allocator.allocate(class.class_size);
395 self.pin(&mut buf);
396 if class.push(buf).is_err() {
397 break;
398 }
399 }
400 }
401
402 /// Total number of buffers across all shared size classes.
403 ///
404 /// Does not include thread-local cached buffers.
405 #[inline]
406 #[must_use]
407 pub fn len(&self) -> usize {
408 self.state.table.total_buffered()
409 }
410
411 /// Whether all shared size classes are empty.
412 ///
413 /// Does not check thread-local caches.
414 #[inline]
415 #[must_use]
416 pub fn is_empty(&self) -> bool {
417 self.state.table.all_empty()
418 }
419
420 /// Drain all buffers from all shared size classes.
421 ///
422 /// Thread-local caches are NOT cleared.
423 pub fn drain(&self) {
424 self.state.table.clear_all();
425 }
426
427 /// Point-in-time snapshot of allocator statistics.
428 ///
429 /// # Example
430 /// ```
431 /// use zeropool::ZeroPool;
432 ///
433 /// let pool = ZeroPool::new().min_buffer_size(0).track_stats(true);
434 /// let buf = pool.alloc(4096);
435 /// drop(buf);
436 ///
437 /// let s = pool.stats();
438 /// assert_eq!(s.gets, 1);
439 /// assert_eq!(s.puts, 1);
440 /// println!("{s}");
441 /// ```
442 #[inline]
443 pub fn stats(&self) -> Stats {
444 snapshot(&self.state.counters, self.state.table.classes())
445 }
446
447 /// Reset all performance counters to zero.
448 pub fn reset_stats(&self) {
449 self.state.counters.reset();
450 }
451
452 /// Pin buffer memory to RAM if configured.
453 #[inline(always)]
454 fn pin(&self, buffer: &mut Vec<u8>) {
455 if !self.state.pinned_memory {
456 return;
457 }
458 if buffer.capacity() == 0 {
459 return;
460 }
461 let _ = region::lock(buffer.as_ptr(), buffer.capacity());
462 }
463}
464
465impl Default for ZeroPool {
466 fn default() -> Self {
467 Self::new()
468 }
469}