From 3bfa59278fab240c459b2b62dd72a691318d03cb Mon Sep 17 00:00:00 2001 From: Lexo Liu Date: Wed, 11 Nov 2026 08:27:33 -0500 Subject: [PATCH] feat(io)!: name stream sockets by registered file slots on io_uring Each compio stream on Linux takes a slot in its ring's file table, and every operation on it (multishot and single receives, copying and zero-copy sends, readiness polls, shutdown) addresses the socket by that slot with IOSQE_FIXED_FILE, so the kernel no longer looks the descriptor up and references its file per operation. The descriptor stays for kernel TLS options, addresses and splice. The slot is allocated with IORING_OP_FILES_UPDATE on a runtime task, so no operation waits for it and none abandoned midway loses it; the first stream on a ring registers the table, sized by the new FileTable setting (runtime-adjustable, default 1024). A full table or a kernel without slot allocation leaves the stream on its descriptor; a table the kernel refuses is the stream's next error. compio builds its operations on descriptors, so an adapter rewrites the entries they build and keeps the rest of each operation, buffer-ring handling included, compio's. Compio::with_zero_copy becomes a builder method beside Compio::with_file_table. Refs #125 --- io/Cargo.toml | 16 +- io/src/compio.rs | 762 ++++++++++++++++++++++++++++++++++++++++++----- io/tests/io.rs | 253 +++++++++++++++- 3 files changed, 949 insertions(+), 82 deletions(-) diff --git a/io/Cargo.toml b/io/Cargo.toml index 69aa0e3..60809b5 100644 --- a/io/Cargo.toml +++ b/io/Cargo.toml @@ -14,7 +14,7 @@ categories = ["asynchronous", "network-programming"] [features] ## compio: io_uring on Linux (falling back to polling where io_uring is unavailable), IOCP on Windows, polling elsewhere. -compio = ["dep:compio", "dep:libc", "dep:windows-sys", "dep:socket2", "dep:futures-core"] +compio = ["dep:compio", "dep:libc", "dep:windows-sys", "dep:socket2", "dep:futures-core", "dep:io-uring", "dep:linux-raw-sys"] ## tokio's reactor, adapted to owned buffers. tokio = ["dep:tokio", "dep:socket2", "dep:libc", "dep:windows-sys"] ## async-io's reactor, adapted to owned buffers. @@ -36,6 +36,20 @@ futures-lite = { version = "2.6.1", default-features = false } libc = { version = "0.2.191", default-features = false } futures-core = { version = "0.3.34", default-features = false, optional = true } +# The submission entries of compio's io_uring driver, which addresses a +# socket by its slot in the ring's file table: the same crate compio builds +# them with, so that an entry built here is one its driver takes. +[target.'cfg(target_os = "linux")'.dependencies.io-uring] +version = "0.7.16" +default-features = false +optional = true + +[target.'cfg(target_os = "linux")'.dependencies.linux-raw-sys] +version = "0.12.1" +default-features = false +features = ["io_uring"] +optional = true + [target.'cfg(unix)'.dependencies] libc = { version = "0.2.191", default-features = false, optional = true } diff --git a/io/src/compio.rs b/io/src/compio.rs index 23f9c56..82f65ad 100644 --- a/io/src/compio.rs +++ b/io/src/compio.rs @@ -42,6 +42,20 @@ //! Smaller writes, the other write methods, Unix-domain streams (which the //! kernel does not send zero-copy) and every other platform copy. //! +//! On Linux over io_uring each stream's socket takes a slot in its ring's +//! file table, and every operation on the stream (receives, multishot ones +//! included, sends, zero-copy ones included, polls, shutdown) names the +//! socket by that slot (`IOSQE_FIXED_FILE`): the kernel then neither looks +//! the descriptor up nor takes a reference to its file per operation. The +//! stream keeps its descriptor for what is not an io_uring operation (kernel +//! TLS options, addresses, `splice(2)`). The slot is taken on the side as the +//! stream is created, so the first operations may still use the descriptor, +//! and is given back once the stream is dropped. A ring registers its table, +//! [`FileTable::slots`] slots large, when its first stream asks for a slot; +//! a socket the table has no room for, or a kernel that cannot allocate +//! slots (before 5.19), leaves the stream on its descriptor. The ring's +//! table belongs to these streams: nothing else on the runtime may use it. +//! //! The poll-style capability is a [`CopyingPoll`]: a completion-based driver //! cannot read into a buffer borrowed for one poll. @@ -56,10 +70,9 @@ use std::{ time::Instant, }; -use ::compio::{ - io::{AsyncRead, AsyncWrite}, - net, -}; +#[cfg(not(target_os = "linux"))] +use ::compio::io::{AsyncRead, AsyncWrite}; +use ::compio::net; use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf}; use crate::{ @@ -73,18 +86,28 @@ use crate::{ use crate::{DialUnix, FromStdUnix, ListenUnix}; /// The compio implementation. compio keeps its runtime per thread, and every -/// operation goes to the current thread's; this holds only the -/// [`ZeroCopySend`] setting, which every stream it creates or adopts shares. +/// operation goes to the current thread's; this holds only the operator's +/// settings, [`ZeroCopySend`] and [`FileTable`], which every stream it +/// creates or adopts shares. #[derive(Debug, Clone, Default)] pub struct Compio { zero_copy: ZeroCopySend, + file_table: FileTable, } impl Compio { - /// An implementation whose streams send zero-copy as `zero_copy` says. + /// This implementation with its streams sending zero-copy as `zero_copy` + /// says. + #[must_use] + pub fn with_zero_copy(self, zero_copy: ZeroCopySend) -> Self { + Self { zero_copy, ..self } + } + + /// This implementation with its rings' file tables sized as `file_table` + /// says. #[must_use] - pub const fn with_zero_copy(zero_copy: ZeroCopySend) -> Self { - Self { zero_copy } + pub fn with_file_table(self, file_table: FileTable) -> Self { + Self { file_table, ..self } } /// The zero-copy setting its streams share. @@ -93,6 +116,12 @@ impl Compio { &self.zero_copy } + /// The file-table setting its streams share. + #[must_use] + pub const fn file_table(&self) -> &FileTable { + &self.file_table + } + /// Adopts a connected compio stream: one a runtime of this thread /// registered, such as a socket accepted elsewhere and handed over as a /// std stream (on Windows a socket joins one completion port only, so a @@ -100,7 +129,58 @@ impl Compio { /// thread adopts its share). #[must_use] pub fn tcp_stream(&self, stream: net::TcpStream) -> TcpStream { - TcpStream::new(stream, self.zero_copy.clone()) + TcpStream::new(stream, self) + } +} + +/// How many sockets a ring's file table holds, on Linux over io_uring: each +/// stream takes a slot while it lives, and a stream the table has no room +/// for uses its descriptor. +/// +/// It is the operator's runtime-adjustable handle: every clone shares one +/// value, across threads. A ring sizes its table once, when its first stream +/// asks for a slot (the kernel cannot resize a table holding sockets), so a +/// change sizes the tables of rings that register theirs afterwards. `0` +/// stops streams created afterwards from taking slots at all, on every ring. +/// +/// The kernel refuses a table larger than the process's `RLIMIT_NOFILE` +/// (`io_sqe_files_register`); the streams of a ring whose table was refused +/// report that error. +#[derive(Debug, Clone)] +pub struct FileTable { + slots: std::sync::Arc, +} + +impl FileTable { + /// The default size: the kernel's default soft `RLIMIT_NOFILE` + /// (`INR_OPEN_CUR`, `include/uapi/linux/fs.h`), the largest table a + /// process that has not raised its limit may register. + pub const DEFAULT_SLOTS: u32 = 1024; + + /// A setting of `slots` slots. + #[must_use] + pub fn new(slots: u32) -> Self { + Self { + slots: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(slots)), + } + } + + /// The slots a ring registers its table with. + #[must_use] + pub fn slots(&self) -> u32 { + self.slots.load(std::sync::atomic::Ordering::Relaxed) + } + + /// Sets the size for every clone of this setting. + pub fn set_slots(&self, slots: u32) { + self.slots + .store(slots, std::sync::atomic::Ordering::Relaxed); + } +} + +impl Default for FileTable { + fn default() -> Self { + Self::new(Self::DEFAULT_SLOTS) } } @@ -177,7 +257,7 @@ impl DialUnix for Compio { async fn connect_unix(&self, path: &std::path::Path) -> io::Result { let stream = net::UnixStream::connect(path).await?; - Ok(UnixStream::new(stream, self.zero_copy.clone())) + Ok(UnixStream::new(stream, self)) } } @@ -213,14 +293,14 @@ impl FromStdUnix for Compio { listener: std::os::unix::net::UnixListener, ) -> io::Result { let listener = net::UnixListener::from_std(listener)?; - Ok(UnixListener::new(listener, self.zero_copy.clone())) + Ok(UnixListener::new(listener, self.clone())) } } impl FromStd for Compio { fn tcp_listener(&self, listener: std::net::TcpListener) -> io::Result { let listener = net::TcpListener::from_std(listener)?; - Ok(TcpListener::new(listener, self.zero_copy.clone())) + Ok(TcpListener::new(listener, self.clone())) } fn udp_socket(&self, socket: std::net::UdpSocket) -> io::Result { @@ -253,7 +333,9 @@ trait Fd {} #[cfg(not(unix))] impl Fd for T {} -/// A stream socket of compio's, which the operations below drive through a +/// A stream socket of compio's. On Linux the stream submits compio's +/// operations on it itself, so that they name the socket by its slot in the +/// ring's file table; elsewhere compio's stream drives them, through a /// shared reference. trait Socket: Clone + Fd + 'static { /// Whether the kernel sends from this socket zero-copy: it sets @@ -264,18 +346,15 @@ trait Socket: Clone + Fd + 'static { /// The descriptor, shared with an operation that outlives a call. #[cfg(target_os = "linux")] fn shared_fd(&self) -> ::compio::driver::SharedFd; + #[cfg(not(target_os = "linux"))] fn read(&self, buf: B) -> impl Future>; + #[cfg(not(target_os = "linux"))] fn write(&self, buf: B) -> impl Future>; + #[cfg(not(target_os = "linux"))] fn write_vectored(&self, buf: B) -> impl Future>; + #[cfg(not(target_os = "linux"))] fn shutdown(&self) -> impl Future>; - /// `sendmsg(2)` of `buf` with the control messages in `control`. - #[cfg(unix)] - fn write_control( - &self, - buf: B, - control: C, - ) -> impl Future>; /// `sendfile(2)` of a region once the socket is writable. #[cfg(target_vendor = "apple")] fn sendfile( @@ -297,41 +376,30 @@ macro_rules! impl_socket { ::compio::driver::ToSharedFd::to_shared_fd(self) } + #[cfg(not(target_os = "linux"))] async fn read(&self, buf: B) -> BufResult { let mut stream = self; AsyncRead::read(&mut stream, buf).await } + #[cfg(not(target_os = "linux"))] async fn write(&self, buf: B) -> BufResult { let mut stream = self; AsyncWrite::write(&mut stream, buf).await } + #[cfg(not(target_os = "linux"))] async fn write_vectored(&self, buf: B) -> BufResult { let mut stream = self; AsyncWrite::write_vectored(&mut stream, buf).await } + #[cfg(not(target_os = "linux"))] async fn shutdown(&self) -> io::Result<()> { let mut stream = self; AsyncWrite::shutdown(&mut stream).await } - #[cfg(unix)] - async fn write_control( - &self, - buf: B, - control: C, - ) -> BufResult { - use ::compio::driver::op::{SendFlags, SendMsg}; - use compio_buf::IntoInner as _; - - let op = SendMsg::new(self.clone(), [buf], control, None, SendFlags::empty()); - let BufResult(res, op) = ::compio::runtime::submit(op).await; - let ([buf], control) = op.into_inner(); - BufResult(res, (buf, control)) - } - #[cfg(target_vendor = "apple")] async fn sendfile(&self, file: &File, offset: u64, len: usize) -> io::Result { use std::os::fd::AsRawFd; @@ -377,14 +445,29 @@ pub struct Stream { /// From what size a write sends zero-copy. #[cfg(target_os = "linux")] zero_copy: ZeroCopySend, + /// The socket's slot in the ring's file table, which every operation + /// names it by once it has one. + #[cfg(target_os = "linux")] + file: files::File, + /// The last receive of its own found the socket empty: the next one waits + /// for data before trying (`IORING_RECVSEND_POLL_FIRST`), as compio's + /// stream does. + #[cfg(target_os = "linux")] + recv_empty: std::cell::Cell, } -impl Stream { - fn new(inner: S, zero_copy: ZeroCopySend) -> Self { - // Only Linux sends zero-copy. +#[expect( + private_bounds, + reason = "every method is private: only this module creates and drives a stream" +)] +impl Stream { + fn new(inner: S, io: &Compio) -> Self { + // Only Linux sends zero-copy and registers files. #[cfg(not(target_os = "linux"))] - drop(zero_copy); + let _ = io; Self { + #[cfg(target_os = "linux")] + file: files::File::new(inner.shared_fd(), io.file_table()), inner, #[cfg(any(target_os = "linux", target_os = "android"))] pipe: std::cell::Cell::new(None), @@ -393,10 +476,28 @@ impl Stream { #[cfg(target_os = "linux")] recv: recv::Recv::default(), #[cfg(target_os = "linux")] - zero_copy, + zero_copy: io.zero_copy.clone(), + #[cfg(target_os = "linux")] + recv_empty: std::cell::Cell::new(false), + } + } + + /// Submits `op`, which operates on this socket: on Linux by the socket's + /// slot in the file table when it has one. + #[cfg(unix)] + async fn submit(&self, op: T) -> BufResult { + #[cfg(target_os = "linux")] + { + self.file.submit(op).await + } + #[cfg(not(target_os = "linux"))] + { + ::compio::runtime::submit(op).await } } +} +impl Stream { /// compio's stream. #[must_use] pub fn get_ref(&self) -> &S { @@ -404,6 +505,71 @@ impl Stream { } } +/// The socket operations, which compio's stream would submit with the +/// descriptor. +#[cfg(target_os = "linux")] +#[expect( + private_bounds, + reason = "every method is private: only this module drives a stream" +)] +impl Stream { + async fn recv(&self, mut buf: B) -> BufResult { + use ::compio::driver::{ + PollFirst as _, + op::{Recv, RecvFlags}, + }; + use compio_buf::IntoInner as _; + + // compio writes from the start but only ever grows the length; the + // contract is a length of exactly what was read. + buf.clear(); + let mut op = Recv::new(self.inner.shared_fd(), buf, RecvFlags::empty()); + if self.recv_empty.get() { + op.poll_first(); + } + let (BufResult(res, op), extra) = self.file.submit_with_extra(op).await; + if let Some(Ok(nonempty)) = extra.map(|extra| extra.sock_nonempty()) { + self.recv_empty.set(!nonempty); + } + let mut buf = op.into_inner(); + if let Ok(n) = res { + // SAFETY: the kernel received `n` bytes into the buffer's start. + unsafe { buf.advance_to(n) }; + } + BufResult(res, buf) + } + + async fn send(&self, buf: B) -> BufResult { + use ::compio::driver::op::{Send, SendFlags}; + use compio_buf::IntoInner as _; + + let op = Send::new(self.inner.shared_fd(), buf, SendFlags::empty()); + let BufResult(res, op) = self.file.submit(op).await; + BufResult(res, op.into_inner()) + } + + async fn send_vectored(&self, buf: B) -> BufResult { + use ::compio::driver::op::{SendFlags, SendVectored}; + use compio_buf::IntoInner as _; + + let op = SendVectored::new(self.inner.shared_fd(), buf, SendFlags::empty()); + let BufResult(res, op) = self.file.submit(op).await; + BufResult(res, op.into_inner()) + } + + async fn shutdown_write(&self) -> io::Result<()> { + use ::compio::driver::op::ShutdownSocket; + + let op = ShutdownSocket::new(self.inner.shared_fd(), std::net::Shutdown::Write); + match self.file.submit(op).await.0 { + // A socket that is not connected has nothing to shut down, as + // compio's stream has it. + Err(e) if e.kind() == io::ErrorKind::NotConnected => Ok(()), + res => res.map(drop), + } + } +} + #[cfg(unix)] impl Stream { /// The local address. @@ -478,17 +644,23 @@ impl SharedStream for Stream { async fn read(&self, buf: B) -> BufResult { #[cfg(target_os = "linux")] - let buf = match self.recv.read_ahead(&self.inner, buf).await { - recv::Copied::Done(result) => return result, - recv::Copied::ReadInstead(buf) => buf, - }; - read_single(&self.inner, buf).await + { + let buf = match self.recv.read_ahead(&self.inner, &self.file, buf).await { + recv::Copied::Done(result) => return result, + recv::Copied::ReadInstead(buf) => buf, + }; + self.recv(buf).await + } + #[cfg(not(target_os = "linux"))] + { + read_single(&self.inner, buf).await + } } async fn read_owned(&self) -> io::Result> { #[cfg(target_os = "linux")] { - self.recv.read_owned(&self.inner).await + self.recv.read_owned(&self.inner, &self.file).await } #[cfg(not(target_os = "linux"))] { @@ -508,17 +680,31 @@ impl SharedStream for Stream { } async fn write(&self, buf: B) -> BufResult { - self.inner.write(buf).await + #[cfg(target_os = "linux")] + { + self.send(buf).await + } + #[cfg(not(target_os = "linux"))] + { + self.inner.write(buf).await + } } async fn write_vectored(&self, buf: B) -> BufResult { - self.inner.write_vectored(buf).await + #[cfg(target_os = "linux")] + { + self.send_vectored(buf).await + } + #[cfg(not(target_os = "linux"))] + { + self.inner.write_vectored(buf).await + } } #[cfg(target_os = "linux")] async fn write_all_owned(&self, buf: B) -> BufResult<(), Option> { if S::ZERO_COPY && buf.total_len() >= self.zero_copy.threshold() { - return zero_copy::write_all(&self.inner, buf).await; + return zero_copy::write_all(self, buf).await; } let BufResult(res, buf) = crate::stream::write_all_vectored(&mut crate::stream::ByRef(self), buf).await; @@ -526,16 +712,30 @@ impl SharedStream for Stream { } async fn shutdown(&self) -> io::Result<()> { - self.inner.shutdown().await + #[cfg(target_os = "linux")] + { + self.shutdown_write().await + } + #[cfg(not(target_os = "linux"))] + { + self.inner.shutdown().await + } } + /// `sendmsg(2)` of `buf` with the control messages in `control`. #[cfg(unix)] async fn write_control( &self, buf: B, control: C, ) -> BufResult { - self.inner.write_control(buf, control).await + use ::compio::driver::op::{SendFlags, SendMsg}; + use compio_buf::IntoInner as _; + + let op = SendMsg::new(self.inner.clone(), [buf], control, None, SendFlags::empty()); + let BufResult(res, op) = self.submit(op).await; + let ([buf], control) = op.into_inner(); + BufResult(res, (buf, control)) } #[cfg(any(target_os = "linux", target_os = "android"))] @@ -581,6 +781,7 @@ impl SharedStream for Stream { } /// One receive into the caller's buffer, with a submission of its own. +#[cfg(not(target_os = "linux"))] async fn read_single(socket: &S, mut buf: B) -> BufResult { // compio writes from the start but only ever grows the length; the // contract is a length of exactly what was read. @@ -721,8 +922,7 @@ mod splice { } Ok(n) => drained += n, Err(e) if e.kind() == io::ErrorKind::WouldBlock => { - let writable = PollOnce::new(stream.inner.clone(), Interest::Writable); - if let Err(e) = ::compio::runtime::submit(writable).await.0 { + if let Err(e) = ready(stream, Interest::Writable).await { return BufResult(Err(e), file); } } @@ -745,7 +945,8 @@ mod splice { /// Waits for the socket to be ready for `interest` with a poll the /// driver arms, which holds no thread. async fn ready(stream: &Stream, interest: Interest) -> io::Result<()> { - ::compio::runtime::submit(PollOnce::new(stream.inner.clone(), interest)) + stream + .submit(PollOnce::new(stream.inner.clone(), interest)) .await .0 .map(drop) @@ -818,6 +1019,405 @@ mod splice { } } +/// Registered files on io_uring: each stream's slot in its ring's file +/// table, and compio's operations naming the socket by it. +/// +/// compio registers a table (`register_files`, where `-1` entries are empty +/// slots) but builds every operation on a descriptor, so [`Addressed`] wraps +/// its operation and rewrites the entry it builds: everything else, the +/// multishot and buffer-ring handling included, stays compio's. +#[cfg(target_os = "linux")] +mod files { + use std::{ + cell::Cell, + io, + os::fd::{AsRawFd as _, RawFd}, + rc::Rc, + task::Poll, + }; + + use ::compio::{ + driver::{ + Decision, Extra, IourOpCode, OpCode, OpEntry, OpType, PollOpCode, SharedFd, TakeBuffer, + }, + runtime::Runtime, + }; + use compio_buf::BufResult; + use linux_raw_sys::io_uring::{IORING_FILE_INDEX_ALLOC, io_uring_sqe}; + + // An entry is the kernel's submission entry, which the ring copies it + // into as it is. + const _: () = assert!(size_of::() == size_of::()); + + /// compio's operation `op` on a stream's socket, naming the socket by its + /// slot when it has one. + pub(super) struct Addressed { + op: T, + /// The slot, and the descriptor `op` names, which the slot replaces. + slot: Option<(u32, RawFd)>, + } + + /// `entry`, naming its socket by `slot` (`IOSQE_FIXED_FILE`) when there + /// is one. + fn fixed(entry: OpEntry, slot: Option<(u32, RawFd)>) -> OpEntry { + match (entry, slot) { + (OpEntry::Submission(entry), Some((index, fd))) => { + let mut entry = entry.flags(io_uring::squeue::Flags::FIXED_FILE); + let sqe = (&raw mut entry).cast::(); + // SAFETY: an entry is a kernel submission entry (asserted + // above); compio sets its fields the same way + // (`set_poll_first`). + unsafe { + debug_assert_eq!((*sqe).fd, fd, "the operation names another descriptor"); + (*sqe).fd = index.cast_signed(); + } + OpEntry::Submission(entry) + } + (entry, _) => entry, + } + } + + // SAFETY: every method is `op`'s own; the entries differ only in naming + // the same socket by its slot, which holds the socket while the stream + // lives. + unsafe impl IourOpCode for Addressed { + type Control = T::Control; + + unsafe fn init(&mut self, control: &mut Self::Control) { + // SAFETY: the caller's guarantee covers `op`, which `self` holds. + unsafe { self.op.init(control) } + } + + fn create_entry(&mut self, control: &mut Self::Control) -> OpEntry { + fixed(self.op.create_entry(control), self.slot) + } + + fn create_entry_fallback(&mut self, control: &mut Self::Control) -> OpEntry { + fixed(self.op.create_entry_fallback(control), self.slot) + } + + fn call_blocking(&mut self, control: &mut Self::Control) -> io::Result { + self.op.call_blocking(control) + } + + unsafe fn set_result( + &mut self, + control: &mut Self::Control, + res: &io::Result, + extra: &Extra, + ) { + // SAFETY: the result is `op`'s, as the caller guarantees. + unsafe { IourOpCode::set_result(&mut self.op, control, res, extra) } + } + + unsafe fn push_multishot( + &mut self, + control: &mut Self::Control, + res: io::Result, + extra: Extra, + ) { + // SAFETY: the result is `op`'s, as the caller guarantees. + unsafe { self.op.push_multishot(control, res, extra) } + } + + fn pop_multishot( + &mut self, + control: &mut Self::Control, + ) -> Option> { + self.op.pop_multishot(control) + } + } + + // SAFETY: every method is `op`'s own; the polling driver uses the + // descriptor, which `op` holds. + unsafe impl PollOpCode for Addressed { + type Control = T::Control; + + unsafe fn init(&mut self, control: &mut Self::Control) { + // SAFETY: the caller's guarantee covers `op`, which `self` holds. + unsafe { self.op.init(control) } + } + + fn pre_submit(&mut self, control: &mut Self::Control) -> io::Result { + self.op.pre_submit(control) + } + + fn op_type(&mut self, control: &mut Self::Control) -> Option { + self.op.op_type(control) + } + + fn operate(&mut self, control: &mut Self::Control) -> Poll> { + self.op.operate(control) + } + + unsafe fn set_result( + &mut self, + control: &mut Self::Control, + res: &io::Result, + extra: &Extra, + ) { + // SAFETY: the result is `op`'s, as the caller guarantees. + unsafe { PollOpCode::set_result(&mut self.op, control, res, extra) } + } + } + + impl TakeBuffer for Addressed { + type Buffer = T::Buffer; + + fn take_buffer(self) -> Option { + self.op.take_buffer() + } + } + + /// A stream's hold on its slot. + pub(super) struct File { + state: Rc>, + /// The socket's descriptor, which the slot replaces in an entry. + fd: RawFd, + } + + #[derive(Default)] + enum State { + /// The installer has not finished; operations name the descriptor. + Installing, + /// The socket holds this slot. + Fixed(u32), + /// No slot; operations name the descriptor. + #[default] + Descriptor, + /// Taking a slot failed other than for want of room: the next + /// operation reports it, and the rest name the descriptor. + Failed(io::Error), + /// The stream is gone: the installer gives back the slot it takes. + Dropped, + } + + impl File { + /// Starts taking a slot for `socket`, on a task of the runtime, so + /// that no operation of the stream waits for it and none abandoned + /// midway loses the slot (and with it the socket, which the slot + /// keeps open). + pub(super) fn new(socket: SharedFd, table: &super::FileTable) -> Self { + let fd = socket.as_raw_fd(); + let state = Rc::new(Cell::new(State::Descriptor)); + let slots = table.slots(); + if slots > 0 { + Runtime::with_current(|runtime| { + if runtime.driver_type().is_iouring() { + state.set(State::Installing); + runtime + .spawn(install(Rc::clone(&state), socket, slots)) + .detach(); + } + }); + } + Self { state, fd } + } + + fn slot(&self) -> io::Result> { + match self.state.take() { + State::Fixed(index) => { + self.state.set(State::Fixed(index)); + Ok(Some(index)) + } + State::Failed(error) => Err(error), + state => { + self.state.set(state); + Ok(None) + } + } + } + + /// `op`, naming the socket by its slot when it has one; or, with + /// `op` back, the error taking a slot failed with. + pub(super) fn address(&self, op: T) -> Result, (io::Error, T)> { + match self.slot() { + Ok(slot) => Ok(Addressed { + op, + slot: slot.map(|index| (index, self.fd)), + }), + Err(error) => Err((error, op)), + } + } + + pub(super) async fn submit(&self, op: T) -> BufResult { + match self.address(op) { + Ok(op) => { + let BufResult(res, op) = ::compio::runtime::submit(op).await; + BufResult(res, op.op) + } + Err((error, op)) => BufResult(Err(error), op), + } + } + + /// [`submit`](Self::submit), with what the driver reports beside the + /// result. + pub(super) async fn submit_with_extra( + &self, + op: T, + ) -> (BufResult, Option) { + match self.address(op) { + Ok(op) => { + let (BufResult(res, op), extra) = + ::compio::runtime::submit(op).with_extra().await; + (BufResult(res, op.op), Some(extra)) + } + Err((error, op)) => (BufResult(Err(error), op), None), + } + } + } + + impl Drop for File { + fn drop(&mut self) { + if let State::Fixed(index) = self.state.replace(State::Dropped) { + // Without a runtime the ring is gone, and its table with it. + let _ = Runtime::try_with_current(|runtime| runtime.spawn(remove(index)).detach()); + } + } + } + + /// Takes a slot for `socket` and leaves the outcome in `state`, or gives + /// the slot back at once when the stream is gone by then. + async fn install(state: Rc>, socket: SharedFd, slots: u32) { + let taken = take_slot(socket, slots).await; + if let State::Dropped = state.take() { + if let Ok(Some(index)) = taken { + remove(index).await; + } + return; + } + state.set(match taken { + Ok(Some(index)) => State::Fixed(index), + Ok(None) => State::Descriptor, + Err(error) => State::Failed(error), + }); + } + + /// The slot the kernel allocates for `socket`, registering the ring's + /// table first if it has none; `None` where there is no room or no + /// allocation. + async fn take_slot(socket: SharedFd, slots: u32) -> io::Result> { + let mut registered = false; + loop { + let op = FilesUpdate::install(socket.clone()); + let BufResult(res, op) = ::compio::runtime::submit(op).await; + let Err(error) = res else { + return Ok(Some(op.slot())); + }; + match error.raw_os_error() { + // The table is full (`io_fixed_fd_install`). + Some(libc::ENFILE) => return Ok(None), + // A kernel that does not allocate slots (before 5.19) reads + // the allocation marker as an offset, which overflows + // (`__io_register_rsrc_update`). + Some(libc::EOVERFLOW) => return Ok(None), + // The ring has no table yet. + Some(libc::ENXIO) if !registered => { + register(slots)?; + registered = true; + } + // A kernel without the operation (before 5.6). + None if error.kind() == io::ErrorKind::Unsupported => return Ok(None), + _ => return Err(error), + } + } + } + + /// Registers the current ring's table, of `slots` empty slots. + fn register(slots: u32) -> io::Result<()> { + let table = vec![-1; slots as usize]; + match Runtime::with_current(|runtime| runtime.register_files(&table)) { + // Another stream's installer registered it meanwhile. + Err(error) if error.raw_os_error() == Some(libc::EBUSY) => Ok(()), + res => res.map_err(|error| { + io::Error::new( + error.kind(), + format!("registering a file table of {slots} slots: {error}"), + ) + }), + } + } + + /// Gives `index` back to the table, which drops its hold on the socket. + async fn remove(index: u32) { + let res = ::compio::runtime::submit(FilesUpdate::remove(index)) + .await + .0; + debug_assert!( + res.is_ok(), + "giving back slot {index} of the file table: {res:?}" + ); + } + + /// `IORING_OP_FILES_UPDATE` of one slot. + struct FilesUpdate { + /// The descriptor to install, or `-1` to empty the slot; once a + /// descriptor is installed, the slot the kernel allocated for it. + fd: RawFd, + offset: i32, + /// The socket being installed, open until the kernel holds it. + _socket: Option>, + } + + impl FilesUpdate { + fn install(socket: SharedFd) -> Self { + Self { + fd: socket.as_raw_fd(), + offset: IORING_FILE_INDEX_ALLOC, + _socket: Some(socket), + } + } + + fn remove(index: u32) -> Self { + Self { + fd: -1, + offset: index.cast_signed(), + _socket: None, + } + } + + /// The slot an install took. + const fn slot(&self) -> u32 { + self.fd.cast_unsigned() + } + } + + fn unsupported() -> io::Error { + io::Error::new(io::ErrorKind::Unsupported, "a file table needs io_uring") + } + + // SAFETY: the entry points at `fd`, which the operation owns and compio + // keeps in place until the operation completes. + unsafe impl IourOpCode for FilesUpdate { + type Control = (); + + fn create_entry(&mut self, (): &mut ()) -> OpEntry { + // The kernel writes the slot it allocates back into `fd`. + io_uring::opcode::FilesUpdate::new((&raw mut self.fd).cast_const(), 1) + .offset(self.offset) + .build() + .into() + } + + fn call_blocking(&mut self, (): &mut ()) -> io::Result { + Err(unsupported()) + } + } + + // SAFETY: nothing is submitted; a stream takes a slot only on io_uring. + unsafe impl PollOpCode for FilesUpdate { + type Control = (); + + fn pre_submit(&mut self, (): &mut ()) -> io::Result { + Err(unsupported()) + } + + fn operate(&mut self, (): &mut ()) -> Poll> { + Poll::Ready(Err(unsupported())) + } + } +} + /// Zero-copy sends on io_uring. #[cfg(target_os = "linux")] mod zero_copy { @@ -831,7 +1431,7 @@ mod zero_copy { use futures_core::stream::FusedStream as _; use futures_lite::StreamExt as _; - use super::Socket; + use super::{Socket, Stream}; /// The caller's buffers, shared by the sends of one write: a short send /// is continued by another while the first still holds what it sent. @@ -850,7 +1450,7 @@ mod zero_copy { /// room in the socket with a poll (`io_net_retry`), so a write is one /// send and one notification unless an error cuts it short. pub(super) async fn write_all( - socket: &S, + stream: &Stream, buf: B, ) -> BufResult<(), Option> { let flags = SendFlags::from_bits_retain(libc::MSG_WAITALL.cast_unsigned()); @@ -864,9 +1464,13 @@ mod zero_copy { } let rest = Shared(Rc::clone(&buf)).slice(written); let res = if zero_copy { - send(SendVectoredZc::new(socket.shared_fd(), rest, flags)).await + let op = SendVectoredZc::new(stream.inner.shared_fd(), rest, flags); + match stream.file.address(op) { + Ok(op) => send(op).await, + Err((e, _)) => Err(e), + } } else { - socket.write_vectored(rest).await.0 + stream.send_vectored(rest).await.0 }; match res { Ok(0) => { @@ -974,13 +1578,14 @@ mod recv { use compio_buf::{BufResult, IoBufMut}; use futures_lite::StreamExt as _; - use super::{RecvBuf, Socket}; + use super::{RecvBuf, Socket, files}; use crate::{ Received, splice::{Pipe, Spliced}, }; - type Op = WithCancel>>>; + type Op = + WithCancel>>>>; /// The receive side of one stream. #[derive(Default)] @@ -1039,9 +1644,10 @@ mod recv { } impl Armed { - fn new(socket: &S, ring: &BufferPool) -> io::Result { + fn new(socket: &S, file: &files::File, ring: &BufferPool) -> io::Result { // A length of 0 receives up to a whole ring buffer at a time. let op = RecvMulti::new(socket.shared_fd(), ring, 0, RecvFlags::empty())?; + let op = file.address(op).map_err(|(e, _)| e)?; let cancel = CancelToken::new(); let op = ::compio::runtime::submit_multi(op) .into_managed(ring.clone()) @@ -1074,7 +1680,7 @@ mod recv { impl State { /// The next received bytes, waiting for them if none are held. - async fn next(&mut self, socket: &S) -> io::Result { + async fn next(&mut self, socket: &S, file: &files::File) -> io::Result { loop { if let Some(buf) = self.held.pop_front() { return Ok(Next::Buf(buf)); @@ -1101,7 +1707,7 @@ mod recv { let Some(ring) = ring(&mut self.ring) else { return Ok(Next::ReadInstead); }; - slot.insert(Armed::new(socket, &ring)?) + slot.insert(Armed::new(socket, file, &ring)?) } }; match armed.op.next().await { @@ -1157,6 +1763,7 @@ mod recv { pub(super) async fn read_ahead( &self, socket: &S, + file: &files::File, mut buf: B, ) -> Copied { if buf.as_uninit().is_empty() { @@ -1164,7 +1771,7 @@ mod recv { return Copied::Done(BufResult(Ok(0), buf)); } let mut state = self.0.borrow_mut(); - match state.next(socket).await { + match state.next(socket, file).await { Ok(Next::Buf(mut held)) => { let n = copy_into(&mut buf, &held); held.advance(n); @@ -1190,8 +1797,9 @@ mod recv { pub(super) async fn read_owned( &self, socket: &S, + file: &files::File, ) -> io::Result> { - Ok(match self.0.borrow_mut().next(socket).await? { + Ok(match self.0.borrow_mut().next(socket, file).await? { Next::Buf(buf) => Received::Lent(buf), Next::Eof => Received::End, Next::ReadInstead => Received::ReadInstead, @@ -1349,17 +1957,17 @@ pub struct TcpListener { inner: net::TcpListener, #[cfg(target_os = "linux")] accepts: multishot::Accepts, - /// The setting each accepted stream shares. - zero_copy: ZeroCopySend, + /// The settings each accepted stream shares. + io: Compio, } impl TcpListener { - fn new(inner: net::TcpListener, zero_copy: ZeroCopySend) -> Self { + fn new(inner: net::TcpListener, io: Compio) -> Self { Self { #[cfg(target_os = "linux")] accepts: multishot::Accepts::new(&inner), inner, - zero_copy, + io, } } @@ -1388,10 +1996,10 @@ impl Listener for TcpListener { // completions, so it reports no peer. let stream = net::TcpStream::from_std(accepted?.into())?; let peer = stream.peer_addr()?; - return Ok((TcpStream::new(stream, self.zero_copy.clone()), peer)); + return Ok((self.io.tcp_stream(stream), peer)); } let (stream, peer) = self.inner.accept().await?; - Ok((TcpStream::new(stream, self.zero_copy.clone()), peer)) + Ok((self.io.tcp_stream(stream), peer)) } fn local_addr(&self) -> io::Result { @@ -1406,18 +2014,18 @@ pub struct UnixListener { inner: net::UnixListener, #[cfg(target_os = "linux")] accepts: multishot::Accepts, - /// The setting each accepted stream shares. - zero_copy: ZeroCopySend, + /// The settings each accepted stream shares. + io: Compio, } #[cfg(unix)] impl UnixListener { - fn new(inner: net::UnixListener, zero_copy: ZeroCopySend) -> Self { + fn new(inner: net::UnixListener, io: Compio) -> Self { Self { #[cfg(target_os = "linux")] accepts: multishot::Accepts::new(&inner), inner, - zero_copy, + io, } } @@ -1448,7 +2056,7 @@ impl Listener for UnixListener { }; #[cfg(not(target_os = "linux"))] let (stream, _) = self.inner.accept().await?; - let stream = UnixStream::new(stream, self.zero_copy.clone()); + let stream = UnixStream::new(stream, &self.io); let peer = stream.peer_addr()?; Ok((stream, peer)) } diff --git a/io/tests/io.rs b/io/tests/io.rs index 4577653..f4227e4 100644 --- a/io/tests/io.rs +++ b/io/tests/io.rs @@ -937,13 +937,18 @@ mod compio { #[cfg(target_os = "linux")] #[test] fn one_accept_submission_serves_many_connections() { - use http_kit_io::{Listen as _, Listener as _, Timer as _, compio::Compio}; + use http_kit_io::{ + Listen as _, Listener as _, Timer as _, + compio::{Compio, FileTable}, + }; const CONNECTIONS: usize = 32; run(async { // The driver arms a poll on its own wake-up eventfd the first time - // it parks; park it once so that submission is not counted. - let io = Compio::default(); + // it parks; park it once so that submission is not counted. Each + // stream would take a slot in the file table with a submission + // of its own; this counts accepts alone. + let io = Compio::default().with_file_table(FileTable::new(0)); io.sleep(std::time::Duration::from_millis(1)).await; let listener = io.bind_tcp(super::loopback()).expect("bind"); @@ -1054,6 +1059,9 @@ mod compio { } }) .await; + // The stream takes its slot with a submission of its own; count + // from once it holds it, so the receive is armed on the slot. + until_held(&socket_name(&stream), 1).await; let before = submitted(); let mut received = Vec::with_capacity(data.len()); let mut reads = 0; @@ -1471,7 +1479,7 @@ mod compio { let sent = data.clone(); run(async move { let setting = ZeroCopySend::new(sent.len() + 1); - let io = Compio::with_zero_copy(setting.clone()); + let io = Compio::default().with_zero_copy(setting.clone()); let listener = io.bind_tcp(super::loopback()).expect("bind"); addr_tx .send(listener.local_addr().expect("local address")) @@ -1501,6 +1509,243 @@ mod compio { expected.extend_from_slice(&[0; 64 * 1024]); super::assert_bytes(&received, &expected); } + + /// The name the kernel gives a socket's file, as `/proc` lists it. + #[cfg(target_os = "linux")] + fn socket_name(socket: &impl std::os::fd::AsFd) -> String { + use std::os::fd::AsRawFd as _; + + let link = format!("/proc/self/fd/{}", socket.as_fd().as_raw_fd()); + std::fs::read_link(link) + .expect("socket link") + .into_os_string() + .into_string() + .expect("socket name") + } + + /// How many slots of the current ring's file table hold the file named + /// `name`, as the kernel lists them (`io_uring_show_fdinfo`). + #[cfg(target_os = "linux")] + fn slots_holding(name: &str) -> usize { + use std::os::fd::AsRawFd as _; + + let ring = ::compio::runtime::Runtime::with_current(|runtime| runtime.as_raw_fd()); + let info = + std::fs::read_to_string(format!("/proc/self/fdinfo/{ring}")).expect("ring fdinfo"); + info.lines() + .skip_while(|line| !line.starts_with("UserFiles:")) + .skip(1) + .take_while(|line| { + line.split(':') + .next() + .is_some_and(|slot| slot.trim().parse::().is_ok()) + }) + .filter(|line| line.trim_end().ends_with(name)) + .count() + } + + /// Waits until `held` slots hold the file named `name`. + #[cfg(target_os = "linux")] + async fn until_held(name: &str, held: usize) { + within(std::time::Duration::from_secs(5), "the file table", async { + while slots_holding(name) != held { + ::compio::time::sleep(std::time::Duration::from_millis(1)).await; + } + }) + .await; + } + + /// Once a stream holds its slot, its operations name the socket by it: + /// with the descriptor pointed elsewhere, the multishot receive, copying + /// and zero-copy sends and the shutdown still reach the socket. Dropping + /// the stream gives the slot back. + #[cfg(target_os = "linux")] + #[test] + fn a_stream_operates_through_its_registered_slot() { + use std::{ + io::{Read as _, Write as _}, + os::fd::{AsFd as _, AsRawFd as _}, + }; + + use http_kit_io::{BufResult, Listen as _, Listener as _, Read as _, Write as _}; + + let greeting = "你好 こんにちは 안녕하세요".as_bytes(); + let data = super::payload()[..256 * 1024].to_vec(); + let (addr_tx, addr_rx) = std::sync::mpsc::channel(); + let peer = std::thread::spawn(move || { + let addr = addr_rx.recv().expect("listener address"); + let mut stream = std::net::TcpStream::connect(addr).expect("connect"); + stream.write_all(greeting).expect("greet"); + let mut received = Vec::new(); + stream.read_to_end(&mut received).expect("read"); + received + }); + + let sent = data.clone(); + run(async move { + let io = http_kit_io::compio::Compio::default(); + assert!(sent.len() >= io.zero_copy().threshold()); + let listener = io.bind_tcp(super::loopback()).expect("bind"); + addr_tx + .send(listener.local_addr().expect("local address")) + .expect("send address"); + let (mut stream, _) = listener.accept().await.expect("accept"); + let name = socket_name(&stream); + until_held(&name, 1).await; + + let null = std::fs::File::open("/dev/null").expect("/dev/null"); + // SAFETY: the stream's descriptor stays open, on /dev/null from + // here on, and the stream closes it when dropped. + let duped = unsafe { libc::dup2(null.as_raw_fd(), stream.as_fd().as_raw_fd()) }; + assert_ne!(duped, -1, "dup2: {}", std::io::Error::last_os_error()); + + let mut echoed = Vec::new(); + while echoed.len() < greeting.len() { + let BufResult(res, buf) = stream.read(Vec::with_capacity(64)).await; + assert_ne!(res.expect("read through the slot"), 0, "early end"); + echoed.extend_from_slice(&buf); + } + let BufResult(res, _) = stream.write_all_owned([echoed]).await; + res.expect("copying write through the slot"); + let BufResult(res, kept) = stream.write_all_owned([sent]).await; + res.expect("zero-copy write through the slot"); + assert!(kept.is_none(), "the write was not zero-copy"); + http_kit_io::Write::shutdown(&mut stream) + .await + .expect("shutdown through the slot"); + drop(stream); + until_held(&name, 0).await; + }); + let mut expected = greeting.to_vec(); + expected.extend_from_slice(&data); + super::assert_bytes(&peer.join().expect("peer thread"), &expected); + } + + /// A socket the ring's table has no room for stays on its descriptor and + /// works as any other. + #[cfg(target_os = "linux")] + #[test] + fn a_socket_the_table_has_no_room_for_uses_its_descriptor() { + use std::io::{Read as _, Write as _}; + + use http_kit_io::{ + BufResult, Listen as _, Listener as _, Read as _, Write as _, + compio::{Compio, FileTable}, + }; + + let greeting = "世界 せかい 세계".as_bytes(); + let (addr_tx, addr_rx) = std::sync::mpsc::channel(); + let peers = std::thread::spawn(move || { + let addr = addr_rx.recv().expect("listener address"); + let peers = [(); 2].map(|()| { + let mut stream = std::net::TcpStream::connect(addr).expect("connect"); + stream.write_all(greeting).expect("greet"); + stream + }); + peers.map(|mut stream| { + let mut received = Vec::new(); + stream.read_to_end(&mut received).expect("read"); + received + }) + }); + + run(async move { + let io = Compio::default().with_file_table(FileTable::new(1)); + let listener = io.bind_tcp(super::loopback()).expect("bind"); + addr_tx + .send(listener.local_addr().expect("local address")) + .expect("send address"); + let (first, _) = listener.accept().await.expect("accept"); + let (second, _) = listener.accept().await.expect("accept"); + let names = [socket_name(&first), socket_name(&second)]; + within( + std::time::Duration::from_secs(5), + "a slot for one of the sockets", + async { + while names.iter().map(|name| slots_holding(name)).sum::() == 0 { + ::compio::time::sleep(std::time::Duration::from_millis(1)).await; + } + }, + ) + .await; + // Both installers have run by now; the one that came second found + // the table full. + ::compio::time::sleep(std::time::Duration::from_millis(20)).await; + let held = names.each_ref().map(|name| slots_holding(name)); + assert_eq!(held.iter().sum::(), 1, "slots held: {held:?}"); + + for mut stream in [first, second] { + let mut echoed = Vec::new(); + while echoed.len() < greeting.len() { + let BufResult(res, buf) = stream.read(Vec::with_capacity(64)).await; + assert_ne!(res.expect("read"), 0, "early end"); + echoed.extend_from_slice(&buf); + } + let BufResult(res, _) = stream.write_all_owned([echoed]).await; + res.expect("write"); + http_kit_io::Write::shutdown(&mut stream) + .await + .expect("shutdown"); + } + }); + for received in peers.join().expect("peer thread") { + super::assert_bytes(&received, greeting); + } + } + + /// A table the kernel refuses to register is the stream's error, not a + /// silent loss of the slot. + #[cfg(target_os = "linux")] + #[test] + fn a_table_the_kernel_refuses_fails_the_stream() { + use std::io::Read as _; + + use http_kit_io::{ + BufResult, Listen as _, Listener as _, Write as _, + compio::{Compio, FileTable}, + }; + + /// One slot more than any ring may register (`IORING_MAX_FIXED_FILES`, + /// `io_uring/rsrc.h`). + const SLOTS: u32 = (1 << 20) + 1; + + let (addr_tx, addr_rx) = std::sync::mpsc::channel(); + let peer = std::thread::spawn(move || { + let addr = addr_rx.recv().expect("listener address"); + let mut stream = std::net::TcpStream::connect(addr).expect("connect"); + let mut received = Vec::new(); + // The stream may end in a reset once it fails. + let _ = stream.read_to_end(&mut received); + }); + + run(async move { + let io = Compio::default().with_file_table(FileTable::new(SLOTS)); + let listener = io.bind_tcp(super::loopback()).expect("bind"); + addr_tx + .send(listener.local_addr().expect("local address")) + .expect("send address"); + let (mut stream, _) = listener.accept().await.expect("accept"); + let error = within( + std::time::Duration::from_secs(5), + "the refused table", + async { + loop { + let BufResult(res, _) = stream.write(b"x".to_vec()).await; + if let Err(error) = res { + break error; + } + ::compio::time::sleep(std::time::Duration::from_millis(1)).await; + } + }, + ) + .await; + assert!( + error.to_string().contains("registering a file table"), + "unexpected error: {error}" + ); + }); + peer.join().expect("peer thread"); + } } #[cfg(feature = "tokio")]