ts_mem_pool/
memory_pool.rs

1use std::sync::mpsc;
2use std::sync::atomic::{
3    AtomicUsize,
4    Ordering,
5};
6use arc_recycled::{
7    ArcRecycled,
8    Recycle,
9};
10
11/// Boxed closure that creates objects of type T
12pub type CreateFn<T> = Box<Fn() -> T>;
13
14/// Memory pool structure
15#[allow(missing_debug_implementations)]
16pub struct MemoryPool<T> {
17    size: AtomicUsize,
18    max: usize,
19    receiver: mpsc::Receiver<Option<T>>,
20    sender: mpsc::Sender<Option<T>>,
21    creator: CreateFn<T>,
22}
23
24impl<T: Recycle> MemoryPool<T> {
25    /// Constructor, must take intial size and maximum size.
26    /// The creator closure is used to initialize the mem slots
27    ///
28    /// # Panics
29    ///
30    /// This function will panic if size > max
31    pub fn create_with(size: usize, max: usize, creator: CreateFn<T>) -> MemoryPool<T> {
32        assert!(size <= max);
33        let (tx, rx) = mpsc::channel();
34        for _ in 0..size {
35            tx.send(Some(creator())).unwrap()
36        }
37
38        MemoryPool {
39            size: AtomicUsize::new(size),
40            max,
41            receiver: rx,
42            sender: tx,
43            creator,
44        }
45    }
46
47    /// This function returns a memory slot from the memory pool
48    ///
49    /// # Panics
50    ///
51    /// This function will panic if it needs to allocate more than max
52    pub fn get(&self) -> ArcRecycled<T> {
53        loop {
54            /// Try to get a mem_slot
55            match self.receiver.try_recv() {
56                /// If got one wrap and return it
57                Ok(Some(mem_slot)) => {
58                    return ArcRecycled::new(mem_slot, self.sender.clone());
59                }
60
61                /// If got None, keep trying
62                Ok(None) => {
63                    self.size.fetch_sub(1, Ordering::Relaxed);
64                }
65
66                /// If channel is empty try to create a new memory slot
67                /// If we have place this works, if not, it panics!
68                Err(mpsc::TryRecvError::Empty) => {
69                    if self.size.fetch_add(1, Ordering::Relaxed) < self.max {
70                        return ArcRecycled::new((self.creator)(), self.sender.clone());
71                    }
72                    else {
73                        panic!("Exceeded memory pool limit");
74                    }
75                }
76
77                /// Unreachable
78                Err(_) => {
79                    unreachable!("If the memory pool is alive, the channel cannot be disconnected")
80                }
81            }
82        }
83    }
84
85    /// This function returns a memory slot from the memory pool
86    /// if size does not exceed max. returns None otherwise
87    pub fn try_get(&self) -> Option<ArcRecycled<T>> {
88        loop {
89            /// Try to get a mem_slot
90            match self.receiver.try_recv() {
91                /// If got one wrap and return it
92                Ok(Some(mem_slot)) => {
93                    return Some(ArcRecycled::new(mem_slot, self.sender.clone()));
94                }
95
96                /// If got None, keep trying
97                Ok(None) => {
98                    self.size.fetch_sub(1, Ordering::Relaxed);
99                }
100
101                /// If channel is empty try to create a new memory slot
102                /// If we have place this works, if not, it panics!
103                Err(mpsc::TryRecvError::Empty) => {
104                    if self.size.fetch_add(1, Ordering::Relaxed) < self.max {
105                        return Some(ArcRecycled::new((self.creator)(), self.sender.clone()));
106                    }
107                    else {
108                        return None;
109                    }
110                }
111
112                /// Unreachable
113                Err(_) => {
114                    unreachable!("If the memory pool is alive, the channel cannot be disconnected")
115                }
116            }
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn creation_test() {
127        let mem = MemoryPool::create_with(5, 10, Box::new(|| { Vec::<f64>::with_capacity(20) }));
128        let _v1 = mem.get();
129        let _v2 = mem.try_get().unwrap();
130    }
131
132    #[test]
133    fn extra_elements_test() {
134        let mem = MemoryPool::create_with(5, 10, Box::new(|| { Vec::<f64>::with_capacity(20) }));
135        let mut vecs = vec![];
136        for _ in 0..10 {
137            vecs.push(mem.get());
138        }
139    }
140
141    #[test]
142    fn recycling_test() {
143        let mem = MemoryPool::create_with(5, 10, Box::new(|| { Vec::<f64>::with_capacity(20) }));
144        /// First use of all 10 elements
145        {
146            let mut vecs = vec![];
147            for _ in 0..10 {
148                vecs.push(mem.get());
149            }
150        }
151
152        /// Second use of all 10 elements
153        {
154            let mut vecs = vec![];
155            for _ in 0..10 {
156                vecs.push(mem.get());
157            }
158        }
159    }
160
161    #[test]
162    #[should_panic]
163    fn too_many_elements_test() {
164        let mem = MemoryPool::create_with(5, 10, Box::new(|| { Vec::<f64>::with_capacity(20) }));
165        let mut vecs = vec![];
166        for _ in 0..11 {
167            vecs.push(mem.get());
168        }
169    }
170
171    #[test]
172    fn too_many_elements_try_test() {
173        let mem = MemoryPool::create_with(5, 10, Box::new(|| { Vec::<f64>::with_capacity(20) }));
174        let mut vecs = vec![];
175        for _ in 0..10 {
176            vecs.push(mem.get());
177        }
178        assert!(mem.try_get().is_none());
179    }
180}
ts_mem_pool/memory_pool.rs

ts_mem_pool/
memory_pool.rs
