bufpool/
lib.rs

1pub mod buf;
2
3use buf::Buf;
4use once_cell::sync::Lazy;
5use std::alloc::alloc;
6use std::alloc::Layout;
7use std::collections::VecDeque;
8use std::mem::size_of;
9use std::panic::RefUnwindSafe;
10use std::panic::UnwindSafe;
11use std::sync::Arc;
12
13// TODO Benchmark parking_lot::Mutex<VecDeque<>> against crossbeam_channel and flume. Also consider one allocator per thread, which could waste a lot of memory but also be very quick.
14#[derive(Clone, Default)]
15struct BufPoolForSize(Arc<parking_lot::Mutex<VecDeque<*mut u8>>>);
16
17unsafe impl Send for BufPoolForSize {}
18unsafe impl Sync for BufPoolForSize {}
19impl UnwindSafe for BufPoolForSize {}
20impl RefUnwindSafe for BufPoolForSize {}
21
22struct BufPoolInner {
23  align: usize,
24  #[cfg(not(feature = "no-pool"))]
25  sizes: Vec<BufPoolForSize>,
26}
27
28#[derive(Clone)]
29pub struct BufPool {
30  inner: Arc<BufPoolInner>,
31}
32
33impl BufPool {
34  pub fn with_alignment(align: usize) -> Self {
35    assert!(align > 0);
36    assert!(align.is_power_of_two());
37    Self {
38      inner: Arc::new(BufPoolInner {
39        align,
40        #[cfg(not(feature = "no-pool"))]
41        sizes: (0..(size_of::<usize>() * 8))
42          .map(|_| Default::default())
43          .collect(),
44      }),
45    }
46  }
47
48  pub fn new() -> Self {
49    Self::with_alignment(size_of::<usize>())
50  }
51
52  fn system_allocate_raw(&self, cap: usize) -> *mut u8 {
53    unsafe { alloc(Layout::from_size_align(cap, self.inner.align).unwrap()) }
54  }
55
56  /// NOTE: This provides a Buf that can grow to `cap`, but has an initial length of zero. Use `allocate_with_zeros` to return something equivalent to `vec![0u8; cap]`.
57  /// `cap` can safely be zero, but it will still cause an allocation of one byte due to rounding.
58  pub fn allocate(&self, cap: usize) -> Buf {
59    // This will round `0` to `1`.
60    let cap = cap.next_power_of_two();
61
62    #[cfg(not(feature = "no-pool"))]
63    let data = if let Some(data) = self.inner.sizes[cap.ilog2() as usize].0.lock().pop_front() {
64      data
65    } else {
66      self.system_allocate_raw(cap)
67    };
68    #[cfg(feature = "no-pool")]
69    let data = self.system_allocate_raw(cap);
70
71    // Failed allocations may return null.
72    assert!(!data.is_null());
73
74    Buf {
75      data,
76      len: 0,
77      cap,
78      pool: self.clone(),
79    }
80  }
81
82  pub fn allocate_from_data(&self, data: impl AsRef<[u8]>) -> Buf {
83    let mut buf = self.allocate(data.as_ref().len());
84    buf.extend_from_slice(data.as_ref());
85    buf
86  }
87
88  pub fn allocate_from_iter(&self, data: impl IntoIterator<Item = u8>, len: usize) -> Buf {
89    let mut buf = self.allocate(len);
90    buf.extend(data);
91    buf
92  }
93
94  /// The returned Buf will have a length equal to the capacity, filled with uninitialised bytes.
95  pub fn allocate_uninitialised(&self, len: usize) -> Buf {
96    let mut buf = self.allocate(len);
97    unsafe { buf.set_len(len) };
98    buf
99  }
100
101  pub fn allocate_with_fill(&self, val: u8, len: usize) -> Buf {
102    let mut buf = self.allocate_uninitialised(len);
103    buf.fill(val);
104    buf
105  }
106
107  pub fn allocate_with_zeros(&self, len: usize) -> Buf {
108    self.allocate_with_fill(0, len)
109  }
110}
111
112pub static BUFPOOL: Lazy<BufPool> = Lazy::new(|| BufPool::new());
bufpool/lib.rs

bufpool/
lib.rs
