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
8pub struct Slab<T> {
9	pool: Mutex<Vec<Arc<T>>>,
10	cap: usize,
11}
12
13impl<T: Default> Slab<T> {
14	pub fn new(cap: usize) -> Self {
15		Self {
16			pool: Mutex::new(Vec::new()),
17			cap,
18		}
19	}
20
21	pub fn acquire(&self) -> Arc<T> {
22		let mut pool = self.pool.lock();
23		while let Some(slab) = pool.pop() {
24			if Arc::strong_count(&slab) == 1 {
25				return slab;
26			}
27		}
28		drop(pool);
29		Arc::new(T::default())
30	}
31
32	pub fn release(&self, slab: Arc<T>) {
33		let mut pool = self.pool.lock();
34		if pool.len() < self.cap {
35			pool.push(slab);
36		}
37	}
38}
39
40impl<T> Slab<T> {
41	pub fn len(&self) -> usize {
42		self.pool.lock().len()
43	}
44
45	pub fn is_empty(&self) -> bool {
46		self.len() == 0
47	}
48}
49
50#[cfg(test)]
51mod tests {
52	use std::sync::Arc;
53
54	use super::Slab;
55
56	#[derive(Clone, Default, Debug)]
57	struct Buf {
58		bytes: Vec<u8>,
59	}
60
61	#[test]
62	fn acquire_from_empty_pool_allocates_fresh() {
63		let slab: Slab<Buf> = Slab::new(8);
64		let a = slab.acquire();
65		assert_eq!(Arc::strong_count(&a), 1);
66		assert_eq!(slab.len(), 0);
67	}
68
69	#[test]
70	fn release_then_acquire_returns_same_allocation() {
71		let slab: Slab<Buf> = Slab::new(8);
72		let mut a = slab.acquire();
73		Arc::make_mut(&mut a).bytes.extend_from_slice(b"hello");
74		let ptr_before = a.bytes.as_ptr();
75		slab.release(a);
76		assert_eq!(slab.len(), 1);
77
78		let b = slab.acquire();
79		// Same allocation reused (capacity preserved). The pool keeps
80		// the data as-is - callers are expected to overwrite via
81		// Arc::make_mut on the next use.
82		assert_eq!(b.bytes.as_ptr(), ptr_before);
83		assert_eq!(slab.len(), 0);
84	}
85
86	#[test]
87	fn shared_slabs_in_pool_are_skipped_on_acquire() {
88		let slab: Slab<Buf> = Slab::new(8);
89		let a = slab.acquire();
90		let _shadow = a.clone(); // bumps strong_count to 2
91		slab.release(a);
92		assert_eq!(slab.len(), 1);
93
94		// acquire skips the shared slab and allocates fresh.
95		let b = slab.acquire();
96		assert_eq!(Arc::strong_count(&b), 1);
97		assert!(slab.is_empty());
98	}
99
100	#[test]
101	fn release_at_cap_drops_overflow() {
102		let slab: Slab<Buf> = Slab::new(2);
103		let a = slab.acquire();
104		let b = slab.acquire();
105		let c = slab.acquire();
106		slab.release(a);
107		slab.release(b);
108		assert_eq!(slab.len(), 2);
109		// Third release exceeds cap; slab is dropped.
110		slab.release(c);
111		assert_eq!(slab.len(), 2);
112	}
113
114	#[test]
115	fn acquire_handles_pool_with_only_shared_slabs() {
116		let slab: Slab<Buf> = Slab::new(8);
117		let a = slab.acquire();
118		let b = slab.acquire();
119		let _shadow_a = a.clone();
120		let _shadow_b = b.clone();
121		slab.release(a);
122		slab.release(b);
123		assert_eq!(slab.len(), 2);
124
125		let c = slab.acquire();
126		assert_eq!(Arc::strong_count(&c), 1);
127		// Pool drained as the iterator skipped shared entries.
128		assert!(slab.is_empty());
129	}
130}
reifydb_core/util/slab.rs

reifydb_core/util/
slab.rs
