block: Add io_uring async I/O struct

Add the shared helpers and UringDataIo queue that drive uring async I/O
operations. UringDataIo is the key component responsible for keeping the
memory pointers active while async operations are ongoing. It uses the
async core added in previous commits.

Later commits will change the block backends to use this instead of the
lower level abstractions directly.

Signed-off-by: Dylan Reid <dgreid@fb.com>

block/src/async_io.rs

@@ -4,16 +4,21 @@
 
 use std::marker::PhantomData;
 use std::os::fd::{AsRawFd, OwnedFd, RawFd};
+mod common;
 mod completion;
 mod guest_memory_target;
 mod operation;
 mod owned_io_buffer;
+#[cfg(feature = "io_uring")]
+mod uring_data_io;
 
 pub use completion::AsyncIoCompletion;
 pub use guest_memory_target::GuestMemoryTarget;
 pub use operation::AsyncIoOperation;
 pub use owned_io_buffer::OwnedIoBuffer;
 use thiserror::Error;
+#[cfg(feature = "io_uring")]
+pub use uring_data_io::UringDataIo;
 use vmm_sys_util::eventfd::EventFd;
 
 use crate::{BatchRequest, SECTOR_SIZE};

block/src/async_io/common.rs

@@ -0,0 +1,41 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
+
+//! Helpers used by both aio and uring async io.
+
+use std::collections::HashSet;
+use std::io;
+
+use super::AsyncIoOperation;
+
+/// Converts an I/O error into the negative errno form used in completions.
+#[allow(dead_code)]
+pub(super) fn errno_result(error: &io::Error) -> i32 {
+    -error.raw_os_error().unwrap_or(libc::EIO)
+}
+
+/// Builds the error returned when a new request reuses in-flight `user_data`.
+pub(super) fn duplicate_user_data_error(user_data: u64) -> io::Error {
+    io::Error::new(
+        io::ErrorKind::AlreadyExists,
+        format!("duplicate async I/O user_data {user_data}"),
+    )
+}
+
+/// Validates that a batch has unique `user_data` not already in flight.
+pub(super) fn validate_batch<F>(mut is_in_flight: F, batch: &[AsyncIoOperation]) -> io::Result<()>
+where
+    F: FnMut(u64) -> bool,
+{
+    let mut seen = HashSet::with_capacity(batch.len());
+
+    for op in batch {
+        let user_data = op.user_data();
+        if is_in_flight(user_data) || !seen.insert(user_data) {
+            return Err(duplicate_user_data_error(user_data));
+        }
+    }
+
+    Ok(())
+}

block/src/async_io/uring_data_io.rs

@@ -0,0 +1,357 @@
+// Copyright © 2021 Intel Corporation
+//
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
+
+use std::collections::{HashMap, VecDeque};
+use std::io;
+use std::os::fd::{AsRawFd, RawFd};
+
+use io_uring::{IoUring, opcode, squeue, types};
+use log::warn;
+use vmm_sys_util::eventfd::EventFd;
+
+use super::common::{duplicate_user_data_error, validate_batch};
+use super::{AsyncIoCompletion, AsyncIoOperation};
+
+/// `io_uring` wrapper for async I/O.
+///
+/// Holds the `IoUring` and its `EventFd`. Tracks ops that are pending.
+pub struct UringDataIo {
+    // Keep `io_uring` before `in_flight`. Rust drops fields in declaration
+    // order, so dropping the ring tears down kernel access before retained
+    // operations release the buffers referenced by their iovecs.
+    io_uring: IoUring,
+    // The `EventFd` for completion signals.
+    eventfd: EventFd,
+    // `in_flight` tracks every user_data value accepted by the kernel. Owned
+    // data operations store `Some(op)` so their iovecs and backing buffers
+    // remain valid until completion; metadata operations store `None`.
+    in_flight: HashMap<u64, Option<AsyncIoOperation>>,
+    // `injected` holds locally produced completions so synchronous failures
+    // and short-circuited requests use the same drain path as kernel CQEs.
+    injected: VecDeque<AsyncIoCompletion>,
+    // `needs_submit_retry` is set when SQEs have been published to the ring,
+    // but the submit syscall failed before confirming kernel ownership.
+    needs_submit_retry: bool,
+}
+
+impl UringDataIo {
+    /// Creates an io_uring queue and registers its completion eventfd.
+    pub fn new(ring_depth: u32) -> io::Result<Self> {
+        let io_uring = IoUring::new(ring_depth)?;
+        let eventfd = EventFd::new(libc::EFD_NONBLOCK)?;
+        io_uring.submitter().register_eventfd(eventfd.as_raw_fd())?;
+
+        Ok(Self {
+            io_uring,
+            eventfd,
+            in_flight: HashMap::new(),
+            injected: VecDeque::new(),
+            needs_submit_retry: false,
+        })
+    }
+
+    /// Returns the eventfd signaled when completions are available.
+    pub fn notifier(&self) -> &EventFd {
+        &self.eventfd
+    }
+
+    /// Submits one owned read or write operation to the queue.
+    pub fn submit_operation(&mut self, fd: RawFd, op: AsyncIoOperation) -> io::Result<()> {
+        self.submit_batch(fd, vec![op])
+    }
+
+    fn reserve_user_data(&mut self, user_data: u64) -> io::Result<()> {
+        if self.in_flight.contains_key(&user_data) {
+            return Err(duplicate_user_data_error(user_data));
+        }
+        self.in_flight.insert(user_data, None);
+
+        Ok(())
+    }
+
+    fn submit_kernel_entry(&mut self, user_data: u64, entry: &squeue::Entry) -> io::Result<()> {
+        self.reserve_user_data(user_data)?;
+
+        let (submitter, mut sq, _) = self.io_uring.split();
+        // SAFETY: the entry has no caller-owned buffer. `user_data` is retained
+        // in `in_flight` until the CQE is consumed.
+        if let Err(e) = unsafe { sq.push(entry) } {
+            self.in_flight.remove(&user_data);
+            return Err(io::Error::other(format!("Submission queue is full: {e:?}")));
+        }
+        sq.sync();
+
+        match submitter.submit() {
+            Ok(_) => self.needs_submit_retry = false,
+            Err(e) => {
+                self.needs_submit_retry = true;
+                warn!("io_uring submit failed after SQE was published: {e}");
+                self.eventfd.write(1).unwrap();
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Submits a batch of owned read and write operations.
+    ///
+    /// If the io_uring submission queue cannot accept the whole batch, each
+    /// operation is completed locally with `-EAGAIN` so callers can observe
+    /// every request through the normal completion path.
+    pub fn submit_batch(&mut self, fd: RawFd, batch: Vec<AsyncIoOperation>) -> io::Result<()> {
+        if batch.is_empty() {
+            return Ok(());
+        }
+
+        validate_batch(|user_data| self.in_flight.contains_key(&user_data), &batch)?;
+
+        let (submitter, mut sq, _) = self.io_uring.split();
+        let available = sq.capacity() - sq.len();
+        if batch.len() > available {
+            // Not enough space for the batch.
+            // Drop sq, which will re-publish an unmodified tail pointer
+            drop(sq);
+            for op in batch {
+                self.injected
+                    .push_back(AsyncIoCompletion::from_operation(op, -libc::EAGAIN));
+            }
+            self.eventfd.write(1).unwrap();
+            return Ok(());
+        }
+
+        let mut signal_completion = false;
+        let mut batch = batch.into_iter();
+        while let Some(op) = batch.next() {
+            let user_data = op.user_data();
+            let entry = Self::build_entry(fd, &op);
+            self.in_flight.insert(user_data, Some(op));
+
+            // SAFETY: the SQ capacity was just checked. Every iovec's pointer is retained in
+            // self.in_flight before the SQ tail is advanced by sync or drop. in_flight only
+            // drops the memory after a completion.
+            if let Err(e) = unsafe { sq.push(&entry) } {
+                Self::handle_push_failure(
+                    &mut self.in_flight,
+                    &mut self.injected,
+                    user_data,
+                    batch.by_ref(),
+                    &e,
+                );
+                signal_completion = true;
+                break;
+            }
+        }
+
+        sq.sync();
+        match submitter.submit() {
+            Ok(_) => self.needs_submit_retry = false,
+            Err(e) => {
+                self.needs_submit_retry = true;
+                warn!("io_uring submit failed after SQEs were published: {e}");
+                signal_completion = true;
+            }
+        }
+        if signal_completion {
+            self.eventfd.write(1).unwrap();
+        }
+
+        Ok(())
+    }
+
+    #[cold]
+    fn handle_push_failure(
+        in_flight: &mut HashMap<u64, Option<AsyncIoOperation>>,
+        injected: &mut VecDeque<AsyncIoCompletion>,
+        user_data: u64,
+        remaining: impl Iterator<Item = AsyncIoOperation>,
+        error: &squeue::PushError,
+    ) {
+        // Since capacity was just checked, this should only happen if the ring
+        // state changed unexpectedly. Keep all affected operations memory safe
+        // by returning local completions through the normal path.
+        let op = in_flight
+            .remove(&user_data)
+            .flatten()
+            .expect("pending operation missing after failed push");
+        injected.push_back(AsyncIoCompletion::from_operation(op, -libc::EAGAIN));
+        for op in remaining {
+            injected.push_back(AsyncIoCompletion::from_operation(op, -libc::EAGAIN));
+        }
+        warn!("io_uring submission queue became full after capacity check: {error:?}");
+    }
+
+    fn build_entry(fd: RawFd, op: &AsyncIoOperation) -> squeue::Entry {
+        let iovecs = op.iovecs();
+        let fd = types::Fd(fd);
+        if op.is_read() {
+            opcode::Readv::new(fd, iovecs.as_ptr(), iovecs.len() as u32)
+                .offset(op.offset() as u64)
+                .build()
+                .user_data(op.user_data())
+        } else {
+            opcode::Writev::new(fd, iovecs.as_ptr(), iovecs.len() as u32)
+                .offset(op.offset() as u64)
+                .build()
+                .user_data(op.user_data())
+        }
+    }
+
+    /// Submits an io_uring NOP carrying `user_data`.
+    pub fn submit_nop(&mut self, user_data: u64) -> io::Result<()> {
+        self.submit_kernel_entry(user_data, &opcode::Nop::new().build().user_data(user_data))
+    }
+
+    /// Submits an fsync operation carrying `user_data`.
+    pub fn submit_fsync(&mut self, fd: RawFd, user_data: u64) -> io::Result<()> {
+        self.submit_kernel_entry(
+            user_data,
+            &opcode::Fsync::new(types::Fd(fd))
+                .build()
+                .user_data(user_data),
+        )
+    }
+
+    /// Submits a fallocate operation carrying `user_data`.
+    pub fn submit_fallocate(
+        &mut self,
+        fd: RawFd,
+        offset: u64,
+        length: u64,
+        mode: i32,
+        user_data: u64,
+    ) -> io::Result<()> {
+        self.submit_kernel_entry(
+            user_data,
+            &opcode::Fallocate::new(types::Fd(fd), length)
+                .offset(offset)
+                .mode(mode)
+                .build()
+                .user_data(user_data),
+        )
+    }
+
+    /// Injects a completion that did not come from a kernel CQE.
+    ///
+    /// The notifier is signaled so callers can drain it with
+    /// [`Self::next_completion`].
+    pub fn inject_completion(&mut self, completion: AsyncIoCompletion) {
+        self.injected.push_back(completion);
+        self.eventfd.write(1).unwrap();
+    }
+
+    /// Returns the next kernel or injected completion if one is available.
+    ///
+    /// Consuming a kernel completion returns ownership of any buffer retained
+    /// by the corresponding operation.
+    pub fn next_completion(&mut self) -> Option<AsyncIoCompletion> {
+        if self.needs_submit_retry {
+            match self.io_uring.submitter().submit() {
+                Ok(_) => self.needs_submit_retry = false,
+                Err(e) => warn!("io_uring retry submit failed for retained SQEs: {e}"),
+            }
+        }
+
+        if let Some(entry) = self.io_uring.completion().next() {
+            let user_data = entry.user_data();
+            return Some(AsyncIoCompletion::new(
+                user_data,
+                entry.result(),
+                self.in_flight
+                    .remove(&user_data)
+                    .flatten()
+                    .and_then(AsyncIoOperation::into_completion_buffer),
+            ));
+        }
+
+        self.injected.pop_front()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::os::fd::AsRawFd;
+    use std::thread::sleep;
+    use std::time::Duration;
+
+    use vmm_sys_util::tempfile::TempFile;
+
+    use super::UringDataIo;
+    use crate::async_io::{AsyncIoCompletion, AsyncIoOperation, OwnedIoBuffer};
+
+    fn wait_for_completion(data_io: &mut UringDataIo) -> AsyncIoCompletion {
+        for _ in 0..1000 {
+            if let Some(completion) = data_io.next_completion() {
+                return completion;
+            }
+            sleep(Duration::from_millis(1));
+        }
+
+        panic!("timed out waiting for io_uring completion");
+    }
+
+    #[test]
+    fn uring_rejects_duplicate_user_data_for_metadata_ops() {
+        let file = TempFile::new().unwrap().into_file();
+        file.set_len(512).unwrap();
+        let fd = file.as_raw_fd();
+        let mut data_io = UringDataIo::new(8).unwrap();
+
+        data_io
+            .submit_operation(
+                fd,
+                AsyncIoOperation::read_to_vec(0, OwnedIoBuffer::from_vec(vec![0; 512]), 7),
+            )
+            .unwrap();
+
+        assert_eq!(
+            data_io.submit_fsync(fd, 7).unwrap_err().kind(),
+            std::io::ErrorKind::AlreadyExists
+        );
+        assert_eq!(
+            data_io.submit_nop(7).unwrap_err().kind(),
+            std::io::ErrorKind::AlreadyExists
+        );
+        assert_eq!(
+            data_io
+                .submit_fallocate(fd, 0, 512, 0, 7)
+                .unwrap_err()
+                .kind(),
+            std::io::ErrorKind::AlreadyExists
+        );
+
+        let completion = wait_for_completion(&mut data_io);
+        assert_eq!(completion.user_data, 7);
+        assert_eq!(completion.result, 512);
+    }
+
+    #[test]
+    fn uring_queue_full_batch_completes_each_operation() {
+        let file = TempFile::new().unwrap().into_file();
+        let fd = file.as_raw_fd();
+        let mut data_io = UringDataIo::new(1).unwrap();
+        let available = {
+            let (_, sq, _) = data_io.io_uring.split();
+            sq.capacity() - sq.len()
+        };
+        let batch_len = available + 1;
+        let batch: Vec<_> = (0..batch_len as u64)
+            .map(|user_data| {
+                AsyncIoOperation::read_to_vec(0, OwnedIoBuffer::from_vec(vec![0; 512]), user_data)
+            })
+            .collect();
+
+        data_io.submit_batch(fd, batch).unwrap();
+
+        let mut completed = Vec::new();
+        while let Some(completion) = data_io.next_completion() {
+            assert_eq!(completion.result, -libc::EAGAIN);
+            assert!(completion.buffer.is_some());
+            completed.push(completion.user_data);
+        }
+        completed.sort_unstable();
+        assert_eq!(completed, (0..batch_len as u64).collect::<Vec<_>>());
+    }
+}