Skip to main content

dope_uring/driver/
provided.rs

1use std::{
2    alloc::{self, Layout},
3    io,
4    sync::atomic::{AtomicU16, Ordering},
5};
6
7use dope_core::driver::ProvidedRingOps;
8use io_uring::types::BufRingEntry;
9
10use crate::driver::{ProvidedBuf, Ring};
11
12pub struct ProvidedRing {
13    tail_pos: u16,
14    last_published: u16,
15    mask: u16,
16    buf_len: u32,
17    buf_base: *mut u8,
18    entries: *mut BufRingEntry,
19    tail: *mut AtomicU16,
20    _buf_mem: Box<[u8]>,
21    entries_layout: Layout,
22}
23
24impl ProvidedRing {
25    const BGID: u16 = 1;
26    const RING_ALIGN: usize = 16384;
27
28    pub(crate) fn new(ring: &Ring, entries: u16, buf_len: usize) -> io::Result<Self> {
29        if !entries.is_power_of_two() {
30            return Err(io::Error::new(
31                io::ErrorKind::InvalidInput,
32                "pbuf_ring entries must be power-of-two",
33            ));
34        }
35        let entry_size = std::mem::size_of::<BufRingEntry>();
36        let alloc_len = entry_size.checked_mul(entries as usize).ok_or_else(|| {
37            io::Error::new(io::ErrorKind::InvalidInput, "pbuf_ring size overflow")
38        })?;
39        let entries_layout = Layout::from_size_align(alloc_len, Self::RING_ALIGN)
40            .map_err(|_| io::Error::other("pbuf_ring layout"))?;
41        let entries_ptr = unsafe { alloc::alloc_zeroed(entries_layout) } as *mut BufRingEntry;
42        if entries_ptr.is_null() {
43            return Err(io::Error::other("pbuf_ring alloc"));
44        }
45
46        let total_buf_len = (entries as usize).checked_mul(buf_len).ok_or_else(|| {
47            io::Error::new(io::ErrorKind::InvalidInput, "pbuf backing size overflow")
48        })?;
49        let mut buf_mem = vec![0u8; total_buf_len].into_boxed_slice();
50        let buf_base = buf_mem.as_mut_ptr();
51
52        for bid in 0..entries {
53            let e = unsafe { &mut *entries_ptr.add(bid as usize) };
54            let addr = unsafe { buf_base.add(bid as usize * buf_len) };
55            e.set_addr(addr as u64);
56            e.set_len(buf_len as u32);
57            e.set_bid(bid);
58        }
59
60        let tail =
61            unsafe { BufRingEntry::tail(entries_ptr as *const _) as *mut u16 } as *mut AtomicU16;
62        unsafe { (*tail).store(entries, Ordering::Release) };
63
64        unsafe {
65            ring.submitter().register_buf_ring_with_flags(
66                entries_ptr as u64,
67                entries,
68                Self::BGID,
69                0,
70            )?;
71        }
72
73        Ok(Self {
74            tail_pos: entries,
75            last_published: entries,
76            mask: entries.wrapping_sub(1),
77            buf_len: buf_len as u32,
78            buf_base,
79            entries: entries_ptr,
80            tail,
81            _buf_mem: buf_mem,
82            entries_layout,
83        })
84    }
85
86    pub(crate) fn group(&self) -> (u16, u32) {
87        (Self::BGID, self.buf_len)
88    }
89
90    pub(crate) fn make_buf(&mut self, bid: u16, len: usize) -> ProvidedBuf {
91        let (ptr, cap) = self.ptr_len(bid);
92        ProvidedBuf::new(bid, ptr, cap, len)
93    }
94
95    pub(crate) fn release(&mut self, bid: u16) {
96        self.defer(bid);
97        unsafe { (*self.tail).store(self.tail_pos, Ordering::Release) };
98        self.last_published = self.tail_pos;
99    }
100
101    pub(crate) fn flush(&mut self) -> u64 {
102        let n = self.tail_pos.wrapping_sub(self.last_published);
103        if n == 0 {
104            return 0;
105        }
106        unsafe { (*self.tail).store(self.tail_pos, Ordering::Release) };
107        self.last_published = self.tail_pos;
108        n as u64
109    }
110}
111
112impl ProvidedRingOps for ProvidedRing {
113    #[inline(always)]
114    fn ptr_len(&self, bid: u16) -> (*mut u8, usize) {
115        let off = bid as usize * self.buf_len as usize;
116        let ptr = unsafe { self.buf_base.add(off) };
117        (ptr, self.buf_len as usize)
118    }
119
120    #[inline(always)]
121    fn defer(&mut self, bid: u16) {
122        let slot = (self.tail_pos & self.mask) as usize;
123        let e = unsafe { &mut *self.entries.add(slot) };
124        let addr = unsafe { self.buf_base.add(bid as usize * self.buf_len as usize) };
125        e.set_addr(addr as u64);
126        e.set_len(self.buf_len);
127        e.set_bid(bid);
128        self.tail_pos = self.tail_pos.wrapping_add(1);
129    }
130}
131
132impl Drop for ProvidedRing {
133    fn drop(&mut self) {
134        unsafe { alloc::dealloc(self.entries as *mut u8, self.entries_layout) };
135    }
136}