dope_core/backend/uring/driver/
mod.rs

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

dope_core/backend/uring/driver/
mod.rs
