From aa22156a03bc20e8d07353f6a18d50209892ce63 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sun, 2 Aug 2026 17:40:13 -0700 Subject: [PATCH 1/3] refactor(quiche): hand the keep-alive timer a waiter like the rest of ez The last poll method in `ez` still taking a `Context` was the keep-alive ticker's, so the driver built one from its waiter to call it. Take the waiter instead and build the `Context` at the one place that genuinely needs it, where tokio's `Interval::poll_tick` demands it. No behaviour change: the `Context` is made from the same waker, one level down. Everything else in `ez` already takes `&kio::Waiter`; the only `Context`s left are the `AsyncRead`/`AsyncWrite` impls, whose signatures are tokio's. Co-Authored-By: Claude Opus 5 --- rs/web-transport-quiche/src/ez/driver.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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; } From b0ad2305ee5b3aafd2aacea7b9d6181fcc0920a1 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Mon, 3 Aug 2026 11:31:16 -0700 Subject: [PATCH 2/3] refactor: build the poll bridge on kio::Park, now that it has shipped MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit kio 0.5.3 released `Park` (the `WaiterCell` of moq-dev/moq#2560, renamed), so the hand-rolled retention in `Parked` goes away: it now wraps `Park` and inherits the reuse `Park` does — a steady-state cell allocates nothing, where this one built a fresh `Waiter` and `Arc` on every poll. What stays is the release-on-`Ready`. `Park` holds its waiter until the next poll or until it drops, which is right for a pending operation but not a finished one: the stream is done with that caller, and holding its waker pins the polling task's allocation until something polls the cell again. Neutering `settle` — that is, using bare `Park` — fails `a_finished_read_releases_the_poller`, so the wrapper earns its keep. 0.5.3 also made `Waiter` `Clone`, which retires the two-step dance in the `AsyncRead`/`AsyncWrite` impls. They took that shape only because the cell's borrow could not be ended before the `&mut self` poll; a cloned waiter shares the parked one's identity, so `hold` hands one out and the borrow ends there. Co-Authored-By: Claude Opus 5 --- rs/web-transport-iroh/Cargo.toml | 2 +- rs/web-transport-noq/Cargo.toml | 2 +- rs/web-transport-quiche/Cargo.toml | 2 +- rs/web-transport-quiche/src/ez/recv.rs | 4 +- rs/web-transport-quiche/src/ez/send.rs | 12 ++-- rs/web-transport-quiche/src/waiters.rs | 80 +++++++++++--------------- rs/web-transport-quinn/Cargo.toml | 2 +- rs/web-transport-quinn/src/waiters.rs | 80 +++++++++++--------------- 8 files changed, 76 insertions(+), 108 deletions(-) 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/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..e19200f 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,49 @@ 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. -/// -/// 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 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(); + } } } 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..f777875 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,49 @@ 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. -/// -/// 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 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(); + } } } From 039999fab51d5b4047d0e6f8545c5ed6c5268f9b Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Mon, 3 Aug 2026 20:35:21 -0700 Subject: [PATCH 3/3] test: verify parked clones start idle --- rs/web-transport-quiche/src/waiters.rs | 32 ++++++++++++++++++++++++-- rs/web-transport-quinn/src/waiters.rs | 32 ++++++++++++++++++++++++-- 2 files changed, 60 insertions(+), 4 deletions(-) diff --git a/rs/web-transport-quiche/src/waiters.rs b/rs/web-transport-quiche/src/waiters.rs index e19200f..73a0d39 100644 --- a/rs/web-transport-quiche/src/waiters.rs +++ b/rs/web-transport-quiche/src/waiters.rs @@ -138,6 +138,9 @@ impl Wake for AcceptWaiters { /// stream whose last read completed and then sat idle, that is the rest of the /// connection. So this releases it on `Ready`. /// +/// `Park` also clones empty, so a cloned handle never inherits another handle's live +/// registration. +/// /// 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)] @@ -160,8 +163,8 @@ impl Parked { /// 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 clone shares the parked waiter's identity, so nothing is lost by taking one, - /// and it ends the borrow on the cell. + /// 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() } @@ -206,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/src/waiters.rs b/rs/web-transport-quinn/src/waiters.rs index f777875..2b72fff 100644 --- a/rs/web-transport-quinn/src/waiters.rs +++ b/rs/web-transport-quinn/src/waiters.rs @@ -136,6 +136,9 @@ impl Wake for AcceptWaiters { /// stream whose last read completed and then sat idle, that is the rest of the /// connection. So this releases it on `Ready`. /// +/// `Park` also clones empty, so a cloned handle never inherits another handle's live +/// registration. +/// /// 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)] @@ -158,8 +161,8 @@ impl Parked { /// 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 clone shares the parked waiter's identity, so nothing is lost by taking one, - /// and it ends the borrow on the cell. + /// 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() } @@ -204,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.