reifydb_core/util/
slab.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::sync::Arc;
5
6use reifydb_runtime::sync::mutex::Mutex;
7
8/// Reusable pool of `Arc<T>` slabs.
9///
10/// `T` must be `Default` so the pool can mint fresh empty slabs when its
11/// reserve is exhausted.
12pub struct Slab<T> {
13	pool: Mutex<Vec<Arc<T>>>,
14	cap: usize,
15}
16
17impl<T: Default> Slab<T> {
18	/// Create a new `Slab` that retains at most `cap` reusable slabs.
19	pub fn new(cap: usize) -> Self {
20		Self {
21			pool: Mutex::new(Vec::new()),
22			cap,
23		}
24	}
25
26	/// Pull a slab from the pool, or allocate a fresh one if the pool
27	/// is empty (or every retained slab is still referenced elsewhere).
28	///
29	/// Returned slabs always have `strong_count == 1`, so the caller
30	/// can mutate them in place via `Arc::make_mut` without the COW
31	/// fork penalty.
32	pub fn acquire(&self) -> Arc<T> {
33		let mut pool = self.pool.lock();
34		while let Some(slab) = pool.pop() {
35			if Arc::strong_count(&slab) == 1 {
36				return slab;
37			}
38			// Slab is still referenced (e.g. an in-flight consumer
39			// has not dropped its clone yet). Drop our local
40			// reference and try the next pool entry; eventually we
41			// either find a unique one or fall through to allocating
42			// fresh.
43		}
44		drop(pool);
45		Arc::new(T::default())
46	}
47
48	/// Return a slab to the pool. If the pool is at its cap the slab is
49	/// dropped instead so memory pressure does not grow without bound.
50	pub fn release(&self, slab: Arc<T>) {
51		let mut pool = self.pool.lock();
52		if pool.len() < self.cap {
53			pool.push(slab);
54		}
55	}
56}
57
58impl<T> Slab<T> {
59	/// Number of slabs currently held in the pool. Primarily for tests.
60	pub fn len(&self) -> usize {
61		self.pool.lock().len()
62	}
63
64	/// `true` if the pool is empty. Primarily for tests.
65	pub fn is_empty(&self) -> bool {
66		self.len() == 0
67	}
68}
69
70#[cfg(test)]
71mod tests {
72	use std::sync::Arc;
73
74	use super::Slab;
75
76	#[derive(Clone, Default, Debug)]
77	struct Buf {
78		bytes: Vec<u8>,
79	}
80
81	#[test]
82	fn acquire_from_empty_pool_allocates_fresh() {
83		let slab: Slab<Buf> = Slab::new(8);
84		let a = slab.acquire();
85		assert_eq!(Arc::strong_count(&a), 1);
86		assert_eq!(slab.len(), 0);
87	}
88
89	#[test]
90	fn release_then_acquire_returns_same_allocation() {
91		let slab: Slab<Buf> = Slab::new(8);
92		let mut a = slab.acquire();
93		Arc::make_mut(&mut a).bytes.extend_from_slice(b"hello");
94		let ptr_before = a.bytes.as_ptr();
95		slab.release(a);
96		assert_eq!(slab.len(), 1);
97
98		let b = slab.acquire();
99		// Same allocation reused (capacity preserved). The pool keeps
100		// the data as-is - callers are expected to overwrite via
101		// Arc::make_mut on the next use.
102		assert_eq!(b.bytes.as_ptr(), ptr_before);
103		assert_eq!(slab.len(), 0);
104	}
105
106	#[test]
107	fn shared_slabs_in_pool_are_skipped_on_acquire() {
108		let slab: Slab<Buf> = Slab::new(8);
109		let a = slab.acquire();
110		let _shadow = a.clone(); // bumps strong_count to 2
111		slab.release(a);
112		assert_eq!(slab.len(), 1);
113
114		// acquire skips the shared slab and allocates fresh.
115		let b = slab.acquire();
116		assert_eq!(Arc::strong_count(&b), 1);
117		assert!(slab.is_empty());
118	}
119
120	#[test]
121	fn release_at_cap_drops_overflow() {
122		let slab: Slab<Buf> = Slab::new(2);
123		let a = slab.acquire();
124		let b = slab.acquire();
125		let c = slab.acquire();
126		slab.release(a);
127		slab.release(b);
128		assert_eq!(slab.len(), 2);
129		// Third release exceeds cap; slab is dropped.
130		slab.release(c);
131		assert_eq!(slab.len(), 2);
132	}
133
134	#[test]
135	fn acquire_handles_pool_with_only_shared_slabs() {
136		let slab: Slab<Buf> = Slab::new(8);
137		let a = slab.acquire();
138		let b = slab.acquire();
139		let _shadow_a = a.clone();
140		let _shadow_b = b.clone();
141		slab.release(a);
142		slab.release(b);
143		assert_eq!(slab.len(), 2);
144
145		let c = slab.acquire();
146		assert_eq!(Arc::strong_count(&c), 1);
147		// Pool drained as the iterator skipped shared entries.
148		assert!(slab.is_empty());
149	}
150}
reifydb_core/util/slab.rs

reifydb_core/util/
slab.rs
