Skip to main content

reifydb_core/util/
slab.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 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		// The slab returns the allocation with its contents intact; callers overwrite through
80		// Arc::make_mut rather than relying on it being cleared.
81		assert_eq!(b.bytes.as_ptr(), ptr_before);
82		assert_eq!(slab.len(), 0);
83	}
84
85	#[test]
86	fn shared_slabs_in_pool_are_skipped_on_acquire() {
87		let slab: Slab<Buf> = Slab::new(8);
88		let a = slab.acquire();
89		let _shadow = a.clone(); // bumps strong_count to 2
90		slab.release(a);
91		assert_eq!(slab.len(), 1);
92
93		let b = slab.acquire();
94		assert_eq!(Arc::strong_count(&b), 1);
95		assert!(slab.is_empty());
96	}
97
98	#[test]
99	fn release_at_cap_drops_overflow() {
100		let slab: Slab<Buf> = Slab::new(2);
101		let a = slab.acquire();
102		let b = slab.acquire();
103		let c = slab.acquire();
104		slab.release(a);
105		slab.release(b);
106		assert_eq!(slab.len(), 2);
107		slab.release(c);
108		assert_eq!(slab.len(), 2);
109	}
110
111	#[test]
112	fn acquire_handles_pool_with_only_shared_slabs() {
113		let slab: Slab<Buf> = Slab::new(8);
114		let a = slab.acquire();
115		let b = slab.acquire();
116		let _shadow_a = a.clone();
117		let _shadow_b = b.clone();
118		slab.release(a);
119		slab.release(b);
120		assert_eq!(slab.len(), 2);
121
122		let c = slab.acquire();
123		assert_eq!(Arc::strong_count(&c), 1);
124		// Skipped shared entries are dropped rather than put back, so the pool drains.
125		assert!(slab.is_empty());
126	}
127}