/home/noah/src/trueno/src/brick/resource_pool.rs
Line | Count | Source |
1 | | //! Resource Pool with Semaphore |
2 | | //! |
3 | | //! AWP-05: Semaphore-based resource pool for managing limited resources. |
4 | | |
5 | | use std::sync::atomic::{AtomicUsize, Ordering}; |
6 | | use std::sync::Mutex; |
7 | | |
8 | | // ---------------------------------------------------------------------------- |
9 | | // AWP-05: Resource Pool with Semaphore |
10 | | // ---------------------------------------------------------------------------- |
11 | | |
12 | | /// Semaphore-based resource pool. |
13 | | /// |
14 | | /// # Example |
15 | | /// ```rust |
16 | | /// use trueno::brick::ResourcePool; |
17 | | /// |
18 | | /// let pool: ResourcePool<Vec<u8>> = ResourcePool::new(4, || Vec::with_capacity(1024)); |
19 | | /// |
20 | | /// // Acquire resources (up to max) |
21 | | /// let r1 = pool.try_acquire().unwrap(); |
22 | | /// let r2 = pool.try_acquire().unwrap(); |
23 | | /// let r3 = pool.try_acquire().unwrap(); |
24 | | /// let r4 = pool.try_acquire().unwrap(); |
25 | | /// |
26 | | /// // Pool is exhausted |
27 | | /// assert!(pool.try_acquire().is_none()); |
28 | | /// |
29 | | /// // Release one |
30 | | /// drop(r1); |
31 | | /// |
32 | | /// // Now we can acquire again |
33 | | /// assert!(pool.try_acquire().is_some()); |
34 | | /// ``` |
35 | | pub struct ResourcePool<T> { |
36 | | /// Maximum concurrent resources. |
37 | | max_resources: usize, |
38 | | /// Available permits. |
39 | | available: AtomicUsize, |
40 | | /// Pooled resources. |
41 | | resources: Mutex<Vec<T>>, |
42 | | /// Factory for creating new resources. |
43 | | factory: Box<dyn Fn() -> T + Send + Sync>, |
44 | | } |
45 | | |
46 | | impl<T> ResourcePool<T> { |
47 | | /// Create a new resource pool. |
48 | 0 | pub fn new(max_resources: usize, factory: impl Fn() -> T + Send + Sync + 'static) -> Self { |
49 | 0 | Self { |
50 | 0 | max_resources, |
51 | 0 | available: AtomicUsize::new(max_resources), |
52 | 0 | resources: Mutex::new(Vec::with_capacity(max_resources)), |
53 | 0 | factory: Box::new(factory), |
54 | 0 | } |
55 | 0 | } |
56 | | |
57 | | /// Get the maximum number of resources. |
58 | | #[must_use] |
59 | 0 | pub fn max_resources(&self) -> usize { |
60 | 0 | self.max_resources |
61 | 0 | } |
62 | | |
63 | | /// Get the number of available permits. |
64 | | #[must_use] |
65 | 0 | pub fn available(&self) -> usize { |
66 | 0 | self.available.load(Ordering::Acquire) |
67 | 0 | } |
68 | | |
69 | | /// Try to acquire a resource (non-blocking). |
70 | 0 | pub fn try_acquire(&self) -> Option<PooledResource<'_, T>> { |
71 | | // Try to get a permit |
72 | | loop { |
73 | 0 | let current = self.available.load(Ordering::Acquire); |
74 | 0 | if current == 0 { |
75 | 0 | return None; |
76 | 0 | } |
77 | 0 | if self |
78 | 0 | .available |
79 | 0 | .compare_exchange(current, current - 1, Ordering::AcqRel, Ordering::Relaxed) |
80 | 0 | .is_ok() |
81 | | { |
82 | 0 | break; |
83 | 0 | } |
84 | | } |
85 | | |
86 | | // Get or create resource |
87 | 0 | let resource = { |
88 | 0 | let mut pool = self.resources.lock().unwrap(); |
89 | 0 | pool.pop().unwrap_or_else(|| (self.factory)()) |
90 | | }; |
91 | | |
92 | 0 | Some(PooledResource { |
93 | 0 | resource: Some(resource), |
94 | 0 | pool: self, |
95 | 0 | }) |
96 | 0 | } |
97 | | |
98 | 0 | fn release(&self, resource: T) { |
99 | | { |
100 | 0 | let mut pool = self.resources.lock().unwrap(); |
101 | 0 | if pool.len() < self.max_resources { |
102 | 0 | pool.push(resource); |
103 | 0 | } |
104 | | // else: drop resource (pool is full) |
105 | | } |
106 | 0 | self.available.fetch_add(1, Ordering::Release); |
107 | 0 | } |
108 | | } |
109 | | |
110 | | impl<T: std::fmt::Debug> std::fmt::Debug for ResourcePool<T> { |
111 | 0 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
112 | 0 | f.debug_struct("ResourcePool") |
113 | 0 | .field("max_resources", &self.max_resources) |
114 | 0 | .field("available", &self.available()) |
115 | 0 | .finish() |
116 | 0 | } |
117 | | } |
118 | | |
119 | | /// A resource acquired from a pool. |
120 | | pub struct PooledResource<'a, T> { |
121 | | resource: Option<T>, |
122 | | pool: &'a ResourcePool<T>, |
123 | | } |
124 | | |
125 | | impl<T> std::ops::Deref for PooledResource<'_, T> { |
126 | | type Target = T; |
127 | 0 | fn deref(&self) -> &T { |
128 | 0 | self.resource.as_ref().unwrap() |
129 | 0 | } |
130 | | } |
131 | | |
132 | | impl<T> std::ops::DerefMut for PooledResource<'_, T> { |
133 | 0 | fn deref_mut(&mut self) -> &mut T { |
134 | 0 | self.resource.as_mut().unwrap() |
135 | 0 | } |
136 | | } |
137 | | |
138 | | impl<T> Drop for PooledResource<'_, T> { |
139 | 0 | fn drop(&mut self) { |
140 | 0 | if let Some(resource) = self.resource.take() { |
141 | 0 | self.pool.release(resource); |
142 | 0 | } |
143 | 0 | } |
144 | | } |
145 | | |
146 | | #[cfg(test)] |
147 | | mod tests { |
148 | | use super::*; |
149 | | |
150 | | #[test] |
151 | | fn test_resource_pool_new() { |
152 | | let pool: ResourcePool<Vec<u8>> = ResourcePool::new(4, Vec::new); |
153 | | assert_eq!(pool.max_resources(), 4); |
154 | | assert_eq!(pool.available(), 4); |
155 | | } |
156 | | |
157 | | #[test] |
158 | | fn test_resource_pool_acquire_release() { |
159 | | let pool: ResourcePool<u32> = ResourcePool::new(2, || 0); |
160 | | |
161 | | let r1 = pool.try_acquire(); |
162 | | assert!(r1.is_some()); |
163 | | assert_eq!(pool.available(), 1); |
164 | | |
165 | | let r2 = pool.try_acquire(); |
166 | | assert!(r2.is_some()); |
167 | | assert_eq!(pool.available(), 0); |
168 | | |
169 | | // Pool exhausted |
170 | | let r3 = pool.try_acquire(); |
171 | | assert!(r3.is_none()); |
172 | | |
173 | | // Release one |
174 | | drop(r1); |
175 | | assert_eq!(pool.available(), 1); |
176 | | |
177 | | // Can acquire again |
178 | | let r4 = pool.try_acquire(); |
179 | | assert!(r4.is_some()); |
180 | | } |
181 | | |
182 | | #[test] |
183 | | fn test_resource_pool_factory_called() { |
184 | | use std::sync::atomic::AtomicUsize; |
185 | | use std::sync::Arc; |
186 | | |
187 | | let call_count = Arc::new(AtomicUsize::new(0)); |
188 | | let cc = Arc::clone(&call_count); |
189 | | let pool: ResourcePool<u32> = |
190 | | ResourcePool::new(2, move || cc.fetch_add(1, Ordering::SeqCst) as u32); |
191 | | |
192 | | // First acquire creates a resource |
193 | | let _r1 = pool.try_acquire().unwrap(); |
194 | | assert_eq!(call_count.load(Ordering::SeqCst), 1); |
195 | | |
196 | | // Second acquire creates another resource |
197 | | let _r2 = pool.try_acquire().unwrap(); |
198 | | assert_eq!(call_count.load(Ordering::SeqCst), 2); |
199 | | } |
200 | | |
201 | | #[test] |
202 | | fn test_resource_pool_reuse() { |
203 | | use std::sync::atomic::AtomicUsize; |
204 | | use std::sync::Arc; |
205 | | |
206 | | let call_count = Arc::new(AtomicUsize::new(0)); |
207 | | let cc = Arc::clone(&call_count); |
208 | | let pool: ResourcePool<u32> = |
209 | | ResourcePool::new(2, move || cc.fetch_add(1, Ordering::SeqCst) as u32); |
210 | | |
211 | | // Acquire and release |
212 | | let r1 = pool.try_acquire().unwrap(); |
213 | | assert_eq!(call_count.load(Ordering::SeqCst), 1); |
214 | | drop(r1); |
215 | | |
216 | | // Acquire again - should reuse existing resource |
217 | | let _r2 = pool.try_acquire().unwrap(); |
218 | | assert_eq!(call_count.load(Ordering::SeqCst), 1); // No new resource created |
219 | | } |
220 | | |
221 | | #[test] |
222 | | fn test_pooled_resource_deref() { |
223 | | let pool: ResourcePool<Vec<u8>> = ResourcePool::new(1, || vec![1, 2, 3]); |
224 | | |
225 | | let resource = pool.try_acquire().unwrap(); |
226 | | assert_eq!(resource.len(), 3); |
227 | | assert_eq!(&*resource, &[1, 2, 3]); |
228 | | } |
229 | | |
230 | | #[test] |
231 | | fn test_pooled_resource_deref_mut() { |
232 | | let pool: ResourcePool<Vec<u8>> = ResourcePool::new(1, Vec::new); |
233 | | |
234 | | let mut resource = pool.try_acquire().unwrap(); |
235 | | resource.push(42); |
236 | | assert_eq!(resource.len(), 1); |
237 | | } |
238 | | |
239 | | #[test] |
240 | | fn test_resource_pool_debug() { |
241 | | let pool: ResourcePool<u32> = ResourcePool::new(4, || 0); |
242 | | let debug = format!("{:?}", pool); |
243 | | assert!(debug.contains("ResourcePool")); |
244 | | assert!(debug.contains("max_resources: 4")); |
245 | | assert!(debug.contains("available: 4")); |
246 | | } |
247 | | |
248 | | /// FALSIFICATION TEST: Verify pool never exceeds max resources |
249 | | /// |
250 | | /// Even under concurrent pressure, the pool must never hand out |
251 | | /// more resources than max_resources. |
252 | | #[test] |
253 | | fn test_falsify_pool_never_exceeds_max() { |
254 | | use std::sync::Arc; |
255 | | |
256 | | let pool = Arc::new(ResourcePool::<u32>::new(5, || 0)); |
257 | | let mut handles = vec![]; |
258 | | |
259 | | // Spawn threads that aggressively try to acquire |
260 | | for _ in 0..10 { |
261 | | let pool_clone = Arc::clone(&pool); |
262 | | handles.push(std::thread::spawn(move || { |
263 | | let mut acquired = vec![]; |
264 | | for _ in 0..100 { |
265 | | if let Some(r) = pool_clone.try_acquire() { |
266 | | acquired.push(r); |
267 | | } |
268 | | // Random sleep to increase contention |
269 | | if acquired.len() > 2 { |
270 | | acquired.pop(); |
271 | | } |
272 | | } |
273 | | acquired.len() |
274 | | })); |
275 | | } |
276 | | |
277 | | for handle in handles { |
278 | | let _ = handle.join(); |
279 | | } |
280 | | |
281 | | // After all threads complete, available should equal max |
282 | | // (all resources released) |
283 | | // Note: This just verifies the pool is consistent |
284 | | let final_available = pool.available(); |
285 | | assert!( |
286 | | final_available <= 5, |
287 | | "FALSIFICATION FAILED: available ({}) > max (5)", |
288 | | final_available |
289 | | ); |
290 | | } |
291 | | |
292 | | /// FALSIFICATION TEST: Verify CAS prevents double-acquire |
293 | | /// |
294 | | /// The compare-exchange loop must prevent two threads from |
295 | | /// acquiring the same permit. |
296 | | #[test] |
297 | | fn test_falsify_cas_prevents_double_acquire() { |
298 | | use std::sync::atomic::AtomicUsize; |
299 | | use std::sync::Arc; |
300 | | |
301 | | let pool = Arc::new(ResourcePool::<u32>::new(1, || 0)); |
302 | | let acquired_count = Arc::new(AtomicUsize::new(0)); |
303 | | let mut handles = vec![]; |
304 | | |
305 | | // Multiple threads try to acquire the single resource simultaneously |
306 | | for _ in 0..10 { |
307 | | let pool_clone = Arc::clone(&pool); |
308 | | let count_clone = Arc::clone(&acquired_count); |
309 | | handles.push(std::thread::spawn(move || { |
310 | | for _ in 0..100 { |
311 | | if let Some(_r) = pool_clone.try_acquire() { |
312 | | let prev = count_clone.fetch_add(1, Ordering::SeqCst); |
313 | | // Only one thread should hold the resource at a time |
314 | | assert!( |
315 | | prev == 0, |
316 | | "FALSIFICATION FAILED: Multiple threads acquired simultaneously" |
317 | | ); |
318 | | std::thread::yield_now(); |
319 | | count_clone.fetch_sub(1, Ordering::SeqCst); |
320 | | } |
321 | | } |
322 | | })); |
323 | | } |
324 | | |
325 | | for handle in handles { |
326 | | handle.join().unwrap(); |
327 | | } |
328 | | } |
329 | | } |