rafx_base/
owned_pool.rs

1use crossbeam_channel::{bounded, Receiver, Sender};
2use std::ops::{Deref, DerefMut};
3
4pub struct OwnedPool<T> {
5    free: Vec<T>,
6    num_borrowed: usize,
7    receiver: Receiver<T>,
8    sender: Sender<T>,
9    reset_fn: fn(&mut T),
10    init_fn: fn() -> T,
11}
12
13impl<T> OwnedPool<T> {
14    pub fn with_capacity(
15        pool_size: usize,
16        init_fn: fn() -> T,
17        reset_fn: fn(&mut T),
18    ) -> Self {
19        let (sender, receiver) = bounded(pool_size);
20        Self {
21            free: Vec::with_capacity(pool_size),
22            num_borrowed: 0,
23            init_fn,
24            reset_fn,
25            sender,
26            receiver,
27        }
28    }
29
30    pub fn borrow(&mut self) -> Pooled<T> {
31        if let Some(mut pooled_val) = self.free.pop() {
32            self.num_borrowed += 1;
33            (self.reset_fn)(&mut pooled_val);
34            Pooled {
35                inner: Some(pooled_val),
36                sender: self.sender.clone(),
37            }
38        } else if self.num_borrowed < self.free.capacity() {
39            self.num_borrowed += 1;
40            Pooled {
41                inner: Some((self.init_fn)()),
42                sender: self.sender.clone(),
43            }
44        } else {
45            panic!("Cannot borrow more than `pool_size` entries.");
46        }
47    }
48
49    pub fn try_recv(&mut self) {
50        for pooled_val in self.receiver.try_iter() {
51            self.free.push(pooled_val);
52            self.num_borrowed -= 1;
53        }
54    }
55}
56
57pub struct Pooled<T> {
58    inner: Option<T>,
59    sender: Sender<T>,
60}
61
62impl<T> Pooled<T> {
63    pub fn as_ref(&self) -> &T {
64        self.inner.as_ref().unwrap()
65    }
66
67    pub fn as_mut(&mut self) -> &mut T {
68        self.inner.as_mut().unwrap()
69    }
70}
71
72impl<T> Deref for Pooled<T> {
73    type Target = T;
74
75    fn deref(&self) -> &Self::Target {
76        self.inner.as_ref().unwrap()
77    }
78}
79
80impl<T> DerefMut for Pooled<T> {
81    fn deref_mut(&mut self) -> &mut Self::Target {
82        self.inner.as_mut().unwrap()
83    }
84}
85
86impl<T> Drop for Pooled<T> {
87    fn drop(&mut self) {
88        let inner = std::mem::take(&mut self.inner);
89        let _ = self.sender.send(inner.unwrap());
90    }
91}