diff --git a/block/src/async_io.rs b/block/src/async_io.rs index b60b3fc92..7337ade54 100644 --- a/block/src/async_io.rs +++ b/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}; diff --git a/block/src/async_io/common.rs b/block/src/async_io/common.rs new file mode 100644 index 000000000..9f631a110 --- /dev/null +++ b/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(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(()) +} diff --git a/block/src/async_io/uring_data_io.rs b/block/src/async_io/uring_data_io.rs new file mode 100644 index 000000000..97e4199e7 --- /dev/null +++ b/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>, + // `injected` holds locally produced completions so synchronous failures + // and short-circuited requests use the same drain path as kernel CQEs. + injected: VecDeque, + // `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 { + 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) -> 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>, + injected: &mut VecDeque, + user_data: u64, + remaining: impl Iterator, + 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 { + 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::>()); + } +}