Skip to main content

dope_uring/driver/
mod.rs

1mod buf;
2pub mod config;
3mod cqe;
4mod current;
5pub mod handle;
6mod provided;
7mod recv_msg_out;
8mod request;
9mod resources;
10mod ring;
11mod submit;
12mod wire;
13
14use std::io;
15use std::marker::PhantomData;
16use std::os::fd::AsRawFd;
17use std::time::Duration;
18
19use io_uring::{opcode, types};
20
21use dope_core::addr::SockAddr;
22use dope_core::driver::DriverOps;
23use dope_core::submit::SendMsgPacket;
24
25use crate::sys::Sqe;
26use crate::sys::fd::FdUringExt;
27use crate::sys::token::Token;
28use crate::{Fd, FixedFd};
29
30pub use self::buf::{FixedBufGuard, FixedBufWrite, ProvidedBuf};
31pub use self::cqe::{CqeKind, CqeRouter, CqeView};
32pub use self::current::{CurrentSlot, current, current_driver_mut, set_current};
33pub use self::provided::ProvidedRing;
34pub use self::recv_msg_out::RecvMsgOut;
35pub use self::request::Request;
36
37pub type DriverRef<'a> = &'a Driver;
38
39pub(crate) use self::request::{RequestId, Requests};
40pub(crate) use self::resources::Resources;
41pub use self::ring::Ring;
42
43use self::cqe::{RawCqe, make_view};
44use self::wire::Wire;
45
46pub struct Driver {
47    pub(crate) requests: Requests,
48    pub(crate) ring: Ring,
49    pub(crate) resources: Resources,
50    pub(crate) provided: ProvidedRing,
51    pub(in crate::driver) cqe_scratch: Vec<self::cqe::RawCqe>,
52
53    _not_send: PhantomData<*const ()>,
54}
55
56impl Drop for Driver {
57    fn drop(&mut self) {
58        let self_ptr = self as *mut Driver;
59        if let Some(slot) = self::current::current()
60            && slot.driver().as_ptr() == self_ptr
61        {
62            self::current::set_current(None);
63        }
64    }
65}
66
67impl Driver {
68    pub(crate) fn new(ring: Ring, cfg: crate::DriverConfig, capacity: usize) -> io::Result<Self> {
69        let provided = ProvidedRing::new(&ring, cfg.provided_buf_entries, cfg.provided_buf_len)
70            .map_err(|e| {
71                io::Error::other(format!(
72                    "dope: required capability setup failed (provided buffer ring): {e}"
73                ))
74            })?;
75        let resources = Resources::new(&ring, &cfg)?;
76        Ok(Self {
77            requests: Requests::with_capacity(capacity),
78            ring,
79            resources,
80            provided,
81            cqe_scratch: Vec::with_capacity(capacity),
82            _not_send: PhantomData,
83        })
84    }
85
86    #[inline(always)]
87    pub fn has_pending(&mut self) -> bool {
88        self.ring.has_submission() || self.ring.has_completion()
89    }
90
91    #[inline(always)]
92    pub fn has_completion(&mut self) -> bool {
93        self.ring.has_completion()
94    }
95
96    #[inline(always)]
97    pub fn park(&mut self) -> io::Result<()> {
98        self.ring.submit_and_wait(1)
99    }
100
101    #[inline(always)]
102    pub fn park_for(&mut self, duration: Duration) -> io::Result<()> {
103        self.ring.submit_with_timeout(1, duration)
104    }
105}
106
107impl Driver {
108    fn submit_retry_error(e: &io::Error) -> bool {
109        matches!(e.raw_os_error(), Some(libc::EAGAIN) | Some(libc::EBUSY))
110    }
111
112    pub(super) fn submit_retry(&mut self) -> io::Result<usize> {
113        match self.ring.submit() {
114            Ok(submitted) => Ok(submitted),
115            Err(ref e) if Self::submit_retry_error(e) => {
116                self.tick()?;
117                Ok(0)
118            }
119            Err(e) => Err(e),
120        }
121    }
122
123    #[inline(always)]
124    fn drive_loop(
125        &mut self,
126        mut on_completion: impl FnMut(&mut Self) -> io::Result<()>,
127    ) -> io::Result<bool> {
128        let mut any = false;
129        for _ in 0..256 {
130            let mut progressed = false;
131            if self.ring.has_submission() {
132                let submitted = self.submit_retry()?;
133                if submitted != 0 {
134                    progressed = true;
135                }
136            }
137            if self.ring.has_completion() {
138                on_completion(self)?;
139                progressed = true;
140            }
141            if !progressed {
142                break;
143            }
144            any = true;
145        }
146        Ok(any)
147    }
148
149    pub fn drive(&mut self) -> io::Result<bool> {
150        self.drive_loop(Self::tick)
151    }
152
153    pub fn drive_with<H: CqeRouter>(&mut self, handler: &mut H) -> io::Result<bool> {
154        self.drive_loop(|s| {
155            s.dispatch_cqes(handler);
156            Ok(())
157        })
158    }
159
160    pub(crate) fn tick(&mut self) -> io::Result<()> {
161        let Self {
162            requests,
163            ring,
164            provided,
165            ..
166        } = self;
167        let mut cq = ring.completion();
168        for item in &mut cq {
169            let result = item.result();
170            let flags = item.flags();
171            if let Wire::Request(id) = wire::decode_wire(item.user_data()) {
172                requests.on_request(id, result, flags);
173            }
174        }
175        cq.sync();
176        let _ = provided.flush();
177        Ok(())
178    }
179
180    #[inline(always)]
181    pub fn cancel_write(&mut self, token: Token) {
182        self.ring.submit_cancel(wire::write(token));
183    }
184
185    #[inline(always)]
186    pub fn cancel_accept(&mut self, token: Token) {
187        self.ring.submit_cancel(wire::accept(token));
188    }
189
190    fn drain_to_scratch(&mut self) {
191        let Self {
192            ring, cqe_scratch, ..
193        } = self;
194        cqe_scratch.clear();
195        let mut cq = ring.completion();
196        for item in &mut cq {
197            cqe_scratch.push(RawCqe {
198                user_data: item.user_data(),
199                result: item.result(),
200                flags: item.flags(),
201            });
202        }
203        cq.sync();
204    }
205
206    pub fn dispatch_cqes<H: CqeRouter>(&mut self, handler: &mut H) {
207        self.drain_to_scratch();
208        let mut scratch = std::mem::take(&mut self.cqe_scratch);
209        for raw in scratch.drain(..) {
210            let wire = wire::decode_wire(raw.user_data);
211            match wire {
212                Wire::Request(id) => self.requests.on_request(id, raw.result, raw.flags),
213                Wire::Ignore => {}
214                Wire::Accept(_) | Wire::Recv(_) | Wire::Write(_) => {
215                    handler.on_cqe(self, make_view(wire, raw.result, raw.flags));
216                }
217            }
218        }
219        self.cqe_scratch = scratch;
220        let _ = self.provided.flush();
221    }
222}
223
224impl Driver {
225    pub fn submit_only(&mut self) -> io::Result<()> {
226        if self.ring.submission_len() == 0 {
227            return Ok(());
228        }
229        for _ in 0..256 {
230            if self.ring.submission_len() == 0 {
231                return Ok(());
232            }
233            self.submit_retry()?;
234        }
235        Ok(())
236    }
237
238    pub(crate) fn submit_sqe_op(&mut self, sqe: crate::Sqe) -> io::Result<RequestId> {
239        if self.ring.is_full() {
240            self.submit_only()?;
241        }
242        let index = self.requests.insert();
243        let sqe = sqe.user_data(wire::request(index));
244        if !self.ring.push(&sqe) {
245            self.requests.rollback(index);
246            return Err(io::Error::new(
247                io::ErrorKind::WouldBlock,
248                "SQ full after flush",
249            ));
250        }
251        Ok(index)
252    }
253}
254
255impl Driver {
256    #[inline(always)]
257    pub(crate) fn try_submit_op<O, F>(
258        &mut self,
259        op: O,
260        sqe_builder: F,
261    ) -> Result<Request<O>, (io::Error, O)>
262    where
263        O: Unpin + 'static,
264        F: FnOnce(&mut O) -> Sqe,
265    {
266        let mut boxed = Box::new(op);
267        let sqe = sqe_builder(&mut *boxed);
268        Request::submit_boxed(self, boxed, sqe).map_err(|(e, b)| (e, *b))
269    }
270
271    #[inline(always)]
272    pub(crate) fn try_submit_op_boxed<O, F>(
273        &mut self,
274        mut boxed: Box<O>,
275        sqe_builder: F,
276    ) -> Result<Request<O>, (io::Error, Box<O>)>
277    where
278        O: Unpin + 'static,
279        F: FnOnce(&mut O) -> Sqe,
280    {
281        let sqe = sqe_builder(&mut *boxed);
282        Request::submit_boxed(self, boxed, sqe)
283    }
284
285    pub fn try_submit_sendmsg_raw_boxed<P>(
286        &mut self,
287        fd: Fd,
288        packet: Box<P>,
289    ) -> Result<Request<P>, (io::Error, Box<P>)>
290    where
291        P: SendMsgPacket + Unpin + 'static,
292    {
293        let raw = fd.as_raw_fd();
294        self.try_submit_op_boxed(packet, move |packet| {
295            io_uring::opcode::SendMsg::new(io_uring::types::Fd(raw), packet.msg_ptr()).build()
296        })
297    }
298}
299
300impl DriverOps for Driver {
301    type Token = Token;
302    type ProvidedRing = ProvidedRing;
303    type IoRequest<O: Unpin + 'static> = Request<O>;
304
305    #[inline(always)]
306    fn provided_ring(&mut self) -> &mut Self::ProvidedRing {
307        &mut self.provided
308    }
309
310    #[inline(always)]
311    fn buf_group(&self) -> u16 {
312        self.provided.group().0
313    }
314
315    #[inline(always)]
316    fn arm_recv_multi(&mut self, token: Self::Token, fd: FixedFd, buf_group: u16) -> bool {
317        let sqe = opcode::RecvMulti::new(types::Fixed(fd.fixed_index()), buf_group)
318            .build()
319            .user_data(wire::recv(token));
320        self.ring.try_push_arm(sqe)
321    }
322
323    #[inline(always)]
324    unsafe fn arm_recv_msg_multi(
325        &mut self,
326        token: Self::Token,
327        fd: FixedFd,
328        buf_group: u16,
329        msghdr: *const libc::msghdr,
330    ) -> bool {
331        let sqe = opcode::RecvMsgMulti::new(types::Fixed(fd.fixed_index()), msghdr, buf_group)
332            .build()
333            .user_data(wire::recv(token));
334        self.ring.try_push_arm(sqe)
335    }
336
337    #[inline(always)]
338    fn arm_accept_multi(&mut self, token: Self::Token, fd: Fd) -> bool {
339        let sqe = opcode::AcceptMulti::new(fd.uring())
340            .flags(libc::SOCK_CLOEXEC)
341            .build()
342            .user_data(wire::accept(token));
343        self.ring.try_push_arm(sqe)
344    }
345
346    #[inline(always)]
347    fn submit_send_tagged(
348        &mut self,
349        token: Self::Token,
350        fd: FixedFd,
351        ptr: *const u8,
352        len: u32,
353    ) -> bool {
354        let sqe = opcode::Send::new(types::Fixed(fd.fixed_index()), ptr, len)
355            .build()
356            .user_data(wire::write(token));
357        self.ring.push(&sqe)
358    }
359
360    #[inline(always)]
361    unsafe fn submit_send_msg_tagged(
362        &mut self,
363        token: Self::Token,
364        fd: FixedFd,
365        msg: *const libc::msghdr,
366    ) -> bool {
367        let sqe = opcode::SendMsg::new(types::Fixed(fd.fixed_index()), msg)
368            .build()
369            .user_data(wire::write(token));
370        self.ring.push(&sqe)
371    }
372
373    #[inline(always)]
374    fn cancel_recv(&mut self, token: Self::Token) {
375        self.ring.submit_cancel(wire::recv(token));
376    }
377
378    fn try_submit_connect_raw(
379        &mut self,
380        fd: Fd,
381        addr: SockAddr,
382    ) -> Result<Self::IoRequest<SockAddr>, (io::Error, SockAddr)> {
383        let raw = fd.as_raw_fd();
384        self.try_submit_op(addr, move |addr| {
385            io_uring::opcode::Connect::new(
386                io_uring::types::Fd(raw),
387                addr.addr_ptr(),
388                addr.addr_len(),
389            )
390            .build()
391        })
392    }
393
394    fn try_submit_sendmsg_raw<P: SendMsgPacket + Unpin + 'static>(
395        &mut self,
396        fd: Fd,
397        packet: P,
398    ) -> Result<Self::IoRequest<P>, (io::Error, P)> {
399        let raw = fd.as_raw_fd();
400        self.try_submit_op(packet, move |packet| {
401            io_uring::opcode::SendMsg::new(io_uring::types::Fd(raw), packet.msg_ptr()).build()
402        })
403    }
404
405    fn try_submit_send_raw<B: AsRef<[u8]> + Unpin + 'static>(
406        &mut self,
407        fd: Fd,
408        buf: B,
409    ) -> Result<Self::IoRequest<B>, (io::Error, B)> {
410        let raw = fd.as_raw_fd();
411        self.try_submit_op(buf, move |buf| {
412            let slice = buf.as_ref();
413            io_uring::opcode::Send::new(
414                io_uring::types::Fd(raw),
415                slice.as_ptr(),
416                slice.len() as u32,
417            )
418            .build()
419        })
420    }
421}