From 52d72d94dcb44cd60a9822603e43842e07c641fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Mon, 31 Aug 2026 12:35:49 -0400 Subject: [PATCH] [UIO] impl `IoUringFile::read_bytes_async` (#10288) * bridge async with a dedicated `tokio_uring` thread * impl `IoBufMut` for `AVec` * [AI] Add tests * [AI] handle O_DIRECT odirect test --- Cargo.lock | 53 ++++- lib/common/common/Cargo.toml | 1 + .../common/src/universal_io/io_uring/mod.rs | 8 +- .../common/src/universal_io/io_uring/tests.rs | 174 ++++++++++++++++- .../src/universal_io/io_uring/tokio_bridge.rs | 181 ++++++++++++++++++ 5 files changed, 404 insertions(+), 13 deletions(-) create mode 100644 lib/common/common/src/universal_io/io_uring/tokio_bridge.rs diff --git a/Cargo.lock b/Cargo.lock index 6368c8ede6..e5947f8c07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -182,7 +182,7 @@ dependencies = [ "futures-core", "futures-util", "mio", - "socket2", + "socket2 0.6.3", "tokio", "tracing", ] @@ -264,7 +264,7 @@ dependencies = [ "serde_json", "serde_urlencoded", "smallvec", - "socket2", + "socket2 0.6.3", "time", "tracing", "url", @@ -1370,7 +1370,7 @@ dependencies = [ "fs4 1.1.0", "fs_extra", "humantime", - "io-uring", + "io-uring 0.7.14", "itertools 0.15.0", "log", "memmap2", @@ -1400,6 +1400,7 @@ dependencies = [ "thiserror 2.0.20", "thread-priority", "tokio", + "tokio-uring", "validator", "walkdir", "zerocopy", @@ -2500,7 +2501,7 @@ dependencies = [ "futures-core", "futures-util", "hashbrown 0.16.1", - "io-uring", + "io-uring 0.7.14", "itertools 0.14.0", "libc", "lz4", @@ -3259,7 +3260,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.3", "tokio", "tower-service", "tracing", @@ -3560,6 +3561,16 @@ dependencies = [ "rustversion", ] +[[package]] +name = "io-uring" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "595a0399f411a508feb2ec1e970a4a30c249351e30208960d58298de8660b0e5" +dependencies = [ + "bitflags 1.3.2", + "libc", +] + [[package]] name = "io-uring" version = "0.7.14" @@ -5773,7 +5784,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2", + "socket2 0.6.3", "thiserror 2.0.20", "tokio", "tracing", @@ -5812,7 +5823,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.3", "tracing", "windows-sys 0.59.0", ] @@ -7151,6 +7162,16 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8e2fb0f499abb4d162f2bedad68f5ef91a1682b5a03596ddb67efd37768d100" +[[package]] +name = "socket2" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "socket2" version = "0.6.3" @@ -7701,7 +7722,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.3", "tokio-macros", "tracing", "windows-sys 0.61.2", @@ -7739,6 +7760,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-uring" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "748482e3e13584a34664a710168ad5068e8cb1d968aa4ffa887e83ca6dd27967" +dependencies = [ + "futures-util", + "io-uring 0.6.4", + "libc", + "slab", + "socket2 0.4.10", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.19" @@ -7821,7 +7856,7 @@ dependencies = [ "hyper-util", "percent-encoding", "pin-project", - "socket2", + "socket2 0.6.3", "sync_wrapper", "tokio", "tokio-rustls", diff --git a/lib/common/common/Cargo.toml b/lib/common/common/Cargo.toml index f8b21da5a7..a540411dda 100644 --- a/lib/common/common/Cargo.toml +++ b/lib/common/common/Cargo.toml @@ -68,6 +68,7 @@ tango-bench = "0.8.0" io-uring = "0.7.12" procfs = { version = "0.18", default-features = false } thread-priority = "3.0.0" +tokio-uring = "0.5.0" [[bench]] name = "bitpacking" diff --git a/lib/common/common/src/universal_io/io_uring/mod.rs b/lib/common/common/src/universal_io/io_uring/mod.rs index 79239c80cb..848ba785c5 100644 --- a/lib/common/common/src/universal_io/io_uring/mod.rs +++ b/lib/common/common/src/universal_io/io_uring/mod.rs @@ -2,6 +2,7 @@ mod error; mod pipeline; mod pool; mod runtime; +mod tokio_bridge; #[cfg(test)] mod tests; @@ -206,11 +207,12 @@ impl UniversalRead for IoUringFile { async fn read_bytes_async( &self, range: Range, - access_pattern: P, + _access_pattern: P, align: usize, ) -> UioResult> { - // TODO(uio): implement real async - self.read_bytes(range, access_pattern, align) + Ok(ACow::Owned( + tokio_bridge::read_bytes_async(self, range, align).await?, + )) } fn len(&self) -> UioResult { diff --git a/lib/common/common/src/universal_io/io_uring/tests.rs b/lib/common/common/src/universal_io/io_uring/tests.rs index 74085373be..7aa7b95476 100644 --- a/lib/common/common/src/universal_io/io_uring/tests.rs +++ b/lib/common/common/src/universal_io/io_uring/tests.rs @@ -1,3 +1,4 @@ +use std::assert_matches; use std::path::Path; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -7,7 +8,7 @@ use nix::libc; use super::super::*; use super::*; use crate::generic_consts::Sequential; -use crate::universal_io::UioResult; +use crate::universal_io::{UioResult, UniversalIoError}; /// Create `path`, populate it with the binary representation of `data`, /// then open and return it. @@ -365,3 +366,174 @@ fn test_io_uring_direct_io() -> UioResult<()> { Ok(()) } + +/// `read_bytes_async` must return the same bytes as the sync path on a +/// buffered (non-`O_DIRECT`) handle. +#[tokio::test] +async fn test_io_uring_read_bytes_async() -> UioResult<()> { + let dir = tempfile::tempdir().unwrap(); + + let data: Vec = (0..KERNEL_PAGE_SIZE * 2 + 1337) + .map(|idx| (idx % 256) as u8) + .collect(); + let file = test_file(&dir.path().join("async.bin"), &data, false)?; + let eof = data.len() as u64; + + // Full file + let bytes = file.read_bytes_async(0..eof, Sequential, 8).await?; + assert_eq!(bytes.as_ref(), &data); + + // Odd, unaligned sub-range + let bytes = file.read_bytes_async(13..1350, Sequential, 8).await?; + assert_eq!(bytes.as_ref(), &data[13..1350]); + + // Tail ending exactly at EOF + let bytes = file.read_bytes_async(eof - 100..eof, Sequential, 8).await?; + assert_eq!(bytes.as_ref(), &data[data.len() - 100..]); + + // Empty range + let bytes = file.read_bytes_async(10..10, Sequential, 8).await?; + assert!(bytes.as_ref().is_empty()); + + Ok(()) +} + +/// `read_bytes_async` on an `O_DIRECT` handle: same contract as +/// [`test_io_uring_direct_io`], including the EOF-truncated tail block. +#[tokio::test] +async fn test_io_uring_read_bytes_async_direct_io() -> UioResult<()> { + let dir = tempfile::tempdir().unwrap(); + + let data: Vec = (0..KERNEL_PAGE_SIZE * 2 + 1337) + .map(|idx| (idx % 256) as u8) + .collect(); + let file = test_file(&dir.path().join("o_direct_async.bin"), &data, true)?; + + for (idx, expected) in data.chunks(KERNEL_PAGE_SIZE).enumerate() { + let start = idx * KERNEL_PAGE_SIZE; + let range = start as u64..(start + expected.len()) as u64; + + let bytes = file + .read_bytes_async(range, Sequential, KERNEL_PAGE_SIZE) + .await?; + assert_eq!( + bytes.as_ref(), + expected, + "O_DIRECT async block {idx} mismatch" + ); + } + + Ok(()) +} + +/// `O_DIRECT` async reads that aren't whole pages: the bridge over-reads to +/// the page boundary and truncates back to the requested length. (The sync +/// pipeline doesn't support this shape — it asserts full-page reads.) +#[tokio::test] +async fn test_io_uring_read_bytes_async_direct_io_partial() -> UioResult<()> { + let dir = tempfile::tempdir().unwrap(); + + let data: Vec = (0..KERNEL_PAGE_SIZE * 2 + 1337) + .map(|idx| (idx % 256) as u8) + .collect(); + let file = test_file(&dir.path().join("o_direct_partial.bin"), &data, true)?; + let eof = data.len() as u64; + + // Sub-page length at a page-aligned offset, mid-file: over-read to the + // page boundary, truncated back to the requested 100 bytes. + let bytes = file + .read_bytes_async(0..100, Sequential, KERNEL_PAGE_SIZE) + .await?; + assert_eq!(bytes.as_ref(), &data[..100]); + + // Sub-page EOF-clamped tail: yields exactly the bytes up to EOF. + let tail_start = (KERNEL_PAGE_SIZE * 2) as u64; + let bytes = file + .read_bytes_async(tail_start..eof, Sequential, KERNEL_PAGE_SIZE) + .await?; + assert_eq!(bytes.as_ref(), &data[KERNEL_PAGE_SIZE * 2..]); + + // Page-aligned range crossing EOF: the short read must error, not + // silently return fewer bytes than requested. + let crossing = file + .read_bytes_async( + tail_start..tail_start + 2 * KERNEL_PAGE_SIZE as u64, + Sequential, + KERNEL_PAGE_SIZE, + ) + .await; + assert!(crossing.is_err(), "O_DIRECT read crossing EOF must error"); + + Ok(()) +} + +/// Invalid ranges must error rather than resolve: past EOF, crossing EOF, +/// and inverted (`start > end`) ranges. +#[tokio::test] +async fn test_io_uring_read_bytes_async_invalid_ranges() -> UioResult<()> { + let dir = tempfile::tempdir().unwrap(); + + let data: Vec = (0..1024).map(|idx| (idx % 256) as u8).collect(); + let file = test_file(&dir.path().join("invalid.bin"), &data, false)?; + let eof = data.len() as u64; + + let past_eof = file + .read_bytes_async(eof + 10..eof + 20, Sequential, 8) + .await; + assert!(past_eof.is_err(), "read past EOF must error"); + + let crossing = file + .read_bytes_async(eof - 10..eof + 10, Sequential, 8) + .await; + assert!(crossing.is_err(), "read crossing EOF must error"); + + // Inverted range: an error, not a silent empty read. + #[expect(clippy::reversed_empty_ranges)] + let inverted = file.read_bytes_async(300..100, Sequential, 8).await; + assert_matches!(inverted, Err(UniversalIoError::OutOfBounds { .. })); + + Ok(()) +} + +/// Concurrent `read_bytes_async` calls from spawned tasks across two files: +/// each must resolve with its own file's bytes. Spawning also pins that the +/// returned future is `Send`. +#[tokio::test(flavor = "multi_thread")] +async fn test_io_uring_read_bytes_async_concurrent() -> UioResult<()> { + const NUM_TASKS: usize = 32; + const CHUNK: usize = 1000; + + let dir = tempfile::tempdir().unwrap(); + + let data_a: Vec = (0..NUM_TASKS * CHUNK) + .map(|idx| (idx % 251) as u8) + .collect(); + let data_b: Vec = (0..NUM_TASKS * CHUNK) + .map(|idx| (idx % 241) as u8) + .collect(); + + let file_a = test_file(&dir.path().join("conc_a.bin"), &data_a, false)?; + let file_b = test_file(&dir.path().join("conc_b.bin"), &data_b, false)?; + + let mut tasks = Vec::new(); + + for idx in 0..NUM_TASKS { + for (file, data) in [(&file_a, &data_a), (&file_b, &data_b)] { + let file = file.clone(); + let expected = data[idx * CHUNK..(idx + 1) * CHUNK].to_vec(); + let range = (idx * CHUNK) as u64..((idx + 1) * CHUNK) as u64; + + tasks.push(tokio::spawn(async move { + let bytes = file.read_bytes_async(range, Sequential, 8).await?; + assert_eq!(bytes.as_ref(), &expected, "concurrent read {idx} mismatch"); + UioResult::Ok(()) + })); + } + } + + for task in tasks { + task.await.expect("task join")?; + } + + Ok(()) +} diff --git a/lib/common/common/src/universal_io/io_uring/tokio_bridge.rs b/lib/common/common/src/universal_io/io_uring/tokio_bridge.rs new file mode 100644 index 0000000000..eb42630467 --- /dev/null +++ b/lib/common/common/src/universal_io/io_uring/tokio_bridge.rs @@ -0,0 +1,181 @@ +//! Alternative Tokio-based IO implementation using io-uring. + +use std::io; +use std::ops::Range; +use std::sync::LazyLock; + +use aligned_vec::{AVec, RuntimeAlign}; +use tokio::sync::mpsc; + +use super::KERNEL_PAGE_SIZE; +use crate::universal_io::{IoUringFile, UioResult, UniversalIoError}; +static URING_BRIDGE: LazyLock = LazyLock::new(UringBridge::spawn); + +struct UringRequest { + file: fs_err::File, + offset: u64, + len: usize, + align: usize, + direct_io: bool, + reply: mpsc::Sender>>, +} + +struct UringBridge { + tx: tokio::sync::mpsc::UnboundedSender, +} + +impl UringBridge { + fn spawn() -> Self { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); + std::thread::Builder::new() + .name("tokio-uring-bridge".into()) + .spawn(move || { + tokio_uring::start(async move { + while let Some(req) = rx.recv().await { + tokio_uring::spawn(async move { + let _ = Self::execute(req).await; + }); + } + }); + }) + .expect("spawn uring bridge thread"); + Self { tx } + } + + async fn execute(req: UringRequest) { + let UringRequest { + file, + offset, + len, + align, + direct_io, + reply, + } = req; + + let file = tokio_uring::fs::File::from_std(file.into_file()); + + let read = || async { + if !direct_io { + let buf = ABuf(AVec::with_capacity(align, len)); + + let (res, buf) = file.read_exact_at(buf, offset).await; + + res?; + + return Ok(buf.0); + } + + // Mirror `IoUringState::read`: `O_DIRECT` requires a page-aligned + // offset and buffer, and a page-multiple submitted length. + assert!( + align.is_multiple_of(KERNEL_PAGE_SIZE), + "O_DIRECT read buffer must be aligned to {KERNEL_PAGE_SIZE} bytes (alignment: {align})", + ); + assert!( + offset.is_multiple_of(KERNEL_PAGE_SIZE as u64), + "O_DIRECT read offset must be aligned to {KERNEL_PAGE_SIZE} bytes (offset: {offset})", + ); + + let kernel_len = len + .checked_next_multiple_of(KERNEL_PAGE_SIZE) + .expect("rounded read length fit within usize"); + + let buf = ABuf(AVec::with_capacity(align, kernel_len)); + let (res, buf) = file.read_at(buf, offset).await; + let bytes_read = res?; + + // A short count only happens when the range crosses EOF: the + // EOF-clamped tail block still yields all `len` requested bytes. + if bytes_read < len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("O_DIRECT read at {offset} returned {bytes_read} of {len} bytes"), + ) + .into()); + } + + // Drop the over-read up to the page boundary. + let mut buffer = buf.0; + buffer.truncate(len); + Ok(buffer) + }; + + match reply.send(read().await).await { + Ok(()) => {} + Err(_err) => { + // The requester was dropped or closed the channel, nothing to do here + } + } + } +} + +pub async fn read_bytes_async( + file: &IoUringFile, + range: Range, + align: usize, +) -> UioResult> { + let (tx, mut rx) = mpsc::channel(1); + + let direct_io = file.direct_io; + let file = file.file.try_clone()?; + + let Some(len) = range.end.checked_sub(range.start) else { + return Err(UniversalIoError::OutOfBounds { + start: range.start, + end: range.end, + elements: 0, + }); + }; + + let req = UringRequest { + file, + offset: range.start, + len: len as usize, + align, + direct_io, + reply: tx, + }; + + URING_BRIDGE + .tx + .send(req) + .map_err(|_err| UniversalIoError::Uninitialized { + description: "uring bridge receiver has been closed".to_owned(), + })?; + + let Some(result) = rx.recv().await else { + return Err(UniversalIoError::Uninitialized { + description: "uring request been dropped".to_owned(), + }); + }; + + result +} + +struct ABuf(AVec); + +unsafe impl tokio_uring::buf::IoBuf for ABuf { + fn stable_ptr(&self) -> *const u8 { + self.0.as_ptr() + } + + fn bytes_init(&self) -> usize { + self.0.len() + } + + fn bytes_total(&self) -> usize { + self.0.capacity() + } +} + +unsafe impl tokio_uring::buf::IoBufMut for ABuf { + fn stable_mut_ptr(&mut self) -> *mut u8 { + self.0.as_mut_ptr() + } + + unsafe fn set_init(&mut self, init_len: usize) { + if self.0.len() < init_len { + unsafe { self.0.set_len(init_len) }; + } + } +}