1use crossbeam::queue::SegQueue;
 2use std::ops::{Deref, DerefMut};
 3use std::sync::Arc;
 4
 5/// A thread-safe, lock-free pool of reusable byte buffers.
 6#[derive(Clone)]
 7pub struct BufPool {
 8    queue: Arc<SegQueue<Vec<u8>>>,
 9}
10
11impl BufPool {
12    /// Creates a pool pre-populated with `capacity` empty buffers.
13    pub fn with_capacity(capacity: usize) -> Self {
14        let queue = Arc::new(SegQueue::new());
15        for _ in 0..capacity {
16            queue.push(Vec::new());
17        }
18        Self { queue }
19    }
20
21    /// Checks out a buffer from the pool. If the pool is empty, allocates a new buffer.
22    pub fn checkout(&self) -> BufGuard {
23        let buf = self.queue.pop().unwrap_or_else(|| Vec::new());
24        BufGuard {
25            buf,
26            pool: self.queue.clone(),
27        }
28    }
29}
30
31/// RAII guard that automatically returns the buffer to the pool when dropped.
32pub struct BufGuard {
33    buf: Vec<u8>,
34    pool: Arc<SegQueue<Vec<u8>>>,
35}
36
37impl BufGuard {
38    /// Extracts the buffer, preventing automatic return to the pool.
39    /// Useful if you need to move the buffer elsewhere.
40    pub fn into_inner(mut self) -> Vec<u8> {
41        let buf = std::mem::take(&mut self.buf);
42        // SAFETY:
43        // 1) We forget `self`, preventing it from
44        // double-free'ing `pool`
45        // 2) The 'pointer' is really a &mut, so
46        // it is aligned, type-valid, and valid for reads.
47        let pool = unsafe { std::ptr::read(&self.pool) };
48        std::mem::forget(self);
49        drop(pool);
50        buf
51    }
52}
53
54impl Drop for BufGuard {
55    fn drop(&mut self) {
56        let mut buf = std::mem::take(&mut self.buf);
57        buf.clear();
58        self.pool.push(buf);
59    }
60}
61
62impl Deref for BufGuard {
63    type Target = Vec<u8>;
64
65    fn deref(&self) -> &Self::Target {
66        &self.buf
67    }
68}
69
70impl DerefMut for BufGuard {
71    fn deref_mut(&mut self) -> &mut Self::Target {
72        &mut self.buf
73    }
74}