refactor: finish the ez waiter conversion, and build the poll bridge on kio::Park

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 = [

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",

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"

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;
             }

rs/web-transport-quiche/src/ez/recv.rs

@@ -403,9 +403,9 @@ impl AsyncRead for RecvStream {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<Result<(), io::Error>> {
-        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),

rs/web-transport-quiche/src/ez/send.rs

@@ -443,9 +443,9 @@ impl AsyncWrite for SendStream {
         buf: &[u8],
     ) -> Poll<Result<usize, io::Error>> {
         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<Result<(), io::Error>> {
-        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()))
     }

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<Waiter>,
-}
+/// 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<T>(
         &mut self,
-        cx: &mut Context<'_>,
+        cx: &Context<'_>,
         poll: impl FnOnce(&Waiter) -> Poll<T>,
     ) -> Poll<T> {
-        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<T>(&mut self, waiter: Waiter, result: &Poll<T>) {
-        // 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<T>(&mut self, result: &Poll<T>) {
+        if result.is_ready() {
+            self.0 = Park::default();
+        }
     }
 }
 

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",

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<Waiter>,
-}
+/// 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<T>(
         &mut self,
-        cx: &mut Context<'_>,
+        cx: &Context<'_>,
         poll: impl FnOnce(&Waiter) -> Poll<T>,
     ) -> Poll<T> {
-        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<T>(&mut self, waiter: Waiter, result: &Poll<T>) {
-        // 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<T>(&mut self, result: &Poll<T>) {
+        if result.is_ready() {
+            self.0 = Park::default();
+        }
     }
 }