diff --git a/rs/web-transport-iroh/Cargo.toml b/rs/web-transport-iroh/Cargo.toml index 555fa09..b5aba5b 100644 --- a/rs/web-transport-iroh/Cargo.toml +++ b/rs/web-transport-iroh/Cargo.toml @@ -22,7 +22,7 @@ qlog = ["iroh/qlog"] bytes = "1" http = "1" iroh = { version = "1", default-features = false, features = ["fast-apple-datapath"] } -kio = "0.5.2" +kio = "0.5.3" n0-error = "1" n0-future = "0.3.1" tokio = { version = "1", default-features = false, features = [ diff --git a/rs/web-transport-noq/Cargo.toml b/rs/web-transport-noq/Cargo.toml index 579e3d5..bb6a664 100644 --- a/rs/web-transport-noq/Cargo.toml +++ b/rs/web-transport-noq/Cargo.toml @@ -26,7 +26,7 @@ qlog = ["noq/qlog"] bytes = "1" futures = "0.3" http = "1" -kio = "0.5.2" +kio = "0.5.3" noq = { version = "1", default-features = false, features = [ "tracing-log", "platform-verifier", diff --git a/rs/web-transport-quiche/Cargo.toml b/rs/web-transport-quiche/Cargo.toml index d0a8027..fefd734 100644 --- a/rs/web-transport-quiche/Cargo.toml +++ b/rs/web-transport-quiche/Cargo.toml @@ -26,7 +26,7 @@ bytes = "1" flume = "0.12" futures = "0.3" http = "1" -kio = "0.5.2" +kio = "0.5.3" rustls-pki-types = "1" thiserror = "2" diff --git a/rs/web-transport-quiche/src/ez/driver.rs b/rs/web-transport-quiche/src/ez/driver.rs index e245d04..0f85a2b 100644 --- a/rs/web-transport-quiche/src/ez/driver.rs +++ b/rs/web-transport-quiche/src/ez/driver.rs @@ -371,7 +371,7 @@ impl KeepAlive { } /// Returns true when a keep-alive is due. - fn poll(&mut self, cx: &mut Context) -> bool { + fn poll(&mut self, waiter: &Waiter) -> bool { let period = self.period; let ticker = self.ticker.get_or_insert_with(|| { // The first tick is one period out; `interval` would instead fire @@ -384,7 +384,10 @@ impl KeepAlive { ticker }); - ticker.poll_tick(cx).is_ready() + // `poll_tick` wants a `Context`; everything else in `ez` is handed a waiter. + ticker + .poll_tick(&mut Context::from_waker(waiter.waker())) + .is_ready() } } @@ -636,7 +639,7 @@ impl Driver { // ack-eliciting, so a tick on a busy connection costs nothing. let mut keep_alive = false; if let Some(k) = self.keep_alive.as_mut() { - if k.poll(&mut Context::from_waker(waiter.waker())) { + if k.poll(waiter) { qconn.send_ack_eliciting()?; keep_alive = true; } diff --git a/rs/web-transport-quiche/src/ez/recv.rs b/rs/web-transport-quiche/src/ez/recv.rs index 854a43d..7b84159 100644 --- a/rs/web-transport-quiche/src/ez/recv.rs +++ b/rs/web-transport-quiche/src/ez/recv.rs @@ -403,9 +403,9 @@ impl AsyncRead for RecvStream { cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.parked.hold(cx); let res = self.poll_read_chunk(&waiter, buf.remaining()); - self.parked.park(waiter, &res); + self.parked.settle(&res); match ready!(res) { Ok(Some(chunk)) => buf.put_slice(&chunk), diff --git a/rs/web-transport-quiche/src/ez/send.rs b/rs/web-transport-quiche/src/ez/send.rs index 5eb8fe5..5dfa05b 100644 --- a/rs/web-transport-quiche/src/ez/send.rs +++ b/rs/web-transport-quiche/src/ez/send.rs @@ -443,9 +443,9 @@ impl AsyncWrite for SendStream { buf: &[u8], ) -> Poll> { let mut buf = io::Cursor::new(buf); - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.parked.hold(cx); let res = self.poll_write_buf(&waiter, &mut buf); - self.parked.park(waiter, &res); + self.parked.settle(&res); match ready!(res) { Ok(n) => Poll::Ready(Ok(n)), @@ -454,9 +454,9 @@ impl AsyncWrite for SendStream { } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.parked.hold(cx); let res = self.poll_flushed(&waiter); - self.parked.park(waiter, &res); + self.parked.settle(&res); res.map_err(|e| io::Error::other(e.to_string())) } @@ -475,9 +475,9 @@ impl AsyncWrite for SendStream { Err(e) => return Poll::Ready(Err(io::Error::other(e.to_string()))), } - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.parked.hold(cx); let res = self.poll_closed(&waiter); - self.parked.park(waiter, &res); + self.parked.settle(&res); res.map_err(|e| io::Error::other(e.to_string())) } diff --git a/rs/web-transport-quiche/src/waiters.rs b/rs/web-transport-quiche/src/waiters.rs index 2e7f1e8..73a0d39 100644 --- a/rs/web-transport-quiche/src/waiters.rs +++ b/rs/web-transport-quiche/src/waiters.rs @@ -6,8 +6,8 @@ //! [`kio::WaiterList`] are for: a slot lives only as long as the [`kio::Waiter`] that //! made it. //! -//! [`Parked`] is the caller's end of such a registration, for a caller that has only a -//! `Context` to work with. [`AcceptWaiters`] is the other end for +//! [`Parked`] is the caller's end of such a registration, a [`kio::Park`] that also lets +//! go of its waiter once the poll finishes. [`AcceptWaiters`] is the other end for //! [`SessionAccept`](crate::SessionAccept), which is shared by every clone of a session //! and has to fan one arrival out to every accepter parked on it. @@ -16,7 +16,7 @@ use std::{ task::{Context, Poll, Wake}, }; -use kio::{Waiter, WaiterList}; +use kio::{Park, Waiter, WaiterList}; #[derive(Default)] struct AcceptState { @@ -128,65 +128,52 @@ impl Wake for AcceptWaiters { } } -/// A [`Waiter`] retained across `poll` calls. +/// A [`kio::Park`] that lets go of its waiter once the poll finishes. /// -/// A registration in [`AcceptWaiters`] stays live only while the caller holds the -/// [`Waiter`] it registered, which is what lets a caller that walks away release its -/// slot. So the handle has to live somewhere the *caller* owns, and a `poll_*` method is -/// handed nothing but a `Context`. This is that somewhere: one cell per operation, held -/// by whoever polls. +/// `Park` retains the waiter until the next poll or until it drops, which is what a +/// pending operation needs: a registration in [`AcceptWaiters`] (or any +/// [`kio::WaiterList`]) lives only as long as the [`Waiter`] that made it. A *finished* +/// poll is the other case. The stream is done with that caller, and going on holding its +/// waker pins the polling task's allocation until something polls the cell again — for a +/// stream whose last read completed and then sat idle, that is the rest of the +/// connection. So this releases it on `Ready`. /// -/// The `async` methods have no need for it — [`kio::wait`] keeps the waiter inside the -/// future it builds, so dropping the future drops the registration. +/// `Park` also clones empty, so a cloned handle never inherits another handle's live +/// registration. /// -/// kio grew a `WaiterCell` for exactly this after 0.5.2 (moq-dev/moq#2560); this can go -/// once that releases. Its `hold` also *reuses* the waiter when the task is unchanged and -/// every registration was already drained, which saves the allocation this one makes on -/// each poll — though it holds the waiter across a `Ready`, so keep retiring it here. -#[derive(Default)] -pub(crate) struct Parked { - waiter: Option, -} +/// The `async` methods have no need for any of it — [`kio::wait`] keeps the waiter inside +/// the future it builds, so dropping the future drops the registration. +#[derive(Clone, Default)] +pub(crate) struct Parked(Park); impl Parked { - /// Run one poll with a retained waiter. - /// - /// The waiter is kept only while the poll is `Pending`. On `Ready` there is nothing - /// left to wake, and holding the waker would pin the polling task's allocation until - /// something polled this cell again — for a stream that finished its last read and - /// then sat idle, that is the rest of the connection. + /// Run one poll with a retained waiter, releasing it if the poll finishes. pub(crate) fn poll( &mut self, - cx: &mut Context<'_>, + cx: &Context<'_>, poll: impl FnOnce(&Waiter) -> Poll, ) -> Poll { - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.hold(cx); let result = poll(&waiter); - self.park(waiter, &result); + self.settle(&result); result } - /// The two-step form of [`poll`](Self::poll), for a poll that needs `&mut self` of - /// the struct holding this cell while the waiter is alive — the borrow checker allows - /// only one of those at a time, so the closure form will not compile there. Build the - /// waiter from the `Context`, poll with it, then hand it here with the result. - pub(crate) fn park(&mut self, waiter: Waiter, result: &Poll) { - // Retiring the previous waiter *after* the poll matters: the new one is already - // registered by then, so there is no window with nothing registered. - self.waiter = result.is_pending().then_some(waiter); - } -} - -impl Clone for Parked { - /// A clone starts unregistered: a registration belongs to the handle that parked it. - fn clone(&self) -> Self { - Self::default() + /// Hold a waiter for this poll, for a body that needs `&mut self` of the struct + /// holding this cell — the borrow checker allows only one of those at a time, so the + /// closure form above will not compile there. Pair it with [`settle`](Self::settle). + /// + /// The returned clone shares the parked waiter's identity, so nothing is lost by + /// taking one, and it ends the borrow on the cell. + pub(crate) fn hold(&mut self, cx: &Context<'_>) -> Waiter { + self.0.hold(cx).clone() } -} -impl std::fmt::Debug for Parked { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.debug_struct("Parked").finish_non_exhaustive() + /// Release the held waiter if the poll finished. + pub(crate) fn settle(&mut self, result: &Poll) { + if result.is_ready() { + self.0 = Park::default(); + } } } @@ -222,6 +209,31 @@ mod tests { } } + #[test] + fn cloning_a_pending_handle_does_not_retain_the_poller() { + let waiters = AcceptWaiters::default(); + let flag = Arc::new(Flag::default()); + let waker = std::task::Waker::from(flag.clone()); + let mut parked = Parked::default(); + + assert!(parked + .poll(&Context::from_waker(&waker), |waiter| { + waiters.register(waiter); + Poll::<()>::Pending + }) + .is_pending()); + + let cloned = parked.clone(); + drop(parked); + waiters.wake_all(); + + assert!( + !flag.woken(), + "the clone retained the original handle's parked poller" + ); + drop(cloned); + } + /// A waker that re-enters the list from `wake`, standing in for an executor that /// polls a resumed task inline: the first thing a resumed accepter does is register /// itself again. diff --git a/rs/web-transport-quinn/Cargo.toml b/rs/web-transport-quinn/Cargo.toml index 0532393..c645e42 100644 --- a/rs/web-transport-quinn/Cargo.toml +++ b/rs/web-transport-quinn/Cargo.toml @@ -26,7 +26,7 @@ qlog = ["quinn/qlog"] bytes = "1" futures = "0.3" http = "1" -kio = "0.5.2" +kio = "0.5.3" quinn = { version = "0.11", default-features = false, features = [ "platform-verifier", diff --git a/rs/web-transport-quinn/src/waiters.rs b/rs/web-transport-quinn/src/waiters.rs index 315e49d..2b72fff 100644 --- a/rs/web-transport-quinn/src/waiters.rs +++ b/rs/web-transport-quinn/src/waiters.rs @@ -6,15 +6,15 @@ //! [`kio::WaiterList`] are for. //! //! [`AcceptWaiters`] is the list, plus the waker the shared accept futures are polled -//! with. [`Parked`] is the caller's end of a registration in that list, for a caller -//! that only has a `Context` to work with. +//! with. [`Parked`] is the caller's end of a registration, a [`kio::Park`] that also +//! lets go of its waiter once the poll finishes. use std::{ sync::{Arc, Mutex}, task::{Context, Poll, Wake}, }; -use kio::{Waiter, WaiterList}; +use kio::{Park, Waiter, WaiterList}; #[derive(Default)] struct AcceptState { @@ -126,65 +126,52 @@ impl Wake for AcceptWaiters { } } -/// A [`Waiter`] retained across `poll` calls. +/// A [`kio::Park`] that lets go of its waiter once the poll finishes. /// -/// A registration in [`AcceptWaiters`] stays live only while the caller holds the -/// [`Waiter`] it registered, which is what lets a caller that walks away release its -/// slot. So the handle has to live somewhere the *caller* owns, and a `poll_*` method is -/// handed nothing but a `Context`. This is that somewhere: one cell per operation, held -/// by whoever polls. +/// `Park` retains the waiter until the next poll or until it drops, which is what a +/// pending operation needs: a registration in [`AcceptWaiters`] (or any +/// [`kio::WaiterList`]) lives only as long as the [`Waiter`] that made it. A *finished* +/// poll is the other case. The stream is done with that caller, and going on holding its +/// waker pins the polling task's allocation until something polls the cell again — for a +/// stream whose last read completed and then sat idle, that is the rest of the +/// connection. So this releases it on `Ready`. /// -/// The `async` methods have no need for it — [`kio::wait`] keeps the waiter inside the -/// future it builds, so dropping the future drops the registration. +/// `Park` also clones empty, so a cloned handle never inherits another handle's live +/// registration. /// -/// kio grew a `WaiterCell` for exactly this after 0.5.2 (moq-dev/moq#2560); this can go -/// once that releases. Its `hold` also *reuses* the waiter when the task is unchanged and -/// every registration was already drained, which saves the allocation this one makes on -/// each poll — though it holds the waiter across a `Ready`, so keep retiring it here. -#[derive(Default)] -pub(crate) struct Parked { - waiter: Option, -} +/// The `async` methods have no need for any of it — [`kio::wait`] keeps the waiter inside +/// the future it builds, so dropping the future drops the registration. +#[derive(Clone, Default)] +pub(crate) struct Parked(Park); impl Parked { - /// Run one poll with a retained waiter. - /// - /// The waiter is kept only while the poll is `Pending`. On `Ready` there is nothing - /// left to wake, and holding the waker would pin the polling task's allocation until - /// something polled this cell again — for a stream that finished its last read and - /// then sat idle, that is the rest of the connection. + /// Run one poll with a retained waiter, releasing it if the poll finishes. pub(crate) fn poll( &mut self, - cx: &mut Context<'_>, + cx: &Context<'_>, poll: impl FnOnce(&Waiter) -> Poll, ) -> Poll { - let waiter = Waiter::new(cx.waker().clone()); + let waiter = self.hold(cx); let result = poll(&waiter); - self.park(waiter, &result); + self.settle(&result); result } - /// The two-step form of [`poll`](Self::poll), for a poll that needs `&mut self` of - /// the struct holding this cell while the waiter is alive — the borrow checker allows - /// only one of those at a time, so the closure form will not compile there. Build the - /// waiter from the `Context`, poll with it, then hand it here with the result. - pub(crate) fn park(&mut self, waiter: Waiter, result: &Poll) { - // Retiring the previous waiter *after* the poll matters: the new one is already - // registered by then, so there is no window with nothing registered. - self.waiter = result.is_pending().then_some(waiter); - } -} - -impl Clone for Parked { - /// A clone starts unregistered: a registration belongs to the handle that parked it. - fn clone(&self) -> Self { - Self::default() + /// Hold a waiter for this poll, for a body that needs `&mut self` of the struct + /// holding this cell — the borrow checker allows only one of those at a time, so the + /// closure form above will not compile there. Pair it with [`settle`](Self::settle). + /// + /// The returned clone shares the parked waiter's identity, so nothing is lost by + /// taking one, and it ends the borrow on the cell. + pub(crate) fn hold(&mut self, cx: &Context<'_>) -> Waiter { + self.0.hold(cx).clone() } -} -impl std::fmt::Debug for Parked { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.debug_struct("Parked").finish_non_exhaustive() + /// Release the held waiter if the poll finished. + pub(crate) fn settle(&mut self, result: &Poll) { + if result.is_ready() { + self.0 = Park::default(); + } } } @@ -220,6 +207,31 @@ mod tests { } } + #[test] + fn cloning_a_pending_handle_does_not_retain_the_poller() { + let waiters = AcceptWaiters::default(); + let flag = Arc::new(Flag::default()); + let waker = std::task::Waker::from(flag.clone()); + let mut parked = Parked::default(); + + assert!(parked + .poll(&Context::from_waker(&waker), |waiter| { + waiters.register(waiter); + Poll::<()>::Pending + }) + .is_pending()); + + let cloned = parked.clone(); + drop(parked); + waiters.wake_all(); + + assert!( + !flag.woken(), + "the clone retained the original handle's parked poller" + ); + drop(cloned); + } + /// A waker that re-enters the list from `wake`, standing in for an executor that /// polls a resumed task inline: the first thing a resumed accepter does is register /// itself again.