Skip to main content

diskann_utils/
object_pool.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! The [`ObjectPool<C>`] class is a thread-safe queue that allows concurrent access and is
7//! used for reusing allocated objects.
8//!
9//! # How this works
10//!
11//! The Pool struct can contain objects for which there is at least one implemention of the
12//! [`AsPooled`] or [`TryAsPooled`] trait.
13//!
14//! The [`AsPooled`] trait is used for types that can be pooled (i.e, created or modified)
15//! without supporting failures, while the [`TryAsPooled`] trait is used for types that can
16//! where creation or modification is fallible.
17
18use std::{
19    collections::VecDeque,
20    mem::ManuallyDrop,
21    ops::{Deref, DerefMut},
22    sync::{Arc, Mutex},
23};
24
25/// A thread-safe queue that allows concurrent access and is used for pooling items.
26#[derive(Debug)]
27pub struct ObjectPool<T> {
28    /// Queue of stored object.
29    queue: Mutex<VecDeque<T>>,
30
31    /// Maximum capacity of the pool. If pool is over maximum capacity, newly items will not
32    /// be pushed back to the pool.
33    capacity: Option<usize>,
34}
35
36impl<T> ObjectPool<T> {
37    /// Create an object pool consisting of `initial_size` object initialized using
38    /// [`AsPooled::create`].
39    ///
40    /// The argument `capacity` can be provided to place an upper bound on the number of
41    /// objects in the pool. If `initial_size > capacity`, the pool will be created with
42    /// excess items.
43    pub fn new<A>(args: A, initial_size: usize, capacity: Option<usize>) -> Self
44    where
45        T: AsPooled<A>,
46        A: Clone,
47    {
48        let queue = (0..initial_size).map(|_| T::create(args.clone())).collect();
49
50        Self {
51            queue: Mutex::new(queue),
52            capacity,
53        }
54    }
55
56    /// Create an empty [`ObjectPool`] with the requested capacity.
57    pub fn with_capacity(capacity: Option<usize>) -> Self {
58        Self {
59            queue: Mutex::new(VecDeque::with_capacity(capacity.unwrap_or(0))),
60            capacity,
61        }
62    }
63
64    /// Create an object pool consisting of `initial_size` object initialized using
65    /// [`TryAsPooled::try_create`].
66    ///
67    /// The argument `capacity` can be provided to place an upper bound on the number of
68    /// objects in the pool. If `initial_size > capacity`, the pool will be created with
69    /// excess items.
70    pub fn try_new<A>(
71        args: A,
72        initial_size: usize,
73        capacity: Option<usize>,
74    ) -> Result<Self, T::Error>
75    where
76        T: TryAsPooled<A>,
77        A: Clone,
78    {
79        let queue = (0..initial_size)
80            .map(|_| T::try_create(args.clone()))
81            .collect::<Result<_, _>>()?;
82        Ok(Self {
83            queue: Mutex::new(queue),
84            capacity,
85        })
86    }
87
88    /// Try to retrieve an object from the pool. If the pool is non-empty, then invoke
89    /// [`TryAsPooled::try_modify`] on the object with the given arguments. Otherwise, invoke
90    /// [`TryAsPooled::try_create`] to create a new object.
91    ///
92    /// If `try_modify` fails, the potentially modified object will not be returned to the
93    /// queue and will instead be dropped.
94    pub fn try_get_ref<A>(&self, args: A) -> Result<PooledRef<'_, T>, T::Error>
95    where
96        T: TryAsPooled<A>,
97    {
98        let item = self.try_get_or_create(args)?;
99        Ok(PooledRef {
100            item: ManuallyDrop::new(item),
101            parent: self,
102        })
103    }
104
105    /// Try to retrieve an object from the pool. If the pool is non-empty, then invoke
106    /// [`AsPooled::modify`] on the object with the given arguments. Otherwise, invoke
107    /// [`AsPooled::create`] to create a new object.
108    ///
109    /// Unlike [`ObjectPool::try_get_ref`], this interface should be infallible.
110    ///
111    /// If `modify` panics, the potentially modified object will not be returned to the
112    /// queue and will instead be dropped.
113    pub fn get_ref<A>(&self, args: A) -> PooledRef<'_, T>
114    where
115        T: AsPooled<A>,
116    {
117        let item = self.get_or_create(args);
118        PooledRef {
119            item: ManuallyDrop::new(item),
120            parent: self,
121        }
122    }
123
124    /// Try to retrieve an object from the pool. If the pool is non-empty, then invoke
125    /// [`TryAsPooled::try_modify`] on the object with the given arguments. Otherwise, invoke
126    /// [`TryAsPooled::try_create`] to create a new object.
127    ///
128    /// If `try_modify` fails, the potentially modified object will not be returned to the
129    /// queue and will instead be dropped.
130    pub fn try_get<A>(self: &Arc<Self>, args: A) -> Result<PooledArc<T>, T::Error>
131    where
132        T: TryAsPooled<A>,
133    {
134        let item = self.try_get_or_create(args)?;
135        Ok(PooledArc {
136            item: ManuallyDrop::new(item),
137            parent: self.clone(),
138        })
139    }
140
141    /// Try to retrieve an object from the pool. If the pool is non-empty, then invoke
142    /// [`AsPooled::modify`] on the object with the given arguments. Otherwise, invoke
143    /// [`AsPooled::create`] to create a new object.
144    ///
145    /// Unlike [`ObjectPool::try_get`], this interface should be infallible.
146    ///
147    /// If `modify` panics, the potentially modified object will not be returned to the
148    /// queue and will instead be dropped.
149    pub fn get<A>(self: &Arc<Self>, args: A) -> PooledArc<T>
150    where
151        T: AsPooled<A>,
152    {
153        let item = self.get_or_create(args);
154        PooledArc {
155            item: ManuallyDrop::new(item),
156            parent: self.clone(),
157        }
158    }
159
160    /// Return the number of items currently in the queue.
161    pub fn len(&self) -> usize {
162        self.lock().len()
163    }
164
165    /// Return whether or not the queue is empty.
166    pub fn is_empty(&self) -> bool {
167        self.len() == 0
168    }
169
170    //-----------------//
171    // Private Methods //
172    //-----------------//
173
174    fn try_get_or_create<A>(&self, args: A) -> Result<T, T::Error>
175    where
176        T: TryAsPooled<A>,
177    {
178        // Important: This needs to be on its own line instead of in the `if let` because
179        // the latter will extend the lifetime of the lock so it is held when `try_modify`
180        // or `try_create` is called.
181        let maybe = self.lock().pop_front();
182        if let Some(mut item) = maybe {
183            item.try_modify(args)?;
184            Ok(item)
185        } else {
186            T::try_create(args)
187        }
188    }
189
190    fn get_or_create<A>(&self, args: A) -> T
191    where
192        T: AsPooled<A>,
193    {
194        // Important: This needs to be on its own line instead of in the `if let` because the
195        // latter will extend the lifetime of the lock so it is held when `modify` or
196        // `create` is called.
197        let maybe = self.lock().pop_front();
198        if let Some(mut item) = maybe {
199            item.modify(args);
200            item
201        } else {
202            T::create(args)
203        }
204    }
205
206    fn lock(&self) -> std::sync::MutexGuard<'_, VecDeque<T>> {
207        match self.queue.lock() {
208            Ok(guard) => guard,
209            Err(poisoned) => {
210                // We trust the implementation of `VecDeque` to keep itself in a consistent
211                // state regardless of panics.
212                //
213                // We endeavor to only call non-panicking methods while holding the lock
214                // anyways.
215                //
216                // In particular, this means that implementations of `AsPooled` and
217                // `TryAsPooled` are allowed to panic, and we cannot call the associated
218                // methods while holding the lock.
219                self.queue.clear_poison();
220                poisoned.into_inner()
221            }
222        }
223    }
224}
225
226/// Attempt to retrieve a pooled object from an [`ObjectPool`], creating a new one if the
227/// queue is empty.
228///
229/// The goal of this trait is to modify an existing object if available such that it is
230/// indistinguishable semantically from a newly created object, allowing user code to be
231/// agnostic of the provenance of the created object.
232pub trait TryAsPooled<A>
233where
234    Self: Sized,
235{
236    /// Any error that can occur during creation or modification.
237    type Error;
238
239    /// Create an instance of `Self` from the argument types.
240    fn try_create(args: A) -> Result<Self, Self::Error>;
241
242    /// Modify an existing object so it behaves semantically identical to an object
243    /// constructed using [`Self::try_create`].
244    ///
245    /// This is often trickier to achieve than first anticipated.
246    ///
247    /// Note that it's up to the user to decide what "semantically identical" means. For
248    /// pooled objects like hash tables, the underlying capacity is not necessarily part of
249    /// the identity of an object - but there are contexts where it matters.
250    fn try_modify(&mut self, args: A) -> Result<(), Self::Error>;
251}
252
253/// Retrieve a pooled object from an [`ObjectPool`], creating a new one if the queue is empty.
254///
255/// The goal of this trait is to modify an existing object if available such that it is
256/// indistinguishable semantically from a newly created object, allowing user code to be
257/// agnostic of the provenance of the created object.
258pub trait AsPooled<A> {
259    /// Create an instance of `Self` from the argument types.
260    fn create(args: A) -> Self;
261
262    /// Modify an existing object so it behaves semantically identical to an object
263    /// constructed using [`Self::create`].
264    ///
265    /// This is often trickier to achieve than first anticipated.
266    ///
267    /// Note that it's up to the user to decide what "semantically identical" means. For
268    /// pooled objects like hash tables, the underlying capacity is not necessarily part of
269    /// the identity of an object - but there are contexts where it matters.
270    fn modify(&mut self, args: A);
271}
272
273/// An [`AsPooled`] initializer for `Vec<T>`.
274///
275/// Creates or modifies an existing array to be the specified length, without specifying
276/// the contents of the returned array.
277///
278/// This lack of specification means that existing `Vec`s that are modified to be the
279/// configured length will contain values from a previous run.
280#[derive(Debug, Clone, Copy)]
281pub struct Undef {
282    pub len: usize,
283}
284
285impl Undef {
286    /// Construct a new [`Undef`] vector initializer.
287    pub fn new(len: usize) -> Self {
288        Self { len }
289    }
290}
291
292impl<T> AsPooled<Undef> for Vec<T>
293where
294    T: Default + Clone,
295{
296    fn create(undef: Undef) -> Self {
297        vec![T::default(); undef.len]
298    }
299
300    fn modify(&mut self, undef: Undef) {
301        self.resize(undef.len, T::default())
302    }
303}
304
305/// [`PooledRef<'a, T>`] is a pooled object wrapper with a **reference** to its parent pool.
306///
307/// When dropped, the contained object will be returned to the pool if there is space.
308#[derive(Debug)]
309pub struct PooledRef<'a, T> {
310    item: ManuallyDrop<T>,
311    parent: &'a ObjectPool<T>,
312}
313
314impl<T> Drop for PooledRef<'_, T> {
315    fn drop(&mut self) {
316        let mut guard = self.parent.lock();
317        if guard.len() < self.parent.capacity.unwrap_or(usize::MAX) {
318            // SAFETY: We do not access self.item again after this.
319            guard.push_back(unsafe { ManuallyDrop::take(&mut self.item) });
320        } else {
321            // NOTE: The implementation of `T::drop` could panic. Release the lock first.
322            std::mem::drop(guard);
323
324            // SAFETY: We do not access self.item again after this.
325            unsafe { ManuallyDrop::drop(&mut self.item) };
326        }
327    }
328}
329
330impl<T> Deref for PooledRef<'_, T> {
331    type Target = T;
332
333    fn deref(&self) -> &Self::Target {
334        &self.item
335    }
336}
337
338impl<T> DerefMut for PooledRef<'_, T> {
339    fn deref_mut(&mut self) -> &mut Self::Target {
340        &mut self.item
341    }
342}
343
344/// [`PooledArc<T>`] is a pooled object wrapper with an **Arc** to its parent pool.
345///
346/// When dropped, the contained object will be returned to the pool if there is space.
347///
348/// Unlike [`PooledRef`], this object is `'static` (it `T: 'static`).
349#[derive(Debug)]
350pub struct PooledArc<T> {
351    item: ManuallyDrop<T>,
352    parent: Arc<ObjectPool<T>>,
353}
354
355impl<T> Drop for PooledArc<T> {
356    fn drop(&mut self) {
357        let mut guard = self.parent.lock();
358        if guard.len() < self.parent.capacity.unwrap_or(usize::MAX) {
359            // SAFETY: We do not access self.item again after this.
360            guard.push_back(unsafe { ManuallyDrop::take(&mut self.item) });
361        } else {
362            // NOTE: The implementation of `T::drop` could panic. Release the lock first.
363            std::mem::drop(guard);
364
365            // SAFETY: We do not access self.item again after this.
366            unsafe { ManuallyDrop::drop(&mut self.item) };
367        }
368    }
369}
370
371impl<T> Deref for PooledArc<T> {
372    type Target = T;
373
374    fn deref(&self) -> &Self::Target {
375        &self.item
376    }
377}
378
379impl<T> DerefMut for PooledArc<T> {
380    fn deref_mut(&mut self) -> &mut Self::Target {
381        &mut self.item
382    }
383}
384
385/// [`PoolOption<T>`] is an enum that can be either a non-pooled item or a pooled item.
386/// This is used to allow the user to choose between using a pooled item or opt out of pooling.
387#[derive(Debug)]
388pub enum PoolOption<T> {
389    NonPooled(T),
390    Pooled(PooledArc<T>),
391}
392
393impl<T> PoolOption<T> {
394    pub fn non_pooled(item: T) -> Self {
395        PoolOption::NonPooled(item)
396    }
397
398    pub fn try_non_pooled_create<A>(args: A) -> Result<Self, T::Error>
399    where
400        T: TryAsPooled<A>,
401    {
402        Ok(PoolOption::NonPooled(T::try_create(args)?))
403    }
404
405    pub fn non_pooled_create<A>(args: A) -> Self
406    where
407        T: AsPooled<A>,
408    {
409        PoolOption::NonPooled(T::create(args))
410    }
411
412    pub fn pooled<A>(pool: &Arc<ObjectPool<T>>, args: A) -> Self
413    where
414        T: AsPooled<A>,
415    {
416        PoolOption::Pooled(pool.get(args))
417    }
418
419    pub fn try_pooled<A>(pool: &Arc<ObjectPool<T>>, args: A) -> Result<Self, T::Error>
420    where
421        T: TryAsPooled<A>,
422    {
423        Ok(PoolOption::Pooled(pool.try_get(args)?))
424    }
425
426    pub fn is_pooled(&self) -> bool {
427        matches!(self, PoolOption::Pooled(_))
428    }
429
430    pub fn is_non_pooled(&self) -> bool {
431        matches!(self, PoolOption::NonPooled(_))
432    }
433}
434
435impl<T> Deref for PoolOption<T> {
436    type Target = T;
437
438    fn deref(&self) -> &Self::Target {
439        match self {
440            PoolOption::NonPooled(item) => item,
441            PoolOption::Pooled(item) => item,
442        }
443    }
444}
445
446impl<T> DerefMut for PoolOption<T> {
447    fn deref_mut(&mut self) -> &mut Self::Target {
448        match self {
449            PoolOption::NonPooled(item) => item,
450            PoolOption::Pooled(item) => item,
451        }
452    }
453}
454
455///////////
456// Tests //
457///////////
458
459#[cfg(test)]
460mod tests {
461    use super::*;
462
463    use crate::assert_contains;
464
465    #[derive(Debug)]
466    struct TestItem {
467        value: Box<u32>,
468        panic_on_drop: bool,
469    }
470
471    impl TestItem {
472        fn new(value: u32) -> Self {
473            Self {
474                value: Box::new(value),
475                panic_on_drop: false,
476            }
477        }
478    }
479
480    impl AsPooled<u32> for TestItem {
481        fn create(value: u32) -> Self {
482            TestItem::new(value)
483        }
484
485        fn modify(&mut self, value: u32) {
486            *self.value = value;
487            self.panic_on_drop = false;
488        }
489    }
490
491    impl TryAsPooled<i32> for TestItem {
492        type Error = ();
493
494        fn try_create(value: i32) -> Result<Self, Self::Error> {
495            match value.try_into() {
496                Ok(v) => Ok(TestItem::new(v)),
497                Err(_) => Err(()),
498            }
499        }
500
501        fn try_modify(&mut self, value: i32) -> Result<(), Self::Error> {
502            match value.try_into() {
503                Ok(v) => {
504                    *self.value = v;
505                    self.panic_on_drop = false;
506                    Ok(())
507                }
508                Err(_) => Err(()),
509            }
510        }
511    }
512
513    impl Drop for TestItem {
514        fn drop(&mut self) {
515            if self.panic_on_drop {
516                panic!("panicking on drop");
517            }
518        }
519    }
520
521    // A struct that panics on API calls to ensure we avoid Mutex poisoning.
522    struct TestPanic;
523
524    impl AsPooled<TestPanic> for TestItem {
525        fn create(_: TestPanic) -> Self {
526            panic!("panicking on create")
527        }
528
529        fn modify(&mut self, _: TestPanic) {
530            panic!("panicking on modify")
531        }
532    }
533
534    impl TryAsPooled<TestPanic> for TestItem {
535        type Error = ();
536
537        fn try_create(_: TestPanic) -> Result<Self, Self::Error> {
538            panic!("panicking on try_create")
539        }
540
541        fn try_modify(&mut self, _: TestPanic) -> Result<(), Self::Error> {
542            panic!("panicking on try_modify")
543        }
544    }
545
546    #[test]
547    fn test_pool_basic_tests() {
548        let pool = ObjectPool::<TestItem>::new(42, 2, None);
549        assert_eq!(pool.len(), 2);
550
551        let item1 = pool.get_ref(100);
552        assert_eq!(*item1.value, 100);
553        assert_eq!(pool.len(), 1);
554
555        let item2 = pool.get_ref(200);
556        assert_eq!(*item2.value, 200);
557        assert_eq!(pool.len(), 0);
558
559        let item = pool.get_ref(300);
560        assert_eq!(*item.value, 300);
561        assert_eq!(pool.len(), 0);
562        {
563            let item = pool.get_ref(400);
564            assert_eq!(*item.value, 400);
565            assert_eq!(pool.len(), 0);
566        }
567        assert_eq!(pool.len(), 1); // Pooled item is pushed back to the pool
568        {
569            let item_a = pool.get_ref(500);
570            assert_eq!(*item_a.value, 500);
571            assert_eq!(pool.len(), 0); // new item not yet returned to pool
572            let item_b = pool.get_ref(600);
573            assert_eq!(*item_b.value, 600);
574            assert_eq!(pool.len(), 0); // another new item not yet returned to pool
575        }
576        assert_eq!(pool.len(), 2); // Both Pooled items are pushed back to the pool
577
578        let pool = ObjectPool::<TestItem>::new(42, 1, None);
579        let item = pool.get_ref(100);
580        assert_eq!(*item.value, 100);
581    }
582
583    #[test]
584    fn test_pool_with_capacity() {
585        let pool = ObjectPool::<TestItem>::with_capacity(Some(4));
586        assert!(pool.is_empty());
587
588        let _ = pool.get_ref(100);
589        assert_eq!(pool.len(), 1);
590
591        {
592            let _v0 = pool.get_ref(100);
593            let _v1 = pool.get_ref(100);
594        }
595        assert_eq!(pool.len(), 2);
596
597        {
598            let _v0 = pool.get_ref(100);
599            let _v1 = pool.get_ref(100);
600            let _v2 = pool.get_ref(100);
601            let _v3 = pool.get_ref(100);
602            let _v4 = pool.get_ref(100);
603            let _v5 = pool.get_ref(100);
604        }
605        assert_eq!(
606            pool.len(),
607            4,
608            "pool should max out at the requested capacity."
609        );
610    }
611
612    #[test]
613    fn test_pool_basic_tests_with_try() {
614        // Create a pool with negative initial size to test error case
615        let pool_result = ObjectPool::<TestItem>::try_new(-1, 2, Some(100));
616        assert!(
617            pool_result.is_err(),
618            "Pool creation should fail with negative args"
619        );
620
621        let pool = ObjectPool::<TestItem>::try_new(42, 2, None).unwrap();
622        assert_eq!(pool.len(), 2);
623        let item1 = pool.try_get_ref(100).unwrap();
624        assert_eq!(*item1.value, 100);
625        assert_eq!(pool.len(), 1);
626        let item2 = pool.try_get_ref(200).unwrap();
627        assert_eq!(*item2.value, 200);
628        assert_eq!(pool.len(), 0);
629        let item = pool.try_get_ref(300).unwrap();
630        assert_eq!(*item.value, 300);
631        assert_eq!(pool.len(), 0);
632        {
633            let item = pool.try_get_ref(400).unwrap();
634            assert_eq!(*item.value, 400);
635            assert_eq!(pool.len(), 0);
636        }
637        assert_eq!(pool.len(), 1); // Pooled item is pushed back to the pool
638        {
639            let item_a = pool.try_get_ref(500).unwrap();
640            assert_eq!(*item_a.value, 500);
641            assert_eq!(pool.len(), 0); // new item not yet returned to pool
642            let item_b = pool.try_get_ref(600).unwrap();
643            assert_eq!(*item_b.value, 600);
644            assert_eq!(pool.len(), 0); // another new item not yet returned to pool
645        }
646        assert_eq!(pool.len(), 2); // Both Pooled items are pushed back to the pool
647
648        let pool = ObjectPool::<TestItem>::try_new(42, 1, Some(100)).unwrap();
649        let item = pool.try_get_ref(100).unwrap();
650        assert_eq!(*item.value, 100);
651    }
652
653    #[test]
654    fn test_pool_with_arc() {
655        let pool = &Arc::new(ObjectPool::<TestItem>::new(42, 1, None));
656        let item = pool.get(100);
657        assert_eq!(*item.value, 100);
658        assert_eq!(pool.len(), 0);
659
660        let item = pool.get(200);
661        assert_eq!(*item.value, 200);
662        assert_eq!(pool.len(), 0);
663        {
664            let item = pool.get(400);
665            assert_eq!(*item.value, 400);
666            assert_eq!(pool.len(), 0);
667        }
668        assert_eq!(pool.len(), 1); // Pooled item is pushed back to the pool
669        let item = pool.try_get_ref(400).unwrap();
670        assert_eq!(*item.value, 400);
671        assert_eq!(pool.len(), 0); // new item not yet returned to pool
672        let item = pool.try_get(500).unwrap();
673        assert_eq!(*item.value, 500);
674    }
675
676    #[test]
677    fn test_pool_max_capacity_ref() {
678        let pool = ObjectPool::<TestItem>::new(42, 1, Some(1));
679        assert_eq!(pool.len(), 1);
680        assert!(!pool.is_empty());
681        assert_eq!(pool.len(), pool.capacity.unwrap()); // size is at max_capacity
682        {
683            let item = pool.get_ref(100);
684            assert_eq!(pool.len(), 0); // item is removed from the pool
685            assert!(pool.is_empty());
686            assert!(pool.len() < pool.capacity.unwrap()); // size is less than max_capacity
687            assert_eq!(*item.value, 100);
688        }
689        assert_eq!(pool.len(), 1); // item is pushed back to the pool
690        assert_eq!(pool.len(), pool.capacity.unwrap()); // size is at max_capacity
691        {
692            let item1 = pool.get_ref(100);
693            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
694            let item2 = pool.get_ref(200);
695            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
696            let item3 = pool.get_ref(300);
697            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
698            assert!(*item1.value == 100 && *item2.value == 200 && *item3.value == 300);
699            // all items are alive
700        }
701        assert_eq!(pool.len(), pool.capacity.unwrap()); // max_capacity is not exceeded
702        assert_eq!(pool.len(), 1); // at most max_capacity items are in the pool
703    }
704
705    #[test]
706    fn test_pool_max_capacity_pooled_item() {
707        let pool = &Arc::new(ObjectPool::<TestItem>::new(42, 1, Some(1)));
708        assert_eq!(pool.len(), 1);
709        assert_eq!(pool.len(), pool.capacity.unwrap()); // size is at max_capacity
710        {
711            let item = pool.get(100);
712            assert_eq!(pool.len(), 0); // item is removed from the pool
713            assert!(pool.len() < pool.capacity.unwrap()); // size is less than max_capacity
714            assert_eq!(*item.value, 100);
715        }
716        assert_eq!(pool.len(), 1); // item is pushed back to the pool
717        assert_eq!(pool.len(), pool.capacity.unwrap()); // size is at max_capacity
718        {
719            let item1 = pool.get(100);
720            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
721            let item2 = pool.get(200);
722            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
723            let item3 = pool.get(300);
724            assert_eq!(pool.len(), 0); // item is not pushed back to the pool
725            assert!(*item1.value == 100 && *item2.value == 200 && *item3.value == 300);
726            // all items are alive
727        }
728        assert_eq!(pool.len(), pool.capacity.unwrap()); // max_capacity is not exceeded
729        assert_eq!(pool.len(), 1); // at most max_capacity items are in the pool
730    }
731
732    #[test]
733    fn test_pool_options() {
734        // Test non_pooled and non_pooled_create
735        let item = PoolOption::non_pooled(TestItem::new(42));
736        assert_eq!(*item.value, 42);
737        assert!(item.is_non_pooled());
738        assert!(!item.is_pooled());
739
740        let item = PoolOption::<TestItem>::non_pooled_create(100);
741        assert_eq!(*item.value, 100);
742        assert!(item.is_non_pooled());
743
744        // Test pooled
745        let pool = Arc::new(ObjectPool::<TestItem>::new(42, 1, None));
746        let item = PoolOption::pooled(&pool, 100);
747        assert_eq!(*item.value, 100);
748        assert!(item.is_pooled());
749        assert!(!item.is_non_pooled());
750
751        // Test try_pooled and try_non_pooled_create
752        let item = PoolOption::try_pooled(&pool, 100).unwrap();
753        assert_eq!(*item.value, 100);
754        assert!(item.is_pooled());
755        let item = PoolOption::<TestItem>::try_non_pooled_create(200).unwrap();
756        assert_eq!(*item.value, 200);
757        assert!(item.is_non_pooled());
758        assert!(!item.is_pooled());
759        let item_result = PoolOption::<TestItem>::try_non_pooled_create(-200);
760        assert!(
761            item_result.is_err(),
762            "Creating non-pooled item with negative args should fail"
763        );
764    }
765
766    #[test]
767    fn test_pool_ref_deref_mut() {
768        // Test PooledRef deref and deref_mut
769        let pool = ObjectPool::<TestItem>::new(42, 1, None);
770        let item = pool.get_ref(100);
771
772        assert_eq!(*item.value, 100);
773        let mut item = pool.get_ref(100);
774        *item.value = 200;
775        assert_eq!(*item.value, 200);
776
777        let item_ref: &TestItem = &item;
778        assert_eq!(*item_ref.value, 200);
779
780        let item_ref_mut: &mut TestItem = &mut item;
781        assert_eq!(*item_ref_mut.value, 200);
782
783        *item_ref_mut.value = 300;
784        assert_eq!(*item_ref_mut.value, 300);
785
786        // Test PooledArc deref and deref_mut
787        let pool = &Arc::new(ObjectPool::<TestItem>::new(42, 1, None));
788        let mut item = pool.get_ref(100);
789        assert_eq!(*item.value, 100);
790
791        *item.value = 200;
792        assert_eq!(*item.value, 200);
793
794        // Test PoolOption deref and deref_mut for both pooled and non-pooled items
795        let pool = Arc::new(ObjectPool::<TestItem>::new(42, 1, None));
796        let mut item = PoolOption::pooled(&pool, 100);
797        assert_eq!(*item.value, 100);
798
799        *item.value = 200;
800        assert_eq!(*item.value, 200);
801
802        let mut item = PoolOption::non_pooled(TestItem::new(42));
803        assert_eq!(*item.value, 42);
804
805        *item.value = 100;
806        assert_eq!(*item.value, 100);
807    }
808
809    //------------------//
810    // Panic Resiliance //
811    //------------------//
812
813    // Here - we test that panics during:
814    //
815    // * Object Creation
816    // * Object Modification
817    // * Object destruction (when the queue is full)
818    //
819    // Do not poison the central mutex.
820    //
821    // The goal is to test code around anything that locks the mutex to ensure it only
822    // invokes non-panicking code, and any code that can potentially panic is run outside
823    // of the lock.
824
825    // Check that the panicking payload is castable to a `&'static str` and that it
826    // contains `msg`.
827    fn check_error(err: &dyn std::any::Any, contains: &str) {
828        match err.downcast_ref::<&'static str>() {
829            Some(msg) => assert_contains!(msg, contains),
830            None => panic!("incorrect downcast type"),
831        }
832    }
833
834    #[test]
835    fn test_panic_during_create() {
836        let pool = ObjectPool::<TestItem>::new(0u32, 0, Some(1));
837
838        // Panic during `create`.
839        let err = std::panic::catch_unwind(|| {
840            let _ = pool.get_ref(TestPanic);
841        })
842        .unwrap_err();
843
844        check_error(&*err, "panicking on create");
845
846        assert!(
847            !pool.queue.is_poisoned(),
848            "lock should be released while calling trait implementations"
849        );
850
851        assert_eq!(pool.len(), 0);
852    }
853
854    #[test]
855    fn test_panic_during_try_create() {
856        let pool = ObjectPool::<TestItem>::new(0u32, 0, Some(1));
857
858        // Panic during `try_create`.
859        let err = std::panic::catch_unwind(|| {
860            let _ = pool.try_get_ref(TestPanic);
861        })
862        .unwrap_err();
863
864        check_error(&*err, "panicking on try_create");
865
866        assert!(
867            !pool.queue.is_poisoned(),
868            "lock should be released while calling trait implementations"
869        );
870
871        assert_eq!(pool.len(), 0);
872    }
873
874    #[test]
875    fn test_panic_during_modify() {
876        let pool = ObjectPool::<TestItem>::new(0u32, 0, Some(1));
877
878        // Append a new item into the pool so it is full.
879        let _ = pool.get_ref(0u32);
880        assert_eq!(pool.len(), 1);
881
882        let err = std::panic::catch_unwind(|| {
883            let _ = pool.get_ref(TestPanic);
884        })
885        .unwrap_err();
886
887        check_error(&*err, "panicking on modify");
888
889        assert!(
890            !pool.queue.is_poisoned(),
891            "lock should be released while calling trait implementations"
892        );
893
894        assert_eq!(
895            pool.len(),
896            0,
897            "we should not return a potentially torn object to the pool"
898        );
899    }
900
901    #[test]
902    fn test_panic_during_try_modify() {
903        let pool = ObjectPool::<TestItem>::new(0u32, 0, Some(1));
904
905        // Append a new item into the pool so it is full.
906        let _ = pool.get_ref(0u32);
907        assert_eq!(pool.len(), 1);
908
909        let err = std::panic::catch_unwind(|| {
910            let _ = pool.try_get_ref(TestPanic);
911        })
912        .unwrap_err();
913
914        check_error(&*err, "panicking on try_modify");
915
916        assert!(
917            !pool.queue.is_poisoned(),
918            "lock should be released while calling trait implementations"
919        );
920
921        assert_eq!(
922            pool.len(),
923            0,
924            "we should not return a potentially torn object to the pool"
925        );
926    }
927
928    // Panic on full drop - `ref` interface.
929    #[test]
930    fn test_panic_during_drop_ref() {
931        let pool = ObjectPool::<TestItem>::new(0u32, 0, Some(1));
932
933        // This sequence ensures that when we try to put the object behind `a` back
934        // into the queue that it is dropped since the queue is full.
935        let mut a = pool.get_ref(0u32);
936        let _ = pool.get_ref(1u32);
937        assert_eq!(pool.len(), 1);
938
939        a.panic_on_drop = true;
940        let err = std::panic::catch_unwind(move || std::mem::drop(a)).unwrap_err();
941        check_error(&*err, "panicking on drop");
942
943        assert!(
944            !pool.queue.is_poisoned(),
945            "lock should be released while calling object drop"
946        );
947        assert_eq!(pool.len(), 1);
948    }
949
950    // Panic on full drop - `Arc` interface.
951    #[test]
952    fn test_panic_during_drop_arc() {
953        let pool = Arc::new(ObjectPool::<TestItem>::new(0u32, 0, Some(1)));
954
955        // This sequence ensures that when we try to put the object behind `a` back
956        // into the queue that it is dropped since the queue is full.
957        let mut a = pool.get(0u32);
958        let _ = pool.get(1u32);
959        assert_eq!(pool.len(), 1);
960
961        a.panic_on_drop = true;
962        let err = std::panic::catch_unwind(move || std::mem::drop(a)).unwrap_err();
963        check_error(&*err, "panicking on drop");
964
965        assert!(
966            !pool.queue.is_poisoned(),
967            "lock should be released while calling object drop"
968        );
969        assert_eq!(pool.len(), 1);
970    }
971
972    // This test uses private member access to induce a poisoned state in the mutex.
973    #[test]
974    fn test_panic_recovery() {
975        let pool = ObjectPool::<TestItem>::new(0u32, 1, Some(1));
976
977        let err = std::panic::catch_unwind(|| {
978            let _guard = pool.queue.lock();
979            panic!("yeet");
980        })
981        .unwrap_err();
982
983        check_error(&*err, "yeet");
984
985        assert!(pool.queue.is_poisoned());
986
987        let _ = pool.get_ref(1u32);
988        assert!(!pool.queue.is_poisoned(), "poison should be cleared");
989    }
990
991    //-------//
992    // Undef //
993    //-------//
994
995    #[test]
996    fn test_undef() {
997        let mut x: Vec<f32> = Vec::<f32>::create(Undef::new(10));
998        assert_eq!(x.len(), 10);
999
1000        x.modify(Undef::new(0));
1001        assert_eq!(x.len(), 0);
1002
1003        x.modify(Undef::new(20));
1004        assert_eq!(x.len(), 20);
1005    }
1006}