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::with_capacity(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 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.allocate_raw(size, size);
216 self.pin(&mut buf);
217 return crate::Buf::new(buf, self, u8::MAX);
218 };
219
220 // ── TLS fast path (lock-free) ──────────────────────────────
221 let tls_result = TlsState::with(|tls| {
222 if !tls.owns(self.state.id) {
223 tls.bind(self.state.id, self.state.tls_cache_size);
224 }
225
226 if let Some(buf) = tls.caches[class_idx].pop() {
227 return Some((buf, true));
228 }
229
230 tls.refill(class_idx, class, self.state.batch_size).map(|buf| (buf, false))
231 });
232
233 let ci = class_idx as u8;
234
235 if let Some((mut buf, from_tls)) = tls_result {
236 if from_tls {
237 if self.state.track_stats {
238 self.state.counters.tls_hits.fetch_add(1, Ordering::Relaxed);
239 }
240 } else if self.state.track_stats {
241 self.state.counters.shared_hits.fetch_add(1, Ordering::Relaxed);
242 }
243 SizeClass::resize(&mut buf, size);
244 return crate::Buf::new(buf, self, ci);
245 }
246
247 // ── Cold path: fresh allocation ────────────────────────────
248 if self.state.track_stats {
249 self.state.counters.allocations.fetch_add(1, Ordering::Relaxed);
250 }
251 let mut buf = self.allocate_raw(class.class_size, size);
252 self.pin(&mut buf);
253 crate::Buf::new(buf, self, ci)
254 }
255
256 /// Return a buffer to the pool for reuse.
257 ///
258 /// `class_hint` is the class index stored in [`Buf`](crate::Buf)
259 /// at allocation time (`u8::MAX` for oversize buffers that bypass pooling).
260 #[inline(always)]
261 pub(crate) fn dealloc(&self, mut buffer: Vec<u8>, class_hint: u8) {
262 if self.state.track_stats {
263 self.state.counters.puts.fetch_add(1, Ordering::Relaxed);
264 }
265 buffer.clear();
266
267 if class_hint == u8::MAX {
268 if self.state.track_stats {
269 self.state.counters.discards.fetch_add(1, Ordering::Relaxed);
270 }
271 return;
272 }
273
274 let cap = buffer.capacity();
275
276 if cap < self.state.min_buffer_size {
277 if self.state.track_stats {
278 self.state.counters.discards.fetch_add(1, Ordering::Relaxed);
279 }
280 return;
281 }
282
283 let class_idx = if cap >= ClassTable::boundary(class_hint as usize) {
284 class_hint as usize
285 } else {
286 let Some((idx, _)) = self.state.table.route_capacity(cap) else {
287 return;
288 };
289 idx
290 };
291
292 self.pin(&mut buffer);
293
294 // ── TLS fast path ──────────────────────────────────────────
295 let overflow = TlsState::with(|tls| {
296 if !tls.owns(self.state.id) {
297 tls.bind(self.state.id, self.state.tls_cache_size);
298 }
299
300 let class = &self.state.table[class_idx];
301
302 if tls.caches[class_idx].len() >= tls.limit {
303 tls.spill(class_idx, class, self.state.batch_size);
304 }
305
306 if tls.caches[class_idx].len() < tls.limit {
307 tls.caches[class_idx].push(buffer);
308 return None;
309 }
310
311 Some(buffer)
312 });
313
314 if let Some(buf) = overflow {
315 let _ = self.state.table[class_idx].push(buf);
316 }
317 }
318
319 /// Warm up the pool by pre-allocating buffers for the given size class.
320 ///
321 /// # Example
322 /// ```
323 /// use zeropool::ZeroPool;
324 ///
325 /// let pool = ZeroPool::new().min_buffer_size(0).track_stats(true);
326 /// pool.warm(16, 64 * 1024); // 16 × 64KB buffers
327 /// ```
328 pub fn warm(&self, count: usize, size: usize) {
329 let Some((_, class)) = self.state.table.route(size) else {
330 return;
331 };
332
333 for _ in 0..count {
334 let mut buf = self.state.allocator.allocate(class.class_size);
335 self.pin(&mut buf);
336 if class.push(buf).is_err() {
337 break;
338 }
339 }
340 }
341
342 /// Total number of buffers across all shared size classes.
343 ///
344 /// Does not include thread-local cached buffers.
345 #[inline]
346 #[must_use]
347 pub fn len(&self) -> usize {
348 self.state.table.total_buffered()
349 }
350
351 /// Whether all shared size classes are empty.
352 ///
353 /// Does not check thread-local caches.
354 #[inline]
355 #[must_use]
356 pub fn is_empty(&self) -> bool {
357 self.state.table.all_empty()
358 }
359
360 /// Drain all buffers from all shared size classes.
361 ///
362 /// Thread-local caches are NOT cleared.
363 pub fn drain(&self) {
364 self.state.table.clear_all();
365 }
366
367 /// Point-in-time snapshot of allocator statistics.
368 ///
369 /// # Example
370 /// ```
371 /// use zeropool::ZeroPool;
372 ///
373 /// let pool = ZeroPool::new().min_buffer_size(0).track_stats(true);
374 /// let buf = pool.alloc(4096);
375 /// drop(buf);
376 ///
377 /// let s = pool.stats();
378 /// assert_eq!(s.gets, 1);
379 /// assert_eq!(s.puts, 1);
380 /// println!("{s}");
381 /// ```
382 #[inline]
383 pub fn stats(&self) -> Stats {
384 snapshot(&self.state.counters, self.state.table.classes())
385 }
386
387 /// Reset all performance counters to zero.
388 pub fn reset_stats(&self) {
389 self.state.counters.reset();
390 }
391
392 /// Allocate a raw buffer via the configured allocator and set its length.
393 #[cold]
394 fn allocate_raw(&self, capacity: usize, len: usize) -> Vec<u8> {
395 let mut buf = self.state.allocator.allocate(capacity);
396 // SAFETY: allocator guarantees capacity >= `capacity` >= `len`.
397 // All u8 bit patterns are valid.
398 unsafe {
399 buf.set_len(len);
400 }
401 buf
402 }
403
404 /// Pin buffer memory to RAM if configured.
405 #[inline(always)]
406 fn pin(&self, buffer: &mut Vec<u8>) {
407 if !self.state.pinned_memory {
408 return;
409 }
410 if buffer.capacity() == 0 {
411 return;
412 }
413 // SAFETY: capacity was allocated; all u8 patterns valid.
414 unsafe { buffer.set_len(buffer.capacity()) };
415 let _ = region::lock(buffer.as_ptr(), buffer.len());
416 buffer.clear();
417 }
418}
419
420impl Default for ZeroPool {
421 fn default() -> Self {
422 Self::new()
423 }
424}