diff --git a/h1/src/client.rs b/h1/src/client.rs index de33e1b..d895212 100644 --- a/h1/src/client.rs +++ b/h1/src/client.rs @@ -15,7 +15,7 @@ use async_channel::{Receiver, Sender}; use bytes::BytesMut; use http::{Method, StatusCode}; use http_kit::{Body, Deadlines, Request, Response}; -use http_kit_io::{Read, SendFile, SpliceRead, SpliceWrite, Split, Timer, Write}; +use http_kit_io::{ReadOwned, SendFile, SpliceRead, SpliceWrite, Split, Timer, Write}; use crate::{ ClientConfig, Error, ParseError, Timeout, Upgraded, body, @@ -270,7 +270,7 @@ pub fn handshake( config: ClientConfig, ) -> (SendRequest, Connection) where - R: Read + SpliceRead, + R: ReadOwned + SpliceRead, W: Write + SendFile + SpliceWrite, T: Timer, { @@ -296,7 +296,7 @@ pub fn handshake_stream( ) -> (SendRequest, Connection) where S: Split, - S::ReadHalf: SpliceRead, + S::ReadHalf: ReadOwned + SpliceRead, S::WriteHalf: SendFile + SpliceWrite, T: Timer, { @@ -312,7 +312,7 @@ enum Next { impl Connection where - R: Read + SpliceRead, + R: ReadOwned + SpliceRead, W: Write + SendFile + SpliceWrite, T: Timer, { @@ -499,7 +499,7 @@ enum Received { /// Reads one response head, every wait for bytes bounded by the same /// `deadline`, so a head that dribbles in cannot extend it. -async fn read_head( +async fn read_head( recv: &mut Recv, limits: &crate::Limits, method: &Method, diff --git a/h1/src/conn.rs b/h1/src/conn.rs index 042555f..3e332da 100644 --- a/h1/src/conn.rs +++ b/h1/src/conn.rs @@ -16,7 +16,8 @@ use futures_lite::{ }; use http_kit::{Body, BodyFrame, PipeOffer}; use http_kit_io::{ - BufResult, IntoInner as _, IoBuf as _, Read, SendFile, SpliceRead, SpliceWrite, Timer, Write, + BufResult, IntoInner as _, IoBuf as _, ReadOwned, Received, SendFile, SpliceRead, SpliceWrite, + Timer, Write, }; use crate::{ @@ -39,7 +40,7 @@ pub(crate) struct Recv { read_size: usize, } -impl Recv { +impl Recv { pub(crate) fn new(reader: R, read_size: usize) -> Self { Self { reader, @@ -51,11 +52,26 @@ impl Recv { /// Reads more bytes onto the end of the buffer, returning how many; `0` /// is the end of the stream. /// - /// The read lands in the buffer's spare capacity, and the bytes are - /// later split off that same allocation, so nothing received is copied. - /// If this future is dropped mid-read, the buffer goes with it; the - /// connection is unusable then and must close. + /// A stream that lends its own receive buffer has the bytes copied onto + /// the end of this one and gets its buffer back at once, so the + /// connection reserves room only for bytes that arrived. Otherwise the + /// read lands in the buffer's spare capacity, and the bytes are later + /// split off that same allocation, so nothing received is copied; if this + /// future is dropped mid-read then, the buffer goes with it, and the + /// connection is unusable and must close. pub(crate) async fn fill(&mut self) -> Result { + loop { + match self.reader.read_owned().await { + Ok(Received::Lent(bytes)) => { + self.buf.extend_from_slice(&bytes); + return Ok(bytes.len()); + } + Ok(Received::End) => return Ok(0), + Ok(Received::ReadInstead) => break, + Err(error) if error.kind() == io::ErrorKind::Interrupted => {} + Err(error) => return Err(error.into()), + } + } self.buf.reserve(self.read_size.max(1)); let mut buf = mem::take(&mut self.buf); loop { @@ -305,7 +321,7 @@ pub(crate) enum Pumped { /// A demand that carries a [`PipeOffer`] is answered with a pipe frame when /// the socket can splice, nothing is buffered and the decoder is in a raw /// run; every other frame is data or trailers, whatever was offered. -pub(crate) async fn pump_body>( +pub(crate) async fn pump_body>( recv: &mut Recv, mut decoder: BodyDecoder, limits: &Limits, @@ -349,7 +365,7 @@ pub(crate) async fn pump_body( +async fn next_frame( recv: &mut Recv, decoder: &mut BodyDecoder, limits: &Limits, diff --git a/h1/src/server.rs b/h1/src/server.rs index 2025f79..aae70b1 100644 --- a/h1/src/server.rs +++ b/h1/src/server.rs @@ -23,7 +23,7 @@ use bytes::{Bytes, BytesMut}; use futures_lite::future::{or, poll_once}; use http::{HeaderValue, Method, StatusCode, Version, response}; use http_kit::{Body, Request, Response}; -use http_kit_io::{Read, SendFile, SpliceRead, SpliceWrite, Split, Timer, Write}; +use http_kit_io::{ReadOwned, SendFile, SpliceRead, SpliceWrite, Split, Timer, Write}; use crate::{ Error, ParseError, ServerConfig, Timeout, Upgraded, body, @@ -109,7 +109,7 @@ impl std::fmt::Debug for ServerConnection { impl ServerConnection where - R: Read + SpliceRead, + R: ReadOwned + SpliceRead, W: Write + SendFile + SpliceWrite, T: Timer, { @@ -424,7 +424,7 @@ where impl ServerConnection where - R: Read + SpliceRead, + R: ReadOwned + SpliceRead, W: Write + SendFile + SpliceWrite, T: Timer, { diff --git a/h1/src/splice.rs b/h1/src/splice.rs index d8f6ea8..f90223d 100644 --- a/h1/src/splice.rs +++ b/h1/src/splice.rs @@ -24,7 +24,8 @@ pub(crate) enum Attempt { /// /// The splice takes at most the run and the offer's `max`, so it never /// reaches past the body into the next message's bytes, which stay in the -/// socket. +/// socket. A stream that received ahead of its reader puts those bytes into +/// the pipe before it splices from the socket. #[cfg(all(feature = "pipe", any(target_os = "linux", target_os = "android")))] pub(crate) async fn try_splice( recv: &mut Recv, diff --git a/h1/tests/splice.rs b/h1/tests/splice.rs index d58518d..d44c03b 100644 --- a/h1/tests/splice.rs +++ b/h1/tests/splice.rs @@ -31,8 +31,8 @@ use http_kit_h1::{ ClientConfig, ClientOutcome, ServerConfig, ServerConnection, ServerOutcome, Service, handshake, }; use http_kit_io::{ - BufResult, Dial, IoBuf, IoBufMut, IoVectoredBuf, Listen, Listener as _, Read, SendFile, - SpliceRead, SpliceWrite, Split, Timer, Write, + BufResult, Dial, IoBuf, IoBufMut, IoVectoredBuf, Listen, Listener as _, Read, ReadOwned, + Received, SendFile, SpliceRead, SpliceWrite, Split, Timer, Write, splice::{FilledPipe, Pipe, Spliced}, }; @@ -70,6 +70,14 @@ impl Read for SpyRead { } } +impl ReadOwned for SpyRead { + type Buf = R::Buf; + + fn read_owned(&mut self) -> impl Future>> { + self.inner.read_owned() + } +} + impl SpliceRead for SpyRead { fn can_splice(&self) -> bool { self.inner.can_splice() diff --git a/h1/tests/support/mod.rs b/h1/tests/support/mod.rs index b38d669..5f54d2e 100644 --- a/h1/tests/support/mod.rs +++ b/h1/tests/support/mod.rs @@ -18,8 +18,8 @@ use std::{ use async_channel::{Receiver, Sender}; use http_kit_io::{ - BufResult, IoBuf, IoBufMut, IoVectoredBuf, Read, SendFile, SpliceRead, SpliceWrite, Timer, - Write, + BufResult, IoBuf, IoBufMut, IoVectoredBuf, NoBuf, Read, ReadOwned, Received, SendFile, + SpliceRead, SpliceWrite, Timer, Write, }; /// A timer whose deadlines never pass, for tests that run no timeouts. @@ -193,6 +193,16 @@ impl Read for PipeReader { } } +/// The reads land in the caller's buffer, which is what the reads it records +/// are about. +impl ReadOwned for PipeReader { + type Buf = NoBuf; + + async fn read_owned(&mut self) -> io::Result> { + Ok(Received::ReadInstead) + } +} + impl Write for PipeWriter { async fn write(&mut self, buf: B) -> BufResult { let result = self.push(buf.as_init().to_vec()); diff --git a/io/src/compio.rs b/io/src/compio.rs index 56a7915..dffeb85 100644 --- a/io/src/compio.rs +++ b/io/src/compio.rs @@ -12,6 +12,24 @@ //! worker; on Apple platforms through `sendfile(2)` once the socket is //! writable. Elsewhere the region is read into a buffer and written. //! +//! On Linux over io_uring a stream receives with a multishot receive into +//! the provided buffer ring of the thread's runtime: one submission stays +//! armed and completes once per arrival, each into a ring buffer, until the +//! kernel ends it. A reader takes a ring buffer through +//! [`ReadOwned`](crate::ReadOwned), or has it copied into its own through +//! [`Read`](crate::Read); either way the buffer goes back to the ring as soon +//! as its bytes are consumed, so an idle connection holds none. When the +//! ring has no free buffer a read receives into the caller's buffer with a +//! submission of its own, and the next one arms again. A splice from the +//! socket first moves what the ring already received, so it never passes +//! those bytes. +//! +//! The ring is compio's buffer pool, one per runtime: size it on the +//! runtime's `ProactorBuilder` (`buffer_pool_size`, `buffer_pool_buffer_len`) +//! for the connections one thread serves. Buffers are consumed whole, one per +//! completion: compio reports no `IORING_CQE_F_BUF_MORE`, which incremental +//! consumption (`IOU_PBUF_RING_INC`) needs. +//! //! The poll-style capability is a [`CopyingPoll`]: a completion-based driver //! cannot read into a buffer borrowed for one poll. @@ -143,6 +161,9 @@ impl Fd for T {} /// A stream socket of compio's, which the operations below drive through a /// shared reference. 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; fn read(&self, buf: B) -> impl Future>; fn write(&self, buf: B) -> impl Future>; fn write_vectored(&self, buf: B) @@ -168,6 +189,11 @@ trait Socket: Clone + Fd + 'static { macro_rules! impl_socket { ($ty:ty) => { impl Socket for $ty { + #[cfg(target_os = "linux")] + fn shared_fd(&self) -> ::compio::driver::SharedFd { + ::compio::driver::ToSharedFd::to_shared_fd(self) + } + async fn read(&self, buf: B) -> BufResult { let mut stream = self; AsyncRead::read(&mut stream, buf).await @@ -242,6 +268,9 @@ pub struct Stream { /// inside the kernel. #[cfg(any(target_os = "linux", target_os = "android"))] nonblocking: std::cell::Cell, + /// The multishot receive and what it received ahead of the reader. + #[cfg(target_os = "linux")] + recv: recv::Recv, } impl Stream { @@ -252,6 +281,8 @@ impl Stream { pipe: std::cell::Cell::new(None), #[cfg(any(target_os = "linux", target_os = "android"))] nonblocking: std::cell::Cell::new(false), + #[cfg(target_os = "linux")] + recv: recv::Recv::default(), } } @@ -339,11 +370,40 @@ impl std::os::fd::AsFd for Stream { } impl SharedStream for Stream { - async fn read(&self, 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. - buf.clear(); - self.inner.read(buf).await + #[cfg(target_os = "linux")] + type Buf = RecvBuf; + #[cfg(not(target_os = "linux"))] + type Buf = crate::NoBuf; + + 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 + } + + async fn read_owned(&self) -> io::Result> { + #[cfg(target_os = "linux")] + { + self.recv.read_owned(&self.inner).await + } + #[cfg(not(target_os = "linux"))] + { + Ok(crate::Received::ReadInstead) + } + } + + async fn stop_read_ahead(&self) -> Vec { + #[cfg(target_os = "linux")] + { + self.recv.stop().await + } + #[cfg(not(target_os = "linux"))] + { + Vec::new() + } } async fn write(&self, buf: B) -> BufResult { @@ -409,6 +469,14 @@ impl SharedStream for Stream { } } +/// One receive into the caller's buffer, with a submission of its own. +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. + buf.clear(); + socket.read(buf).await +} + impl_stream!([S: Socket] Stream); impl IntoPoll for Stream { @@ -575,6 +643,10 @@ mod splice { /// Socket to pipe: waits for readability, then one non-blocking /// `splice(2)`. io_uring would run the splice on an io-wq worker, which /// an idle peer would hold. + /// + /// What the multishot receive took from the socket ahead of the reader + /// goes into the pipe first, from memory; only once nothing is held, and + /// the receive is no longer armed to take more, does the socket splice. pub(super) async fn to_pipe( stream: &Stream, pipe: crate::splice::Pipe, @@ -586,6 +658,11 @@ mod splice { pipe, )); } + #[cfg(target_os = "linux")] + let pipe = match stream.recv.splice_held(pipe, max).await { + super::recv::Held::Spliced(spliced) => return spliced, + super::recv::Held::None(pipe) => pipe, + }; if let Err(e) = nonblocking(stream) { return Err((e, pipe)); } @@ -630,6 +707,357 @@ mod splice { } } +/// Bytes a multishot receive put in a buffer of the runtime's ring, lent to +/// the reader by [`ReadOwned`](crate::ReadOwned). Dropping it hands the +/// buffer back to the ring, so a reader holds it only as long as it needs the +/// bytes. +#[cfg(target_os = "linux")] +pub struct RecvBuf { + buf: ::compio::driver::BufferRef, + /// How much of the buffer was already read. + start: usize, +} + +#[cfg(target_os = "linux")] +impl RecvBuf { + const fn new(buf: ::compio::driver::BufferRef) -> Self { + Self { buf, start: 0 } + } + + const fn advance(&mut self, n: usize) { + self.start += n; + } +} + +#[cfg(target_os = "linux")] +impl std::ops::Deref for RecvBuf { + type Target = [u8]; + + fn deref(&self) -> &[u8] { + &self.buf[self.start..] + } +} + +#[cfg(target_os = "linux")] +impl fmt::Debug for RecvBuf { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RecvBuf") + .field("len", &self.len()) + .finish_non_exhaustive() + } +} + +/// Multishot receive on io_uring into the runtime's provided buffer ring. +/// +/// The operation belongs to the stream, not to a read, so a read a caller +/// abandons (one raced against a timeout) leaves it armed with whatever it +/// already received. It is ended (cancelled, then drained, so nothing it +/// received is lost) only before the socket goes to someone else: a splice, +/// or for good when the reader stops reading ahead. +#[cfg(target_os = "linux")] +mod recv { + use std::{cell::RefCell, collections::VecDeque, io}; + + use ::compio::{ + driver::{ + SharedFd, + op::{RecvFlags, RecvMulti}, + }, + runtime::{ + BufferPool, CancelToken, Runtime, StreamExt as _, SubmitMultiManaged, WithCancel, + }, + }; + use compio_buf::{BufResult, IoBufMut}; + use futures_lite::StreamExt as _; + + use super::{RecvBuf, Socket}; + use crate::{ + Received, + splice::{Pipe, Spliced}, + }; + + type Op = WithCancel>>>; + + /// The receive side of one stream. + #[derive(Default)] + pub(super) struct Recv(RefCell); + + #[derive(Default)] + struct State { + /// The thread's ring, looked up on the first read: `Some(None)` where + /// there is none (the runtime is not on io_uring, or the kernel + /// refused the ring). + ring: Option>, + armed: Option, + /// Bytes received ahead of the reader, in order; the first may be + /// partly read. + held: VecDeque, + /// How the receive ended, reported once `held` is read. + end: Option, + /// The last receive ended for want of a free ring buffer: the next + /// read takes the caller's buffer instead of arming again at once. + exhausted: bool, + /// The reader stopped reading ahead. + stopped: bool, + } + + /// The armed receive and the token that cancels it. + struct Armed { + op: Op, + cancel: CancelToken, + } + + enum End { + /// Sticky, like the end of a socket. + Eof, + /// Reported once; the next read arms again. + Error(io::Error), + } + + enum Next { + Buf(RecvBuf), + Eof, + ReadInstead, + } + + /// What a copying read came to. + pub(super) enum Copied { + Done(BufResult), + /// Read into the caller's buffer with a receive of its own. + ReadInstead(B), + } + + /// What a splice found held. + pub(super) enum Held { + Spliced(Result), + /// Nothing: the socket holds everything not yet read. + None(Pipe), + } + + impl Armed { + fn new(socket: &S, 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 cancel = CancelToken::new(); + let op = ::compio::runtime::submit_multi(op) + .into_managed(ring.clone()) + .with_cancel(cancel.clone()); + Ok(Self { op, cancel }) + } + } + + /// Whether a receive ended because the ring had no free buffer: compio's + /// io_uring driver reports `ENOBUFS` as a [`io::ErrorKind::ResourceBusy`] + /// error of its own, without the OS code (compio-driver 0.12, + /// `create_result`), which a real `EBUSY` would carry. + fn ring_dry(error: &io::Error) -> bool { + error.kind() == io::ErrorKind::ResourceBusy && error.raw_os_error().is_none() + } + + /// The thread's ring, looked up once per stream. + fn ring(ring: &mut Option>) -> Option { + ring.get_or_insert_with(|| { + Runtime::with_current(|runtime| { + if runtime.driver_type().is_iouring() { + runtime.buffer_pool().ok() + } else { + None + } + }) + }) + .clone() + } + + impl State { + /// The next received bytes, waiting for them if none are held. + async fn next(&mut self, socket: &S) -> io::Result { + loop { + if let Some(buf) = self.held.pop_front() { + return Ok(Next::Buf(buf)); + } + match self.end.take() { + Some(End::Eof) => { + self.end = Some(End::Eof); + return Ok(Next::Eof); + } + Some(End::Error(error)) => return Err(error), + None => {} + } + if self.stopped || self.exhausted { + if self.armed.is_some() { + self.disarm().await; + continue; + } + self.exhausted = false; + return Ok(Next::ReadInstead); + } + let armed = match &mut self.armed { + Some(armed) => armed, + slot @ None => { + let Some(ring) = ring(&mut self.ring) else { + return Ok(Next::ReadInstead); + }; + slot.insert(Armed::new(socket, &ring)?) + } + }; + match armed.op.next().await { + // The kernel ended the receive with its last completion: + // arm another. + None => self.armed = None, + Some(Ok(Some(buf))) if !buf.is_empty() => { + return Ok(Next::Buf(RecvBuf::new(buf))); + } + Some(Ok(_)) => { + self.armed = None; + self.end = Some(End::Eof); + } + Some(Err(error)) if ring_dry(&error) => { + self.armed = None; + self.exhausted = true; + } + Some(Err(error)) => { + self.armed = None; + return Err(error); + } + } + } + } + + /// Ends the armed receive, keeping everything it received. + async fn disarm(&mut self) { + while let Some(armed) = &mut self.armed { + if !armed.cancel.is_cancelled() { + armed.cancel.clone().cancel(); + } + match armed.op.next().await { + None => self.armed = None, + Some(Ok(Some(buf))) if !buf.is_empty() => { + self.held.push_back(RecvBuf::new(buf)); + } + Some(Ok(_)) => self.end = Some(End::Eof), + Some(Err(error)) + if ring_dry(&error) || error.raw_os_error() == Some(libc::ECANCELED) => {} + Some(Err(error)) => self.end = Some(End::Error(error)), + } + } + } + } + + impl Recv { + /// [`Read::read`](crate::Read::read): copies what was received into + /// `buf`, keeping what does not fit. + #[expect( + clippy::await_holding_refcell_ref, + reason = "the stream's one reader holds its receive side for the read, released when the read ends or is dropped" + )] + pub(super) async fn read_ahead( + &self, + socket: &S, + mut buf: B, + ) -> Copied { + if buf.as_uninit().is_empty() { + buf.clear(); + return Copied::Done(BufResult(Ok(0), buf)); + } + let mut state = self.0.borrow_mut(); + match state.next(socket).await { + Ok(Next::Buf(mut held)) => { + let n = copy_into(&mut buf, &held); + held.advance(n); + if !held.is_empty() { + state.held.push_front(held); + } + Copied::Done(BufResult(Ok(n), buf)) + } + Ok(Next::Eof) => { + buf.clear(); + Copied::Done(BufResult(Ok(0), buf)) + } + Ok(Next::ReadInstead) => Copied::ReadInstead(buf), + Err(error) => Copied::Done(BufResult(Err(error), buf)), + } + } + + /// [`ReadOwned::read_owned`](crate::ReadOwned::read_owned). + #[expect( + clippy::await_holding_refcell_ref, + reason = "the stream's one reader holds its receive side for the read, released when the read ends or is dropped" + )] + pub(super) async fn read_owned( + &self, + socket: &S, + ) -> io::Result> { + Ok(match self.0.borrow_mut().next(socket).await? { + Next::Buf(buf) => Received::Lent(buf), + Next::Eof => Received::End, + Next::ReadInstead => Received::ReadInstead, + }) + } + + /// [`ReadAhead::stop_read_ahead`](crate::ReadAhead::stop_read_ahead). + #[expect( + clippy::await_holding_refcell_ref, + reason = "the stream's one reader holds its receive side while the receive drains" + )] + pub(super) async fn stop(&self) -> Vec { + let mut state = self.0.borrow_mut(); + state.stopped = true; + state.disarm().await; + let mut bytes = Vec::with_capacity(state.held.iter().map(|buf| buf.len()).sum()); + for buf in state.held.drain(..) { + bytes.extend_from_slice(&buf); + } + bytes + } + + /// Puts held bytes, at most `max`, into `pipe` ahead of anything the + /// socket splices. With none held, ends the armed receive first, so + /// that the socket holds everything not yet read. + #[expect( + clippy::await_holding_refcell_ref, + reason = "the stream's one reader holds its receive side while the receive drains" + )] + pub(super) async fn splice_held(&self, pipe: Pipe, max: usize) -> Held { + let mut state = self.0.borrow_mut(); + if state.held.is_empty() { + state.disarm().await; + } + if let Some(mut buf) = state.held.pop_front() { + let n = buf.len().min(max); + let filled = pipe.fill(&buf[..n]).map(|filled| { + buf.advance(filled.len()); + Spliced::Filled(filled) + }); + if !buf.is_empty() { + state.held.push_front(buf); + } + return Held::Spliced(filled); + } + match state.end.take() { + Some(End::Eof) => { + state.end = Some(End::Eof); + Held::Spliced(Ok(Spliced::End(pipe))) + } + Some(End::Error(error)) => Held::Spliced(Err((error, pipe))), + None => Held::None(pipe), + } + } + } + + /// Copies as much of `bytes` as fits into `buf`, from its start, and sets + /// its length to that. + fn copy_into(buf: &mut B, bytes: &[u8]) -> usize { + let room = buf.as_uninit(); + let n = room.len().min(bytes.len()); + // SAFETY: `room` has at least `n` bytes, and a ring buffer never + // overlaps the caller's. + unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), room.as_mut_ptr().cast(), n) }; + // SAFETY: the first `n` bytes were initialized just now. + unsafe { buf.set_len(n) }; + n + } +} + /// Multishot accept on io_uring: one submission stays armed on the /// listener and completes once per connection, until the kernel ends it. /// diff --git a/io/src/lib.rs b/io/src/lib.rs index 6d7f89d..f9db740 100644 --- a/io/src/lib.rs +++ b/io/src/lib.rs @@ -11,6 +11,10 @@ //! # Traits //! //! - [`Read`], [`Write`], [`Split`] and [`Listener`]: stream sockets. +//! - [`ReadOwned`]: reads that lend the stream's own receive buffer (a ring +//! buffer the kernel received into, on io_uring), and [`ReadAhead`] to stop +//! a stream receiving ahead of its reader before the kernel takes the +//! socket. //! - [`StreamSocket`]: the bounds a connected stream carries for the protocol //! crates. //! - [`TcpOptions`]: socket options of a connected TCP stream. @@ -85,7 +89,10 @@ pub use crate::{ poll::{CopyingPoll, IntoPoll}, resolve::{Answer, Family, Resolve}, splice::{SpliceRead, SpliceWrite}, - stream::{Listener, Read, ReadHalf, SendFile, Split, TcpOptions, Write, WriteHalf}, + stream::{ + Listener, NoBuf, Read, ReadAhead, ReadHalf, ReadOwned, Received, SendFile, Split, + TcpOptions, Write, WriteHalf, + }, time::Timer, }; @@ -98,12 +105,20 @@ pub use crate::stream::WriteControl; /// It is the bound the protocol crates put on the halves of a stream, named /// once. It is implemented for every type that meets it. pub trait StreamSocket: - Read + Write + Split + 'static + ReadOwned + + ReadAhead + + Write + + Split + + 'static { } impl StreamSocket for S where - S: Read + Write + Split + 'static + S: ReadOwned + + ReadAhead + + Write + + Split + + 'static { } diff --git a/io/src/readiness.rs b/io/src/readiness.rs index 208dae1..6b37ccb 100644 --- a/io/src/readiness.rs +++ b/io/src/readiness.rs @@ -184,6 +184,8 @@ pub(crate) fn socket_addr(addr: SockAddr) -> io::Result { } impl SharedStream for Stream { + type Buf = crate::NoBuf; + async fn read(&self, mut buf: B) -> BufResult { let res = self.0.read_io(|sock| sock.recv(buf.as_uninit())).await; if let Ok(n) = res { @@ -193,6 +195,17 @@ impl SharedStream for Stream { BufResult(res, buf) } + /// A readiness reactor receives into the caller's buffer once the socket + /// is readable; it has none of its own to lend. + async fn read_owned(&self) -> io::Result> { + Ok(crate::Received::ReadInstead) + } + + /// Nothing is received before the reader asks. + async fn stop_read_ahead(&self) -> Vec { + Vec::new() + } + async fn write(&self, buf: B) -> BufResult { let res = self.0.write_io(|sock| send(sock, buf.as_init())).await; BufResult(res, buf) diff --git a/io/src/stream.rs b/io/src/stream.rs index 7060612..65e2771 100644 --- a/io/src/stream.rs +++ b/io/src/stream.rs @@ -1,6 +1,6 @@ //! Stream sockets. -use std::{borrow::Borrow, fmt, fs::File, future::Future, io, rc::Rc}; +use std::{borrow::Borrow, fmt, fs::File, future::Future, io, ops::Deref, rc::Rc}; use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBuf}; @@ -12,9 +12,68 @@ pub trait Read { /// spare capacity (`buf.slice(buf.buf_len()..)`). /// /// `Ok(0)` into a buffer with spare capacity means end of stream. + /// + /// A stream that receives ahead of its reader ([`ReadOwned`]) copies what + /// it already received first, and keeps what does not fit for the next + /// read. fn read(&mut self, buf: B) -> impl Future>; } +/// What [`ReadOwned::read_owned`] came back with. +#[derive(Debug)] +pub enum Received { + /// Received bytes in a buffer the stream owns. Dropping it hands the + /// buffer back to the stream. + Lent(B), + /// The end of the stream. + End, + /// The stream has no buffer of its own to lend now: read into one of the + /// caller's with [`Read::read`]. + ReadInstead, +} + +/// Reads that lend the stream's own receive buffer, so a reader waiting for +/// bytes holds no buffer of its own. +/// +/// The compio implementation on Linux receives with a multishot receive into +/// the provided buffer ring of the thread's runtime: one submission serves +/// many receives, and a connection pins a buffer only from the moment bytes +/// arrive until its reader is done with them. Every other implementation +/// answers [`Received::ReadInstead`], and so does the compio one whenever no +/// ring buffer is free. +pub trait ReadOwned: Read { + /// The received bytes as lent out. + type Buf: Deref; + + /// Waits for received bytes and lends them out. + fn read_owned(&mut self) -> impl Future>>; +} + +/// The [`ReadOwned::Buf`] of a stream that never lends a buffer: there is +/// none. +#[derive(Debug, Clone, Copy)] +pub enum NoBuf {} + +impl Deref for NoBuf { + type Target = [u8]; + + fn deref(&self) -> &[u8] { + match *self {} + } +} + +/// A stream that may receive ahead of its reader, and can stop. +/// +/// Whoever hands the socket's receive side to the kernel (kernel TLS) needs +/// every byte not yet received to still be in the socket. +pub trait ReadAhead { + /// Stops receiving ahead of the reader, for good: from now on the socket + /// is read only when the reader asks. Returns the bytes received ahead + /// and not yet read, in order; the end of the stream or an error the + /// receive ended with is reported by the next read. + fn stop_read_ahead(&mut self) -> impl Future>; +} + /// Writes from owned buffers. pub trait Write { /// Writes some of `buf`'s initialized bytes, returning how many. @@ -201,7 +260,11 @@ pub(crate) mod sealed { /// halves can drive one socket concurrently. Every socket this crate /// implements is one. pub trait SharedStream: 'static { + type Buf: std::ops::Deref; + fn read(&self, buf: B) -> impl Future>; + fn read_owned(&self) -> impl Future>>; + fn stop_read_ahead(&self) -> impl Future>; fn write(&self, buf: B) -> impl Future>; fn write_vectored( &self, @@ -250,6 +313,26 @@ macro_rules! impl_stream { } } + impl<$($generics)*> $crate::stream::ReadOwned for $ty { + type Buf = <$ty as $crate::stream::sealed::SharedStream>::Buf; + + fn read_owned( + &mut self, + ) -> impl ::std::future::Future< + Output = ::std::io::Result<$crate::stream::Received>, + > { + $crate::stream::sealed::SharedStream::read_owned(&*self) + } + } + + impl<$($generics)*> $crate::stream::ReadAhead for $ty { + fn stop_read_ahead( + &mut self, + ) -> impl ::std::future::Future> { + $crate::stream::sealed::SharedStream::stop_read_ahead(&*self) + } + } + impl<$($generics)*> $crate::stream::Write for $ty { fn write( &mut self, @@ -406,6 +489,20 @@ impl Read for ReadHalf { } } +impl ReadOwned for ReadHalf { + type Buf = S::Buf; + + fn read_owned(&mut self) -> impl Future>> { + self.0.read_owned() + } +} + +impl ReadAhead for ReadHalf { + fn stop_read_ahead(&mut self) -> impl Future> { + self.0.stop_read_ahead() + } +} + impl Write for WriteHalf { fn write(&mut self, buf: B) -> impl Future> { self.0.write(buf) diff --git a/io/tests/io.rs b/io/tests/io.rs index 4c9f104..2cb7e87 100644 --- a/io/tests/io.rs +++ b/io/tests/io.rs @@ -990,6 +990,199 @@ mod compio { }); } + /// A runtime whose provided buffer ring has `buffers` buffers of `len` + /// bytes each. + #[cfg(target_os = "linux")] + fn run_with_ring(buffers: u16, len: usize, future: F) -> F::Output { + let mut proactor = ::compio::driver::ProactorBuilder::new(); + proactor + .buffer_pool_size(std::num::NonZeroU16::new(buffers).expect("a ring has buffers")) + .buffer_pool_buffer_len(len); + ::compio::runtime::RuntimeBuilder::new() + .with_proactor(proactor) + .build() + .expect("compio runtime") + .block_on(future) + } + + /// One accepted connection, with the peer on a thread of its own. + #[cfg(target_os = "linux")] + async fn accept_one( + peer: impl FnOnce(std::net::TcpStream) + Send + 'static, + ) -> (http_kit_io::compio::TcpStream, std::thread::JoinHandle<()>) { + use http_kit_io::{Listen as _, Listener as _, compio::Compio}; + + let listener = Compio.bind_tcp(super::loopback()).expect("bind"); + let addr = listener.local_addr().expect("local address"); + let peer = + std::thread::spawn(move || peer(std::net::TcpStream::connect(addr).expect("connect"))); + let (stream, _) = listener.accept().await.expect("accept"); + (stream, peer) + } + + /// One multishot receive serves every read of a connection. The ring + /// holds a quarter of what the connection carries, so each buffer must + /// go back to it as soon as its bytes are read; a buffer kept would run + /// the ring dry, end the receive and take a submission to arm it again. + #[cfg(target_os = "linux")] + #[test] + fn one_receive_submission_serves_many_reads() { + use std::io::{Read as _, Write as _}; + + use http_kit_io::{ReadOwned as _, Received, Timer as _, compio::Compio}; + + const BUFFERS: u16 = 4; + const BUFFER_LEN: usize = 4096; + const MESSAGE: usize = 1024; + const MESSAGES: usize = 64; + + let data = super::payload()[..MESSAGE * MESSAGES].to_vec(); + let sent = data.clone(); + run_with_ring(BUFFERS, BUFFER_LEN, async move { + // 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. + Compio.sleep(std::time::Duration::from_millis(1)).await; + // Each message waits for the one before it to be read, so the + // receive always finds a free buffer. + let (mut stream, peer) = accept_one(move |mut client| { + for message in sent.chunks(MESSAGE) { + client.write_all(message).expect("send"); + client.read_exact(&mut [0]).expect("wait for the read"); + } + }) + .await; + let before = submitted(); + let mut received = Vec::with_capacity(data.len()); + let mut reads = 0; + while received.len() < data.len() { + let Received::Lent(buf) = stream.read_owned().await.expect("read") else { + panic!("the ring ran dry or the peer closed early"); + }; + reads += 1; + received.extend_from_slice(&buf); + drop(buf); + if received.len() % MESSAGE == 0 { + // The acknowledgement is a send, which is a submission; + // count it out. + let ack = submitted(); + let http_kit_io::BufResult(res, _) = + http_kit_io::Write::write(&mut stream, vec![0]).await; + res.expect("acknowledge"); + assert_eq!(submitted() - ack, 1, "an acknowledgement is one send"); + } + } + let acks = u32::try_from(MESSAGES).expect("message count"); + assert_eq!( + submitted() - before - acks, + 1, + "receive submissions for {reads} reads" + ); + assert!(reads >= MESSAGES, "{reads} reads"); + super::assert_bytes(&received, &data); + peer.join().expect("peer thread"); + }); + } + + /// A connection that carries far more than the ring holds, sent as fast + /// as the peer can, completes: buffers go back to the ring as they are + /// read, and a read that finds the ring dry receives into the caller's + /// buffer instead. Reads smaller than a ring buffer take it in parts. + #[cfg(target_os = "linux")] + #[test] + fn a_connection_reads_more_than_the_ring_holds() { + use std::io::Write as _; + + let data = super::payload(); + let sent = data.clone(); + run_with_ring(2, 4096, async move { + let (mut stream, peer) = accept_one(move |mut client| { + client.write_all(&sent).expect("send"); + }) + .await; + let mut received = Vec::with_capacity(data.len()); + let mut buf = Vec::with_capacity(1500); + loop { + let http_kit_io::BufResult(res, filled) = + http_kit_io::Read::read(&mut stream, buf).await; + if res.expect("read") == 0 { + break; + } + received.extend_from_slice(&filled); + buf = filled; + } + super::assert_bytes(&received, &data); + peer.join().expect("peer thread"); + }); + } + + /// A splice from a socket that has received ahead of its reader moves + /// those bytes first, so it never passes them. + #[cfg(target_os = "linux")] + #[test] + fn a_splice_takes_the_bytes_received_ahead_first() { + use std::io::Write as _; + + use http_kit_io::{ + SpliceRead as _, + splice::{Pipe, Spliced}, + }; + + let data = "世界 せかい 세계 ".repeat(512).into_bytes(); + let sent = data.clone(); + run_with_ring(8, 4096, async move { + let (mut stream, peer) = accept_one(move |mut client| { + client.write_all(&sent).expect("send"); + }) + .await; + // A small read leaves the rest of its ring buffer held. + let http_kit_io::BufResult(res, head) = + http_kit_io::Read::read(&mut stream, Vec::with_capacity(10)).await; + assert_eq!(res.expect("read"), 10); + let mut received = head; + let mut pipe = Pipe::new().expect("pipe"); + loop { + match stream.splice_to_pipe(pipe, usize::MAX).await { + Ok(Spliced::Filled(filled)) => { + let (bytes, empty) = filled.read_out().expect("read the pipe"); + received.extend_from_slice(&bytes); + pipe = empty; + } + Ok(Spliced::End(_)) => break, + Err((error, _)) => panic!("splice: {error}"), + } + } + super::assert_bytes(&received, &data); + peer.join().expect("peer thread"); + }); + } + + /// Stopping the read-ahead hands back what was received and not read, + /// and the reads after it continue from the socket where that ended. + #[cfg(target_os = "linux")] + #[test] + fn stopping_read_ahead_hands_back_what_was_received() { + use std::io::Write as _; + + use http_kit_io::ReadAhead as _; + + let data = "世界 せかい 세계 ".repeat(4096).into_bytes(); + let sent = data.clone(); + run_with_ring(8, 4096, async move { + let (mut stream, peer) = accept_one(move |mut client| { + client.write_all(&sent).expect("send"); + }) + .await; + let http_kit_io::BufResult(res, head) = + http_kit_io::Read::read(&mut stream, Vec::with_capacity(10)).await; + assert_eq!(res.expect("read"), 10); + let mut received = head; + received.extend_from_slice(&stream.stop_read_ahead().await); + received.extend_from_slice(&super::read_to_end(&mut stream).await); + super::assert_bytes(&received, &data); + peer.join().expect("peer thread"); + }); + } + /// The io-wq workers of this process that are doing something other than /// waiting for work, each as its thread id and kernel wait channel. #[cfg(any(target_os = "linux", target_os = "android"))] diff --git a/tls/src/kernel/linux.rs b/tls/src/kernel/linux.rs index 30721d4..7e1153e 100644 --- a/tls/src/kernel/linux.rs +++ b/tls/src/kernel/linux.rs @@ -33,8 +33,8 @@ use std::{ }; use http_kit_io::{ - BufResult, IoBuf, IoBufMut, IoVectoredBuf, Read, SendFile, SpliceRead, SpliceWrite, Write, - WriteControl, + BufResult, IoBuf, IoBufMut, IoVectoredBuf, Read, ReadAhead, SendFile, SpliceRead, SpliceWrite, + Write, WriteControl, splice::{FilledPipe, Pipe, Spliced}, }; use ktls_sys::bindings as uapi; @@ -47,7 +47,7 @@ use crate::{ engine::{Engine, copy_into, tls_error}, record::{self, Boundary, CLOSE_NOTIFY, Cipher, Control, KEY_UPDATE, MAX_FRAGMENT, Message}, side::Side, - stream::{KernelRead, KernelWrite}, + stream::{KernelRead, KernelWrite, TransportRead}, }; /// What the halves share: rustls's key schedule, which both directions @@ -73,10 +73,10 @@ pub(crate) enum Switched { /// kernel supports its suite. /// /// The kernel must start at a record boundary: everything buffered is -/// decrypted here first, and a record that only partly arrived is read to -/// its end, no further. The plaintext comes back to be read before the -/// socket. -pub(crate) async fn switch( +/// decrypted here first, what the transport received ahead of its reader +/// included, and a record that only partly arrived is read to its end, no +/// further. The plaintext comes back to be read before the socket. +pub(crate) async fn switch( mut engine: Engine, reader: &mut R, writer: &mut W, @@ -120,6 +120,11 @@ pub(crate) async fn switch( ))); } } + // The kernel decrypts only what is still in the socket. Ciphertext the + // transport already received goes to rustls with the rest, and from here + // on the socket is read only as far as the boundary below asks. + let ahead = as KernelRead>::stop_read_ahead(reader).await; + engine.incoming.extend_from_slice(&ahead); loop { engine.pump()?; if engine.peer_closed { @@ -809,13 +814,17 @@ impl Writer { } } -impl KernelRead for Reader { +impl KernelRead for Reader { const SOCKET: bool = true; fn adopt(reader: Self) -> Self { reader } + fn stop_read_ahead(reader: &mut R) -> impl Future> { + reader.stop_read_ahead() + } + fn get_ref(&self) -> &R { Self::get_ref(self) } diff --git a/tls/src/stream.rs b/tls/src/stream.rs index 98793c8..f41e075 100644 --- a/tls/src/stream.rs +++ b/tls/src/stream.rs @@ -5,8 +5,8 @@ use std::{borrow::Borrow, fmt, fs::File, future::Future, io}; #[cfg(any(target_os = "linux", target_os = "android"))] use http_kit_io::splice::{FilledPipe, Pipe, Spliced}; use http_kit_io::{ - BufResult, IoBuf, IoBufMut, IoVectoredBuf, Read, SendFile, SpliceRead, SpliceWrite, Split, - Write, + BufResult, IoBuf, IoBufMut, IoVectoredBuf, NoBuf, Read, ReadOwned, Received, SendFile, + SpliceRead, SpliceWrite, Split, Write, }; use rustls::{ProtocolVersion, SupportedCipherSuite}; @@ -86,6 +86,12 @@ pub trait KernelRead: Sized { #[cfg(all(target_os = "linux", feature = "kernel"))] fn adopt(reader: crate::kernel::linux::Reader) -> Self; + /// [`ReadAhead::stop_read_ahead`](http_kit_io::ReadAhead::stop_read_ahead) + /// on the socket's reading half before the kernel takes it; reached only + /// where [`SOCKET`](Self::SOCKET). + #[cfg(all(target_os = "linux", feature = "kernel"))] + fn stop_read_ahead(reader: &mut R) -> impl Future>; + fn get_ref(&self) -> &R; fn read(&mut self, buf: B) -> impl Future>; @@ -141,6 +147,11 @@ impl KernelRead for Never { unreachable!("only a socket's reading half is adopted by the kernel") } + #[cfg(all(target_os = "linux", feature = "kernel"))] + async fn stop_read_ahead(_reader: &mut R) -> Vec { + unreachable!("only a socket's reading half is handed to the kernel") + } + fn get_ref(&self) -> &R { match *self {} } @@ -247,7 +258,9 @@ impl TransportWrite for W { } #[cfg(all(target_os = "linux", feature = "kernel"))] -impl TransportRead for R { +impl TransportRead + for R +{ type Kernel = crate::kernel::linux::Reader; } @@ -488,6 +501,16 @@ impl Read for TlsReadHalf { } } +/// Plaintext is decrypted into the reader's buffer: there is none of the +/// stream's own to lend. +impl ReadOwned for TlsReadHalf { + type Buf = NoBuf; + + async fn read_owned(&mut self) -> io::Result> { + Ok(Received::ReadInstead) + } +} + impl Write for TlsWriteHalf { async fn write(&mut self, buf: B) -> BufResult { match &mut self.0 { @@ -616,6 +639,14 @@ impl Read for TlsStream { } } +impl ReadOwned for TlsStream { + type Buf = NoBuf; + + fn read_owned(&mut self) -> impl Future>> { + self.read.read_owned() + } +} + impl Write for TlsStream { fn write( &mut self,