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};
11
12use io_uring::cqueue::buffer_select;
13use io_uring_buf_ring::IoUringBufRing;
14
15/// Buffer pool
16///
17/// A buffer pool to allow user no need to specify a specific buffer to do the
18/// IO operation
19pub struct BufferPool {
20    buf_ring: IoUringBufRing<Vec<u8>>,
21}
22
23impl Debug for BufferPool {
24    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("BufferPool").finish_non_exhaustive()
26    }
27}
28
29impl BufferPool {
30    pub(crate) fn new(buf_ring: IoUringBufRing<Vec<u8>>) -> Self {
31        Self { buf_ring }
32    }
33
34    pub(crate) fn buffer_group(&self) -> u16 {
35        self.buf_ring.buffer_group()
36    }
37
38    pub(crate) fn into_inner(self) -> IoUringBufRing<Vec<u8>> {
39        self.buf_ring
40    }
41
42    /// ## Safety
43    /// * `available_len` should be the returned value from the op.
44    pub(crate) unsafe fn get_buffer(
45        &self,
46        flags: u32,
47        available_len: usize,
48    ) -> io::Result<BorrowedBuffer<'_>> {
49        let buffer_id = buffer_select(flags).ok_or_else(|| {
50            io::Error::new(
51                io::ErrorKind::InvalidInput,
52                format!("flags {flags} is invalid"),
53            )
54        })?;
55
56        unsafe { self.buf_ring.get_buf(buffer_id, available_len) }
57            .map(BorrowedBuffer)
58            .ok_or_else(|| io::Error::other(format!("cannot find buffer {buffer_id}")))
59    }
60
61    pub(crate) fn reuse_buffer(&self, flags: u32) {
62        // It ignores invalid flags.
63        if let Some(buffer_id) = buffer_select(flags) {
64            // SAFETY: 0 is always valid length. We just want to get the buffer once and
65            // return it immediately.
66            unsafe { self.buf_ring.get_buf(buffer_id, 0) };
67        }
68    }
69}
70
71/// Buffer borrowed from buffer pool
72///
73/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
74/// filled data
75pub struct BorrowedBuffer<'a>(io_uring_buf_ring::BorrowedBuffer<'a, Vec<u8>>);
76
77impl Debug for BorrowedBuffer<'_> {
78    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
79        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
80    }
81}
82
83impl Deref for BorrowedBuffer<'_> {
84    type Target = [u8];
85
86    fn deref(&self) -> &Self::Target {
87        self.0.deref()
88    }
89}
90
91impl DerefMut for BorrowedBuffer<'_> {
92    fn deref_mut(&mut self) -> &mut Self::Target {
93        self.0.deref_mut()
94    }
95}
96
97impl AsRef<[u8]> for BorrowedBuffer<'_> {
98    fn as_ref(&self) -> &[u8] {
99        self.deref()
100    }
101}
102
103impl AsMut<[u8]> for BorrowedBuffer<'_> {
104    fn as_mut(&mut self) -> &mut [u8] {
105        self.deref_mut()
106    }
107}
108
109impl Borrow<[u8]> for BorrowedBuffer<'_> {
110    fn borrow(&self) -> &[u8] {
111        self.deref()
112    }
113}
114
115impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
116    fn borrow_mut(&mut self) -> &mut [u8] {
117        self.deref_mut()
118    }
119}