compio_driver/buffer_pool/
iour.rs

1//! The io-uring buffer pool. It is backed by a [`Vec`] of [`Vec<u8>`].
2//! The kernel selects the buffer and returns `flags`. The crate
3//! [`io_uring_buf_ring`] handles the returning of buffer on drop.
4
5use std::{
6    borrow::{Borrow, BorrowMut},
7    fmt::{Debug, Formatter},
8    io,
9    ops::{Deref, DerefMut},
10    rc::Rc,
11};
12
13use io_uring_buf_ring::IoUringBufRing;
14#[cfg(not(fusion))]
15pub use {BufferPool as IoUringBufferPool, OwnedBuffer as IoUringOwnedBuffer};
16
17struct BufferPoolInner {
18    buf_ring: IoUringBufRing<Vec<u8>>,
19}
20
21impl BufferPoolInner {
22    fn reuse_buffer(&self, buffer_id: u16) {
23        // SAFETY: 0 is always valid length. We just want to get the buffer once and
24        // return it immediately.
25        unsafe { self.buf_ring.get_buf(buffer_id, 0) };
26    }
27}
28
29/// Buffer pool
30///
31/// A buffer pool to allow user no need to specify a specific buffer to do the
32/// IO operation
33#[derive(Clone)]
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(buf_ring: IoUringBufRing<Vec<u8>>) -> Self {
46        Self {
47            inner: Rc::new(BufferPoolInner { buf_ring }),
48        }
49    }
50
51    pub(crate) fn buffer_group(&self) -> u16 {
52        self.inner.buf_ring.buffer_group()
53    }
54
55    pub(crate) fn into_inner(self) -> Result<IoUringBufRing<Vec<u8>>, Self> {
56        Rc::try_unwrap(self.inner)
57            .map(|inner| inner.buf_ring)
58            .map_err(|inner| Self { inner })
59    }
60
61    #[doc(hidden)]
62    pub unsafe fn get_buffer(&self, buffer_id: u16, available_len: usize) -> OwnedBuffer {
63        OwnedBuffer {
64            pool: self.inner.clone(),
65            params: Some((buffer_id, available_len)),
66        }
67    }
68
69    /// ## Safety
70    /// * `len` should be the returned value from the op.
71    pub(crate) unsafe fn create_proxy(
72        &self,
73        slice: OwnedBuffer,
74        len: usize,
75    ) -> io::Result<BorrowedBuffer<'_>> {
76        let Some((buffer_id, available_len)) = slice.leak() else {
77            return Err(io::Error::other("no buffer selected"));
78        };
79        debug_assert_eq!(available_len, len);
80        unsafe { self.inner.buf_ring.get_buf(buffer_id, available_len) }
81            .map(BorrowedBuffer)
82            .ok_or_else(|| io::Error::other(format!("cannot find buffer {buffer_id}")))
83    }
84
85    pub(crate) fn reuse_buffer(&self, buffer_id: u16) {
86        self.inner.reuse_buffer(buffer_id);
87    }
88}
89
90#[doc(hidden)]
91pub struct OwnedBuffer {
92    pool: Rc<BufferPoolInner>,
93    params: Option<(u16, usize)>,
94}
95
96impl OwnedBuffer {
97    pub fn leak(mut self) -> Option<(u16, usize)> {
98        self.params.take()
99    }
100}
101
102impl Drop for OwnedBuffer {
103    fn drop(&mut self) {
104        if let Some((buffer_id, _)) = self.params {
105            self.pool.reuse_buffer(buffer_id);
106        }
107    }
108}
109
110/// Buffer borrowed from buffer pool
111///
112/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
113/// filled data
114pub struct BorrowedBuffer<'a>(io_uring_buf_ring::BorrowedBuffer<'a, Vec<u8>>);
115
116impl Debug for BorrowedBuffer<'_> {
117    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
119    }
120}
121
122impl Deref for BorrowedBuffer<'_> {
123    type Target = [u8];
124
125    fn deref(&self) -> &Self::Target {
126        self.0.deref()
127    }
128}
129
130impl DerefMut for BorrowedBuffer<'_> {
131    fn deref_mut(&mut self) -> &mut Self::Target {
132        self.0.deref_mut()
133    }
134}
135
136impl AsRef<[u8]> for BorrowedBuffer<'_> {
137    fn as_ref(&self) -> &[u8] {
138        self.deref()
139    }
140}
141
142impl AsMut<[u8]> for BorrowedBuffer<'_> {
143    fn as_mut(&mut self) -> &mut [u8] {
144        self.deref_mut()
145    }
146}
147
148impl Borrow<[u8]> for BorrowedBuffer<'_> {
149    fn borrow(&self) -> &[u8] {
150        self.deref()
151    }
152}
153
154impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
155    fn borrow_mut(&mut self) -> &mut [u8] {
156        self.deref_mut()
157    }
158}
compio_driver/buffer_pool/iour.rs

compio_driver/buffer_pool/
iour.rs
