graven/
pool.rs

1//! Buffer pools.
2
3use bytes::BytesMut;
4
5/// Source of reusable buffers. `take` returns an empty buffer; `give` returns it.
6pub trait BufPool {
7    fn take(&self) -> BytesMut;
8    fn give(&self, b: BytesMut);
9    /// Preferred buffer size; `Builder::flush` reserves this much when a
10    /// buffer runs low.
11    fn chunk_size(&self) -> usize {
12        4096
13    }
14}
15
16/// Mutex-protected free list of chunk buffers.
17///
18/// Finished messages are `split()` off the front of a chunk, so consecutive
19/// messages pack into the same allocation (bump-arena style). `take` only
20/// asks for a fresh chunk when less than a quarter of one is left; that
21/// `reserve` reuses the allocation for free once every message cut from it
22/// has been dropped, so steady state is zero allocations, and with messages
23/// in flight it costs about one allocation per chunk instead of per message.
24/// Memory pinned by in-flight messages is roughly their own size plus one chunk.
25#[cfg(feature = "std")]
26pub struct SimplePool {
27    free: std::sync::Mutex<alloc::vec::Vec<BytesMut>>,
28    capacity: usize,
29}
30
31#[cfg(feature = "std")]
32impl SimplePool {
33    pub fn new(initial_capacity: usize) -> Self {
34        SimplePool { free: std::sync::Mutex::new(alloc::vec::Vec::with_capacity(16)), capacity: initial_capacity }
35    }
36}
37
38#[cfg(feature = "std")]
39impl BufPool for SimplePool {
40    fn take(&self) -> BytesMut {
41        let popped = self.free.lock().map(|mut f| f.pop()).unwrap_or(None);
42        match popped {
43            Some(mut b) => {
44                b.clear();
45                if b.capacity() < self.capacity / 4 {
46                    b.reserve(self.capacity);
47                }
48                b
49            }
50            None => BytesMut::with_capacity(self.capacity),
51        }
52    }
53    fn chunk_size(&self) -> usize {
54        self.capacity
55    }
56    fn give(&self, b: BytesMut) {
57        if let Ok(mut f) = self.free.lock()
58            && f.len() < f.capacity() {
59                f.push(b);
60            }
61    }
62}