bufpool_fixed/
lib.rs

1pub mod buf;
2
3use buf::FixedBuf;
4use off64::usz;
5use std::alloc::alloc_zeroed;
6use std::alloc::Layout;
7use std::cmp::max;
8use std::collections::VecDeque;
9use std::mem::size_of;
10use std::sync::Arc;
11
12// 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.
13#[derive(Clone, Default)]
14struct BufPoolForSize(Arc<parking_lot::Mutex<VecDeque<usize>>>);
15
16struct Inner {
17  align: usize,
18  sizes: Vec<BufPoolForSize>,
19}
20
21/// Thread-safe pool of `FixedBuf` values, which are byte arrays with a fixed length.
22/// This can be cheaply cloned to share the same underlying pool around.
23/// The maximum length is 2^64, and the minimum alignment is 64. This allows storing the pointer and capacity in one `usize`, making it much faster to move the `FixedBuf` value around.
24#[derive(Clone)]
25pub struct FixedBufPool {
26  inner: Arc<Inner>,
27}
28
29impl FixedBufPool {
30  pub fn with_alignment(align: usize) -> Self {
31    assert!(align > 64);
32    assert!(align.is_power_of_two());
33    let mut sizes = Vec::new();
34    for _ in 0..64 {
35      sizes.push(Default::default());
36    }
37    Self {
38      inner: Arc::new(Inner { align, sizes }),
39    }
40  }
41
42  pub fn new() -> Self {
43    Self::with_alignment(max(64, size_of::<usize>()))
44  }
45
46  pub fn allocate_from_data(&self, data: impl AsRef<[u8]>) -> FixedBuf {
47    let mut buf = self.allocate_with_zeros(data.as_ref().len());
48    buf.copy_from_slice(data.as_ref());
49    buf
50  }
51
52  /// `cap` must be a power of two. It can safely be zero, but it will still cause an allocation of one byte due to rounding.
53  pub fn allocate_with_zeros(&self, cap: usize) -> FixedBuf {
54    // FixedBuf values do not have a length + capacity, so check that `cap` will be fully used.
55    assert!(cap.is_power_of_two());
56    // This will round `0` to `1`.
57    let cap = cap.next_power_of_two();
58    // Release lock ASAP.
59    let existing = self.inner.sizes[usz!(cap.ilog2())].0.lock().pop_front();
60    let ptr_and_cap = if let Some(ptr_and_cap) = existing {
61      ptr_and_cap
62    } else {
63      let ptr = unsafe { alloc_zeroed(Layout::from_size_align(cap, self.inner.align).unwrap()) };
64      // Failed allocations may return null.
65      assert!(!ptr.is_null());
66      let raw = ptr as usize;
67      assert_eq!(raw & (self.inner.align - 1), 0);
68      raw | usz!(cap.ilog2())
69    };
70    FixedBuf {
71      ptr_and_cap,
72      pool: self.clone(),
73    }
74  }
75}
bufpool_fixed/lib.rs

bufpool_fixed/
lib.rs
