compio_driver/buffer_pool/
fallback.rs

1//! The fallback buffer pool. It is backed by a [`VecDeque`] of [`Vec<u8>`].
2//! An [`OwnedBuffer`] is selected when the op is created. It keeps a strong
3//! reference to the buffer pool. The [`BorrowedBuffer`] is created after the op
4//! returns successfully.
5
6use std::{
7    borrow::{Borrow, BorrowMut},
8    cell::RefCell,
9    collections::VecDeque,
10    fmt::{Debug, Formatter},
11    io,
12    mem::ManuallyDrop,
13    ops::{Deref, DerefMut},
14    rc::Rc,
15};
16
17use compio_buf::{IntoInner, IoBuf, IoBufMut, SetBufInit, Slice};
18
19struct BufferPoolInner {
20    buffers: RefCell<VecDeque<Vec<u8>>>,
21}
22
23impl BufferPoolInner {
24    pub(crate) fn add_buffer(&self, mut buffer: Vec<u8>) {
25        buffer.clear();
26        self.buffers.borrow_mut().push_back(buffer)
27    }
28}
29
30/// Buffer pool
31///
32/// A buffer pool to allow user no need to specify a specific buffer to do the
33/// IO operation
34pub struct BufferPool {
35    inner: Rc<BufferPoolInner>,
36}
37
38impl Debug for BufferPool {
39    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("BufferPool").finish_non_exhaustive()
41    }
42}
43
44impl BufferPool {
45    pub(crate) fn new(buffer_len: u16, buffer_size: usize) -> Self {
46        // To match the behavior of io-uring, extend the number of buffers.
47        let buffers = (0..buffer_len.next_power_of_two())
48            .map(|_| Vec::with_capacity(buffer_size))
49            .collect();
50
51        Self {
52            inner: Rc::new(BufferPoolInner {
53                buffers: RefCell::new(buffers),
54            }),
55        }
56    }
57
58    /// Select an [`OwnedBuffer`] when the op creates.
59    #[doc(hidden)]
60    pub fn get_buffer(&self, len: usize) -> io::Result<OwnedBuffer> {
61        let buffer = self
62            .inner
63            .buffers
64            .borrow_mut()
65            .pop_front()
66            .ok_or_else(|| io::Error::other("buffer ring has no available buffer"))?;
67        let len = if len == 0 {
68            buffer.capacity()
69        } else {
70            buffer.capacity().min(len)
71        };
72        Ok(OwnedBuffer::new(buffer.slice(..len), self.inner.clone()))
73    }
74
75    /// Return the buffer to the pool.
76    pub(crate) fn add_buffer(&self, buffer: Vec<u8>) {
77        self.inner.add_buffer(buffer);
78    }
79
80    /// ## Safety
81    /// * `len` should be valid.
82    #[doc(hidden)]
83    pub unsafe fn create_proxy(&self, mut slice: OwnedBuffer, len: usize) -> BorrowedBuffer<'_> {
84        unsafe {
85            slice.set_buf_init(len);
86        }
87        BorrowedBuffer::new(slice.into_inner(), self)
88    }
89}
90
91#[doc(hidden)]
92pub struct OwnedBuffer {
93    buffer: ManuallyDrop<Slice<Vec<u8>>>,
94    pool: ManuallyDrop<Rc<BufferPoolInner>>,
95}
96
97impl OwnedBuffer {
98    fn new(buffer: Slice<Vec<u8>>, pool: Rc<BufferPoolInner>) -> Self {
99        Self {
100            buffer: ManuallyDrop::new(buffer),
101            pool: ManuallyDrop::new(pool),
102        }
103    }
104}
105
106unsafe impl IoBuf for OwnedBuffer {
107    fn as_buf_ptr(&self) -> *const u8 {
108        self.buffer.as_buf_ptr()
109    }
110
111    fn buf_len(&self) -> usize {
112        self.buffer.buf_len()
113    }
114
115    fn buf_capacity(&self) -> usize {
116        self.buffer.buf_capacity()
117    }
118}
119
120unsafe impl IoBufMut for OwnedBuffer {
121    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
122        self.buffer.as_buf_mut_ptr()
123    }
124}
125
126impl SetBufInit for OwnedBuffer {
127    unsafe fn set_buf_init(&mut self, len: usize) {
128        self.buffer.set_buf_init(len);
129    }
130}
131
132impl Drop for OwnedBuffer {
133    fn drop(&mut self) {
134        // Safety: `take` is called only once here.
135        self.pool
136            .add_buffer(unsafe { ManuallyDrop::take(&mut self.buffer) }.into_inner());
137        // Safety: `drop` is called only once here.
138        unsafe { ManuallyDrop::drop(&mut self.pool) };
139    }
140}
141
142impl IntoInner for OwnedBuffer {
143    type Inner = Slice<Vec<u8>>;
144
145    fn into_inner(mut self) -> Self::Inner {
146        // Safety: `self` is forgotten in this method.
147        let buffer = unsafe { ManuallyDrop::take(&mut self.buffer) };
148        // The buffer is taken, we only need to drop the Rc.
149        // Safety: `self` is forgotten in this method.
150        unsafe { ManuallyDrop::drop(&mut self.pool) };
151        std::mem::forget(self);
152        buffer
153    }
154}
155
156/// Buffer borrowed from buffer pool
157///
158/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
159/// filled data
160pub struct BorrowedBuffer<'a> {
161    buffer: ManuallyDrop<Slice<Vec<u8>>>,
162    pool: &'a BufferPool,
163}
164
165impl<'a> BorrowedBuffer<'a> {
166    pub(crate) fn new(buffer: Slice<Vec<u8>>, pool: &'a BufferPool) -> Self {
167        Self {
168            buffer: ManuallyDrop::new(buffer),
169            pool,
170        }
171    }
172}
173
174impl Debug for BorrowedBuffer<'_> {
175    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
176        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
177    }
178}
179
180impl Drop for BorrowedBuffer<'_> {
181    fn drop(&mut self) {
182        // Safety: `take` is called only once here.
183        let buffer = unsafe { ManuallyDrop::take(&mut self.buffer) };
184        self.pool.add_buffer(buffer.into_inner());
185    }
186}
187
188impl Deref for BorrowedBuffer<'_> {
189    type Target = [u8];
190
191    fn deref(&self) -> &Self::Target {
192        self.buffer.deref()
193    }
194}
195
196impl DerefMut for BorrowedBuffer<'_> {
197    fn deref_mut(&mut self) -> &mut Self::Target {
198        self.buffer.deref_mut()
199    }
200}
201
202impl AsRef<[u8]> for BorrowedBuffer<'_> {
203    fn as_ref(&self) -> &[u8] {
204        self.deref()
205    }
206}
207
208impl AsMut<[u8]> for BorrowedBuffer<'_> {
209    fn as_mut(&mut self) -> &mut [u8] {
210        self.deref_mut()
211    }
212}
213
214impl Borrow<[u8]> for BorrowedBuffer<'_> {
215    fn borrow(&self) -> &[u8] {
216        self.deref()
217    }
218}
219
220impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
221    fn borrow_mut(&mut self) -> &mut [u8] {
222        self.deref_mut()
223    }
224}