simoom_rt/
pool.rs

1//! Per-loop buffer pool in size classes.
2//!
3//! Buffers are plain `Vec<u8>`s whose capacity is exactly one of the class sizes. [`BufPool::get`]
4//! hands out an empty buffer of the smallest class that fits; [`BufPool::put`] takes it back for
5//! reuse, up to a cap per class (beyond that, or for a capacity that is not a class size, the
6//! buffer is freed). A pool belongs to one loop thread, so there is no locking.
7
8/// Default class sizes: 2 KiB, 16 KiB, 64 KiB.
9pub const DEFAULT_CLASSES: [usize; 3] = [2048, 16 * 1024, 64 * 1024];
10
11/// Counters of a [`BufPool`].
12#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
13pub struct PoolStats {
14    /// Buffers handed out that were newly allocated.
15    pub allocated: u64,
16    /// Buffers handed out that were reused from the pool.
17    pub reused: u64,
18    /// Buffers taken back and kept.
19    pub returned: u64,
20    /// Buffers taken back and freed (pool full or odd size).
21    pub freed: u64,
22    /// Buffers kept in the pool now.
23    pub cached: usize,
24    /// Bytes kept in the pool now.
25    pub cached_bytes: usize,
26}
27
28/// A pool of byte buffers in size classes. See the [module docs](self).
29#[derive(Debug)]
30pub struct BufPool {
31    classes: [usize; 3],
32    free: [Vec<Vec<u8>>; 3],
33    max_cached: [usize; 3],
34    stats: PoolStats,
35}
36
37impl BufPool {
38    /// A pool with the given class sizes (ascending) and at most `max_cached` buffers kept per
39    /// class (the larger classes keep a quarter and a sixteenth of that).
40    pub fn new(classes: [usize; 3], max_cached: usize) -> BufPool {
41        assert!(classes[0] > 0 && classes[0] < classes[1] && classes[1] < classes[2], "class sizes must ascend");
42        BufPool {
43            classes,
44            free: [Vec::new(), Vec::new(), Vec::new()],
45            max_cached: [max_cached, (max_cached / 4).max(1), (max_cached / 16).max(1)],
46            stats: PoolStats::default(),
47        }
48    }
49
50    /// The class sizes.
51    pub fn classes(&self) -> [usize; 3] {
52        self.classes
53    }
54
55    /// An empty buffer with capacity of at least `min` bytes: the smallest class that fits, or
56    /// exactly `min` when it is larger than every class.
57    #[inline]
58    pub fn get(&mut self, min: usize) -> Vec<u8> {
59        for (i, &size) in self.classes.iter().enumerate() {
60            if min <= size {
61                if let Some(b) = self.free[i].pop() {
62                    self.stats.reused += 1;
63                    self.stats.cached -= 1;
64                    self.stats.cached_bytes -= size;
65                    return b;
66                }
67                self.stats.allocated += 1;
68                return Vec::with_capacity(size);
69            }
70        }
71        self.stats.allocated += 1;
72        Vec::with_capacity(min)
73    }
74
75    /// Takes a buffer back. It is cleared and kept if its capacity is a class size and that
76    /// class is not full; otherwise it is freed. Out of line: it is called from many places in
77    /// the loop, and a call costs nothing next to the free or push it ends in.
78    #[inline(never)]
79    pub fn put(&mut self, mut buf: Vec<u8>) {
80        let cap = buf.capacity();
81        for (i, &size) in self.classes.iter().enumerate() {
82            if cap == size {
83                if self.free[i].len() < self.max_cached[i] {
84                    buf.clear();
85                    self.free[i].push(buf);
86                    self.stats.returned += 1;
87                    self.stats.cached += 1;
88                    self.stats.cached_bytes += size;
89                } else {
90                    self.stats.freed += 1;
91                }
92                return;
93            }
94        }
95        self.stats.freed += 1;
96    }
97
98    /// Frees cached buffers until at most `keep` remain in each class (for trimming after a
99    /// burst).
100    pub fn trim(&mut self, keep: usize) {
101        for i in 0..3 {
102            while self.free[i].len() > keep {
103                self.free[i].pop();
104                self.stats.cached -= 1;
105                self.stats.cached_bytes -= self.classes[i];
106                self.stats.freed += 1;
107            }
108        }
109    }
110
111    /// Counters.
112    pub fn stats(&self) -> PoolStats {
113        self.stats
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn reuses_buffers_by_class() {
123        let mut p = BufPool::new(DEFAULT_CLASSES, 4);
124        let a = p.get(100);
125        assert_eq!(a.capacity(), 2048);
126        let a_ptr = a.as_ptr();
127        p.put(a);
128        let mut b = p.get(2048);
129        assert_eq!(b.as_ptr(), a_ptr, "the same allocation comes back");
130        b.extend_from_slice(b"dirty");
131        p.put(b);
132        let c = p.get(1);
133        assert!(c.is_empty(), "returned buffers are cleared");
134        assert_eq!(p.get(3000).capacity(), 16 * 1024);
135        assert_eq!(p.get(70_000).capacity(), 70_000);
136        let s = p.stats();
137        assert_eq!((s.allocated, s.reused), (3, 2));
138    }
139
140    #[test]
141    fn caps_the_cache_and_frees_odd_sizes() {
142        let mut p = BufPool::new(DEFAULT_CLASSES, 2);
143        let bufs: Vec<Vec<u8>> = (0..5).map(|_| p.get(10)).collect();
144        for b in bufs {
145            p.put(b);
146        }
147        assert_eq!(p.stats().cached, 2);
148        assert_eq!(p.stats().freed, 3);
149        p.put(Vec::with_capacity(100));
150        assert_eq!(p.stats().freed, 4);
151        assert_eq!(p.stats().cached_bytes, 4096);
152        p.trim(0);
153        assert_eq!(p.stats().cached, 0);
154    }
155}
simoom_rt/pool.rs

simoom_rt/
pool.rs
