Skip to main content

dope_core/backend/uring/driver/
mod.rs

1pub(crate) mod files;
2
3use crate::backend::uring::sqe;
4
5use std::io::{self, Error, ErrorKind};
6use std::os::fd::AsRawFd;
7use std::process::abort;
8
9use io_uring::IoUring;
10
11use self::files::FileTable;
12use crate::backend::uring::provided::ring::Ring;
13use crate::driver::route::Routes;
14use crate::driver::token::{
15    KIND_SHIFT, KeyTag, ROUTE_FRAMEWORK, SLOT_MASK, Token, TokenSlab, kind,
16};
17use crate::driver::{Config, PushError};
18use crate::io::fd::FdSlot;
19use crate::io::{BUFFER, BUFFER_SHIFT};
20use io_uring::squeue::Entry;
21use io_uring::types::{CancelBuilder, Timespec};
22use sqe::Sqe;
23
24use super::platform::gso::Gso;
25use crate::platform::raw::host::HOST;
26use crate::io::file::RawMetadata;
27use crate::platform::Platform;
28use crate::platform::snapshot::Snapshot;
29
30const SETSOCKOPT_CAP: usize = 4096;
31const _: () = assert!(SETSOCKOPT_CAP <= SLOT_MASK as usize + 1);
32type SetsockoptTag = KeyTag<ROUTE_FRAMEWORK, { kind::SETSOCKOPT }>;
33pub(crate) enum Disposition {
34    Drop,
35    DropBuffer(u16),
36    Internal,
37    Public(u64),
38}
39
40pub struct Uring {
41    pub(crate) uring: IoUring,
42    pub(crate) setsockopt: TokenSlab<libc::c_int, SetsockoptTag>,
43    pub(crate) files: FileTable,
44    pub(crate) provided: Ring,
45    next_slot: u32,
46    accept_slots: u32,
47    pub(crate) routes: Routes,
48}
49
50struct RingSetup {
51    cq_entries: u32,
52    ring_entries: u32,
53    defer_taskrun: bool,
54}
55
56impl RingSetup {
57    fn new(cq_entries: u32, defer_taskrun: bool, ring_entries: u32) -> Self {
58        Self {
59            cq_entries,
60            ring_entries,
61            defer_taskrun,
62        }
63    }
64
65    fn attempt(
66        &self,
67        no_sqarray: bool,
68        single_issuer: bool,
69        defer_taskrun: bool,
70    ) -> io::Result<IoUring> {
71        let mut builder = IoUring::builder();
72        builder.setup_submit_all();
73        if single_issuer {
74            builder.setup_single_issuer();
75        }
76        if no_sqarray {
77            builder.setup_no_sqarray();
78        }
79        builder.setup_cqsize(self.cq_entries);
80        if defer_taskrun {
81            builder.setup_defer_taskrun();
82        }
83        builder.build(self.ring_entries)
84    }
85
86    fn build(&self) -> io::Result<IoUring> {
87        let may_fallback = |error: &io::Error| error.raw_os_error() == Some(libc::EINVAL);
88        let mut built = self.attempt(true, true, self.defer_taskrun);
89        if built.as_ref().err().is_some_and(may_fallback) {
90            built = self.attempt(false, true, self.defer_taskrun);
91        }
92        if built.as_ref().err().is_some_and(may_fallback) {
93            built = self.attempt(false, false, false);
94        }
95        let uring = built?;
96
97        if !uring.params().is_feature_ext_arg() {
98            return Err(Error::new(ErrorKind::Unsupported, "IORING_FEAT_EXT_ARG"));
99        }
100        if !uring.params().is_feature_nodrop() {
101            return Err(Error::new(ErrorKind::Unsupported, "IORING_FEAT_NODROP"));
102        }
103        Ok(uring)
104    }
105
106    fn register_alloc_range(ring: &IoUring, len: u32) -> io::Result<()> {
107        #[repr(C)]
108        struct Range {
109            off: u32,
110            len: u32,
111            resv: u64,
112        }
113        const FILE_ALLOC_RANGE: libc::c_long = 25;
114        let range = Range {
115            off: 0,
116            len,
117            resv: 0,
118        };
119        let rc = unsafe {
120            libc::syscall(
121                libc::SYS_io_uring_register,
122                AsRawFd::as_raw_fd(ring) as libc::c_long,
123                FILE_ALLOC_RANGE,
124                &raw const range as usize as libc::c_long,
125                0 as libc::c_long,
126            )
127        };
128        if rc < 0 {
129            return Err(Error::last_os_error());
130        }
131        Ok(())
132    }
133}
134
135impl Uring {
136    pub(crate) fn new(cfg: &Config) -> io::Result<(Self, usize)> {
137        let uring = RingSetup::new(cfg.cq_entries, cfg.defer_taskrun, cfg.ring_entries).build()?;
138        match uring
139            .submitter()
140            .register_sync_cancel(None, CancelBuilder::user_data(u64::MAX).all())
141        {
142            Ok(()) => {}
143            Err(error) if error.kind() == ErrorKind::NotFound => {}
144            Err(error) => return Err(error),
145        }
146        uring
147            .submitter()
148            .register_files_sparse(cfg.fixed_file_slots)?;
149        RingSetup::register_alloc_range(&uring, cfg.accept_slots)?;
150        let provided = Ring::new(&uring.submitter(), cfg.provided.entries, cfg.provided.len)?;
151
152        let slots = cfg.fixed_file_slots as usize;
153        Ok((
154            Uring {
155                setsockopt: TokenSlab::with_capacity(SETSOCKOPT_CAP),
156                files: FileTable::new(slots),
157                uring,
158                provided,
159                next_slot: cfg.fixed_file_slots,
160                accept_slots: cfg.accept_slots,
161                routes: Routes::new(),
162            },
163            slots,
164        ))
165    }
166}
167
168impl Drop for Uring {
169    fn drop(&mut self) {
170        self.shutdown();
171    }
172}
173
174impl Uring {
175    pub(crate) fn shutdown(&mut self) {
176        let _ = self.uring.submit();
177        match self
178            .uring
179            .submitter()
180            .register_sync_cancel(None, CancelBuilder::any())
181        {
182            Ok(()) => {}
183            Err(error) if error.kind() == ErrorKind::NotFound => {}
184            Err(_) => abort(),
185        }
186        loop {
187            let mut drained = false;
188            {
189                let Self {
190                    uring,
191                    setsockopt,
192                    files,
193                    provided,
194                    routes,
195                    ..
196                } = self;
197                let mut cq = uring.completion();
198                for item in cq.by_ref() {
199                    drained = true;
200                    if let Disposition::DropBuffer(bid) = Self::complete_cqe(
201                        setsockopt,
202                        files,
203                        routes,
204                        item.user_data(),
205                        item.result(),
206                        item.flags(),
207                    ) {
208                        provided.defer(bid);
209                    }
210                }
211                cq.sync();
212            }
213            if !drained {
214                break;
215            }
216            let _ = self.uring.submit();
217        }
218    }
219
220    pub(crate) fn entry_push(uring: &mut IoUring, entry: &Entry) -> Result<(), PushError> {
221        if unsafe { uring.submission().push(entry) }.is_ok() {
222            return Ok(());
223        }
224        uring.submit().map_err(|_| PushError)?;
225        unsafe { uring.submission().push(entry) }.map_err(|_| PushError)
226    }
227
228    pub(crate) fn await_one(&mut self) -> io::Result<i32> {
229        loop {
230            self.uring.submitter().submit_and_wait(1)?;
231            let mut found = None;
232            {
233                let Self {
234                    uring,
235                    setsockopt,
236                    files,
237                    provided,
238                    routes,
239                    ..
240                } = self;
241                let mut cq = uring.completion();
242                for item in cq.by_ref() {
243                    let result = item.result();
244                    match Self::complete_cqe(
245                        setsockopt,
246                        files,
247                        routes,
248                        item.user_data(),
249                        result,
250                        item.flags(),
251                    ) {
252                        Disposition::Drop => {}
253                        Disposition::DropBuffer(bid) => provided.defer(bid),
254                        Disposition::Internal | Disposition::Public(_) => {
255                            found = Some(result);
256                        }
257                    }
258                }
259                cq.sync();
260            }
261            self.flush_deferred_close();
262            self.flush_ready_create();
263            if let Some(result) = found {
264                return Ok(result);
265            }
266        }
267    }
268
269    pub(crate) fn alloc_fixed_range(&mut self, len: u32) -> io::Result<u32> {
270        let base = self
271            .next_slot
272            .checked_sub(len)
273            .filter(|&b| b >= self.accept_slots)
274            .ok_or_else(|| {
275                Error::new(ErrorKind::OutOfMemory, "dope: fixed-file slots exhausted")
276            })?;
277        self.next_slot = base;
278        Ok(base)
279    }
280
281    fn release_setsockopt(
282        setsockopt: &mut TokenSlab<libc::c_int, SetsockoptTag>,
283        token: Token,
284    ) -> bool {
285        let Some(parts) = token.parts::<SetsockoptTag>() else {
286            return false;
287        };
288        setsockopt.remove_parts(parts.slab());
289        true
290    }
291
292    pub(crate) fn complete_cqe(
293        setsockopt: &mut TokenSlab<libc::c_int, SetsockoptTag>,
294        files: &mut FileTable,
295        routes: &Routes,
296        user_data: u64,
297        result: i32,
298        flags: u32,
299    ) -> Disposition {
300        let Some(token) = Token::try_from_raw(user_data) else {
301            return if flags & BUFFER != 0 {
302                Disposition::DropBuffer((flags >> BUFFER_SHIFT) as u16)
303            } else {
304                Disposition::Drop
305            };
306        };
307        if Self::release_setsockopt(setsockopt, token) {
308            return Disposition::Internal;
309        }
310        let op_kind = (user_data >> KIND_SHIFT) as u8;
311        if token.route() == ROUTE_FRAMEWORK {
312            return match op_kind {
313                kind::CLOSE_PREP => Disposition::Drop,
314                kind::CLOSE => {
315                    files.complete_close(token.slot());
316                    Disposition::Drop
317                }
318                kind::CREATE => files
319                    .complete_create(token.slot(), result)
320                    .map_or(Disposition::Drop, Disposition::Public),
321                _ => Disposition::Public(user_data),
322            };
323        }
324        if op_kind == kind::ACCEPT {
325            let poisoned = routes.is_poisoned(token.route());
326            files.mark_accepted(result, poisoned);
327            if poisoned {
328                return Disposition::Drop;
329            }
330        } else if op_kind == kind::OPEN && result >= 0 && routes.is_poisoned(token.route()) {
331            unsafe {
332                libc::close(result);
333            }
334            return Disposition::Drop;
335        } else if op_kind == kind::RECV && flags & BUFFER != 0 && routes.is_poisoned(token.route())
336        {
337            return Disposition::DropBuffer((flags >> BUFFER_SHIFT) as u16);
338        }
339        Disposition::Public(user_data)
340    }
341
342    fn push_close_at(uring: &mut IoUring, slot: FdSlot) -> bool {
343        let shut = Sqe::shutdown_linked_at(slot, libc::SHUT_RDWR);
344        let close = Sqe::close_at(slot);
345        let entries = [shut.entry().clone(), close.entry().clone()];
346        unsafe { uring.submission().push_multiple(&entries) }.is_ok()
347    }
348
349    pub(crate) fn flush_deferred_close(&mut self) {
350        let Self { uring, files, .. } = self;
351        files.flush_deferred_close(|slot| Self::push_close_at(uring, slot));
352    }
353
354    pub(crate) fn flush_ready_create(&mut self) {
355        let Self { uring, files, .. } = self;
356        files.flush_ready(|sqe| Self::entry_push(uring, sqe.entry()).is_ok());
357    }
358
359    pub(crate) fn close_fd(&mut self, slot: FdSlot) {
360        let Self { uring, files, .. } = self;
361        files.release(slot, |slot| {
362            Self::push_close_at(uring, slot)
363                || (uring.submit().is_ok() && Self::push_close_at(uring, slot))
364        });
365    }
366}
367
368impl Platform for Uring {
369    type Sqe = Sqe;
370    type Gso = Gso;
371    type StatBuf = libc::statx;
372    type TimerSpec = Timespec;
373
374    fn entropy() -> io::Result<[u64; 2]> {
375        HOST.entropy()
376    }
377
378    fn parse_meta(raw: &Self::StatBuf) -> io::Result<RawMetadata> {
379        HOST.parse_meta(raw)
380    }
381
382    fn snapshot() -> io::Result<Snapshot> {
383        HOST.snapshot()
384    }
385}