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        self.buf_ring
57            .get_buf(buffer_id, available_len)
58            .map(BorrowedBuffer)
59            .ok_or_else(|| io::Error::other(format!("cannot find buffer {buffer_id}")))
60    }
61
62    pub(crate) fn reuse_buffer(&self, flags: u32) {
63        // It ignores invalid flags.
64        if let Some(buffer_id) = buffer_select(flags) {
65            // Safety: 0 is always valid length. We just want to get the buffer once and
66            // return it immediately.
67            unsafe { self.buf_ring.get_buf(buffer_id, 0) };
68        }
69    }
70}
71
72/// Buffer borrowed from buffer pool
73///
74/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
75/// filled data
76pub struct BorrowedBuffer<'a>(io_uring_buf_ring::BorrowedBuffer<'a, Vec<u8>>);
77
78impl Debug for BorrowedBuffer<'_> {
79    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
80        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
81    }
82}
83
84impl Deref for BorrowedBuffer<'_> {
85    type Target = [u8];
86
87    fn deref(&self) -> &Self::Target {
88        self.0.deref()
89    }
90}
91
92impl DerefMut for BorrowedBuffer<'_> {
93    fn deref_mut(&mut self) -> &mut Self::Target {
94        self.0.deref_mut()
95    }
96}
97
98impl AsRef<[u8]> for BorrowedBuffer<'_> {
99    fn as_ref(&self) -> &[u8] {
100        self.deref()
101    }
102}
103
104impl AsMut<[u8]> for BorrowedBuffer<'_> {
105    fn as_mut(&mut self) -> &mut [u8] {
106        self.deref_mut()
107    }
108}
109
110impl Borrow<[u8]> for BorrowedBuffer<'_> {
111    fn borrow(&self) -> &[u8] {
112        self.deref()
113    }
114}
115
116impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
117    fn borrow_mut(&mut self) -> &mut [u8] {
118        self.deref_mut()
119    }
120}