Repository: tokio-rs/tokio
File: tokio/src/runtime/io/driver/uring.rs (Rust)

Code before:
        let mut ops = std::mem::take(&mut self.ops);

        // Remove all completed ops since we don't need to wait for them.
        ops.retain(|_, lifecycle| !matches!(lifecycle, Lifecycle::Completed(_)));

        while !ops.is_empty() {
            // Wait until at least one completion is available.
            self.ring_mut()
                .submit_and_wait(1)
                .expect("Internal error when dropping driver");

            for cqe in self.ring_mut().completion() {
                let idx = cqe.user_data() as usize;
                ops.remove(idx);
            }
        }
    }
}

impl Handle {
    fn add_uring_source(&self, uringfd: RawFd) -> io::Result<()> {
        let mut source = SourceFd(&uringfd);
        self.registry
            .register(&mut source, TOKEN_WAKEUP, Interest::READABLE.to_mio())
    }

    pub(crate) fn get_uring(&self) -> &Mutex<UringContext> {
        &self.uring_context
    }

    fn set_uring_state(&self, state: State) {
        self.uring_state.store(state.as_usize(), Ordering::Release);
    }

    /// Check if the io_uring context is initialized. If not, it will try to initialize it.
    /// Then, check if the provided opcode is supported.
    ///
    /// If both the context initialization succeeds and the opcode is supported,
    /// this returns `Ok(true)`.
    /// If either io_uring is unsupported or the opcode is unsupported,
    /// this returns `Ok(false)`.
    /// An error is returned if an io_uring syscall returns an unexpected error value.
    pub(crate) fn check_and_init(&self, opcode: u8) -> io::Result<bool> {
        match State::from_usize(self.uring_state.load(Ordering::Acquire)) {
            State::Uninitialized => match self.try_init_and_check_opcode(opcode) {
                Ok(opcode_supported) => {
                    self.set_uring_state(State::Initialized);
                    Ok(opcode_supported)
                }
                // If the system doesn't support io_uring, we set the state to Unsupported.
                Err(e) if e.raw_os_error() == Some(libc::ENOSYS) => {

Review comment (suggestion):
Maybe a dumb question, is this `.as_ref` required? It should be ok to consume the `probe`.

Code after:
    /// Check if the io_uring context is initialized. If not, it will try to initialize it.
    /// Then, check if the provided opcode is supported.
    ///
    /// If both the context initialization succeeds and the opcode is supported,
    /// this returns `Ok(true)`.
    /// If either io_uring is unsupported or the opcode is unsupported,
    /// this returns `Ok(false)`.
    /// An error is returned if an io_uring syscall returns an unexpected error value.
    pub(crate) fn check_and_init(&self, opcode: u8) -> io::Result<bool> {
        let probe = self.uring_probe.get_or_try_init(|| {
            let mut probe = Probe::new();
            match self.try_init(&mut probe) {
                Ok(()) => Ok(Some(probe)),
                // If the system doesn't support io_uring, we set the probe to `None`.
                Err(e) if e.raw_os_error() == Some(libc::ENOSYS) => Ok(None),
                // If we get EPERM, io-uring syscalls may be blocked (for example, by seccomp).
                // In this case, we try to fall back to spawn_blocking for this and future operations.
                // See also: https://github.com/tokio-rs/tokio/issues/7691
                Err(e) if e.raw_os_error() == Some(libc::EPERM) => Ok(None),
                // For other system errors, we just return it.
                Err(e) => Err(e),
            }
        })?;

        Ok(probe
            .as_ref()
            .is_some_and(|probe| probe.is_supported(opcode)))
    }

    /// Initialize the io_uring context if it hasn't been initialized yet.
    fn try_init(&self, probe: &mut Probe) -> io::Result<()> {
        let mut guard = self.get_uring().lock();
        if guard.try_init(probe)? {
            self.add_uring_source(guard.ring().as_raw_fd())?;
        }

        Ok(())
    }

    /// Register an operation with the io_uring.
    ///
    /// If this is the first io_uring operation, it will also initialize the io_uring context.
    /// If io_uring isn't supported, this function returns an `ENOSYS` error, so the caller can
    /// perform custom handling, such as falling back to an alternative mechanism.
    ///
    /// # Safety
    ///
    /// Callers must ensure that parameters of the entry (such as buffer) are valid and will
    /// be valid for the entire duration of the operation, otherwise it may cause memory problems.
    pub(crate) unsafe fn register_op(&self, entry: Entry, waker: Waker) -> io::Result<usize> {
        // Note: Maybe this check can be removed if upstream callers consistently use `check_and_init`.