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