sized_object_pool/
lib.rs

1use dynamic_pool::DynamicPool;
2pub use dynamic_pool::{DynamicPoolItem, DynamicReset};
3use thiserror::Error;
4
5#[derive(Error, Debug)]
6pub enum SizedPoolError {
7    #[error("the given size exceeds the maximum allowed size of the pool")]
8    SizeExceedMaxSize,
9}
10
11pub trait SizedAllocatable {
12    fn new(size: usize) -> Self;
13    fn size(&self) -> usize;
14}
15
16pub struct SizedPool<T: SizedAllocatable + DynamicReset> {
17    /// each entry represents an allocation queue of 2**n bytes block
18    sub_pools: Vec<DynamicPool<T>>,
19}
20
21impl<T: SizedAllocatable + DynamicReset> SizedPool<T> {
22    /// cap is the capacity of each subpool, pool_size_power_of_two is the number of subpools,
23    /// init_fn takes the pool_size (the power of two) as input and outputs the reusable resource
24    pub fn new(cap: usize, pool_size_power_of_two: u32, max_pool_size: usize) -> Self {
25        let mut pools = Vec::new();
26        for pool_power in 0..pool_size_power_of_two {
27            let pool =
28                DynamicPool::new(cap, max_pool_size, move || T::new(2_usize.pow(pool_power)));
29            pools.push(pool);
30        }
31        Self { sub_pools: pools }
32    }
33
34    fn get_subpool_location(&self, size: usize) -> usize {
35        (size.next_power_of_two().trailing_zeros()) as usize
36    }
37
38    fn get_subpool(&self, size: usize) -> Result<&DynamicPool<T>, SizedPoolError> {
39        self.sub_pools
40            .get(self.get_subpool_location(size))
41            .ok_or(SizedPoolError::SizeExceedMaxSize)
42    }
43
44    pub fn try_pull(&self, size: usize) -> Result<DynamicPoolItem<T>, SizedPoolError> {
45        let pool = self.get_subpool(size)?;
46        match pool.try_take() {
47            None => {
48                tracing::debug!("not enough items in pool, allocating");
49                Ok(pool.take())
50            }
51            Some(x) => Ok(x),
52        }
53    }
54}
55
56#[cfg(test)]
57mod test {
58    use super::*;
59
60    #[derive(Debug)]
61    struct TestItem {
62        size: usize,
63    }
64
65    impl SizedAllocatable for TestItem {
66        fn new(size: usize) -> Self {
67            Self { size }
68        }
69
70        fn size(&self) -> usize {
71            self.size
72        }
73    }
74
75    impl DynamicReset for TestItem {
76        fn reset(&mut self) {}
77    }
78
79    #[test]
80    fn test_allocate() {
81        let pool: SizedPool<TestItem> = SizedPool::new(0, 40, 1024);
82        let mut items = Vec::new();
83        for _ in 0..2048 {
84            items.push(pool.try_pull(10).unwrap());
85        }
86    }
87}
sized_object_pool/lib.rs

sized_object_pool/
lib.rs
