From b2f891814943e274bfa01325ed8d9ec784976ba4 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Mon, 27 Jul 2026 13:07:36 -0700 Subject: [PATCH 1/2] feat(kio): add a poll-native Deadline and adopt it in moq-net Six sites in moq-net hand-rolled a wall-clock deadline inside a `kio::wait` closure, each pinning an async block whose `None` arm was `std::future::pending()`, then fusing it with a separate `fired` flag and an `is_some()` guard. Three moving parts to express "an optional deadline that fires once". `kio::time::Deadline` replaces all of it. Both clock backends behind `web_async::time` (tokio natively, wasmtimer in the browser) already expose `Sleep::is_elapsed` and `Sleep::reset`, so the fuse comes for free and re-arming reuses the allocation instead of boxing a fresh future every loop turn. Construction of the inner `Sleep` is deferred to the first poll, since on native it panics without a live tokio time driver, and only the poll is guaranteed to run inside the executor. The probe interval in lite/publisher no longer pins `interval.tick()` either; `Interval::poll_tick` is already poll-native on both backends. `kio::tokio::Sleep` is superseded: it is native-only, cannot be re-armed or disarmed, and has no fuse. It had no callers in or out of the module, and nothing in the workspace enabled the `tokio` feature. Hidden and deprecated here; removal is a breaking change and belongs on dev. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 1 + rs/kio/Cargo.toml | 6 +- rs/kio/src/lib.rs | 4 + rs/kio/src/time.rs | 209 ++++++++++++++++++++++++++++++ rs/kio/src/tokio.rs | 3 + rs/moq-net/Cargo.toml | 2 +- rs/moq-net/src/ietf/control.rs | 4 +- rs/moq-net/src/ietf/subscriber.rs | 4 +- rs/moq-net/src/lite/publisher.rs | 35 ++--- rs/moq-net/src/model/origin.rs | 49 +------ 10 files changed, 246 insertions(+), 71 deletions(-) create mode 100644 rs/kio/src/time.rs diff --git a/Cargo.lock b/Cargo.lock index 4e04b3b5dd..cf0ff1d4a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3859,6 +3859,7 @@ version = "0.5.0" dependencies = [ "smallvec", "tokio", + "web-async", ] [[package]] diff --git a/rs/kio/Cargo.toml b/rs/kio/Cargo.toml index 9af0543701..1615267bfc 100644 --- a/rs/kio/Cargo.toml +++ b/rs/kio/Cargo.toml @@ -13,13 +13,15 @@ keywords = ["async", "producer", "consumer", "state", "sync"] categories = ["asynchronous", "concurrency"] [features] -# Opt-in `tokio::time::Sleep`, a poll-driven wall-clock wait backed by tokio. Off by -# default so kio stays runtime-free. Enable it when driving `poll_*` functions on tokio. +# Opt-in poll-driven wall-clock deadlines, backed by `web-async` (tokio on native, +# wasmtimer in the browser). Off by default so kio stays runtime-free. +time = ["dep:web-async"] tokio = ["dep:tokio"] [dependencies] smallvec = "1.15" tokio = { workspace = true, features = ["time"], optional = true } +web-async = { workspace = true, optional = true } [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt", "test-util", "time"] } diff --git a/rs/kio/src/lib.rs b/rs/kio/src/lib.rs index fd0f2f5c8e..da2d073979 100644 --- a/rs/kio/src/lib.rs +++ b/rs/kio/src/lib.rs @@ -24,7 +24,11 @@ mod producer; mod shared; mod weak; +#[cfg(feature = "time")] +pub mod time; + #[cfg(feature = "tokio")] +#[doc(hidden)] pub mod tokio; #[cfg(test)] diff --git a/rs/kio/src/time.rs b/rs/kio/src/time.rs new file mode 100644 index 0000000000..e10e20dc9b --- /dev/null +++ b/rs/kio/src/time.rs @@ -0,0 +1,209 @@ +//! Poll-driven wall-clock deadlines. +//! +//! Behind the `time` feature. Built on [`web_async::time`], which is `tokio::time` on +//! native and `wasmtimer` in the browser, so the rest of kio stays runtime-free. +//! +//! On native, a timer must first be polled inside a tokio runtime with the time driver +//! enabled, or tokio panics. Reading the clock ([`Instant::now`]) has no such +//! requirement. Since the clock is tokio's, `tokio::time::pause()` advances these +//! deadlines in tests. + +use std::{pin::Pin, task::Poll}; + +pub use web_async::time::{Duration, Instant}; + +use crate::Waiter; + +/// A wall-clock deadline driven by kio's poll model. +/// +/// Arm it with an [`Instant`], poll it from a `poll_*` function, and re-arm or disarm it +/// as the deadline moves. A disarmed deadline never fires, and an elapsed one stays +/// ready until it is armed for a different instant. +/// +/// ```no_run +/// # async fn example(next_expiry: Option) { +/// let mut deadline = kio::time::Deadline::new(); +/// deadline.set(next_expiry); +/// kio::wait(|waiter| deadline.poll(waiter)).await; +/// # } +/// ``` +pub struct Deadline { + at: Option, + + // Allocated on the first poll after arming, then re-armed in place via `Sleep::reset`. + // Construction is deferred because on native it panics without a live tokio time + // driver, and only the poll is guaranteed to run inside the executor. + sleep: Option>>, +} + +impl Deadline { + /// A disarmed deadline, which never fires until [`set`](Self::set) arms it. + pub fn new() -> Self { + Self { at: None, sleep: None } + } + + /// A deadline armed for `at`. + pub fn at(at: Instant) -> Self { + Self { + at: Some(at), + sleep: None, + } + } + + /// A deadline armed for `duration` from now. + pub fn after(duration: Duration) -> Self { + Self::at(Instant::now() + duration) + } + + /// Arm, re-arm, or disarm (`None`) the deadline. + /// + /// Setting the instant it already holds does nothing, so a poll loop can recompute + /// its deadline every turn without restarting the countdown. + pub fn set(&mut self, at: Option) { + if self.at == at { + return; + } + self.at = at; + + // Reuse the allocation when there is one; `reset` also clears `is_elapsed`. + if let (Some(at), Some(sleep)) = (at, &mut self.sleep) { + sleep.as_mut().reset(at); + } + } + + /// The instant this fires at, or `None` while disarmed. + pub fn deadline(&self) -> Option { + self.at + } + + /// Poll the deadline, registering `waiter` so the poll re-fires once it elapses. + /// + /// `Ready` once the instant has passed, `Pending` before then and while disarmed. + pub fn poll(&mut self, waiter: &Waiter) -> Poll<()> { + // Disarmed: register nothing. Only `set` can arm it, and the caller driving this + // poll is the one that calls `set`. + let Some(at) = self.at else { return Poll::Pending }; + + let sleep = self + .sleep + .get_or_insert_with(|| Box::pin(web_async::time::sleep_until(at))); + + // Fused, so a caller that keeps polling after the deadline keeps seeing `Ready` + // rather than re-polling a completed future. + if sleep.is_elapsed() { + return Poll::Ready(()); + } + + waiter.poll_future(sleep.as_mut()) + } + + /// Wait for the deadline to elapse. Parks forever while disarmed. + pub async fn wait(&mut self) { + crate::wait(|waiter| self.poll(waiter)).await + } +} + +impl Default for Deadline { + fn default() -> Self { + Self::new() + } +} + +impl std::fmt::Debug for Deadline { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Deadline").field("at", &self.at).finish() + } +} + +#[cfg(test)] +mod tests { + use std::task::Waker; + + use super::*; + + /// Poll once without parking, for asserting `Pending` without hanging the test. + fn poll_once(deadline: &mut Deadline) -> Poll<()> { + let waiter = Waiter::new(Waker::noop().clone()); + deadline.poll(&waiter) + } + + #[tokio::test(start_paused = true)] + async fn fires_at_its_deadline() { + let at = Instant::now() + Duration::from_secs(5); + let mut deadline = Deadline::at(at); + + deadline.wait().await; + assert!(Instant::now() >= at, "returned before the deadline"); + } + + #[tokio::test(start_paused = true)] + async fn disarmed_never_fires() { + let mut deadline = Deadline::new(); + assert!(poll_once(&mut deadline).is_pending()); + + // Auto-advance would fire any armed timer well inside this window. + let res = tokio::time::timeout(Duration::from_secs(60), deadline.wait()).await; + assert!(res.is_err(), "a disarmed deadline fired"); + } + + #[tokio::test(start_paused = true)] + async fn stays_ready_once_elapsed() { + let mut deadline = Deadline::after(Duration::from_secs(1)); + deadline.wait().await; + + // Re-polling a completed timer must keep reporting the deadline as passed. + assert!(poll_once(&mut deadline).is_ready()); + assert!(poll_once(&mut deadline).is_ready()); + } + + #[tokio::test(start_paused = true)] + async fn re_arming_to_the_same_instant_does_not_restart() { + let at = Instant::now() + Duration::from_secs(1); + let mut deadline = Deadline::at(at); + deadline.wait().await; + + // The instant really has passed, so an idempotent `set` must not rewind it. + deadline.set(Some(at)); + assert!(poll_once(&mut deadline).is_ready()); + } + + #[tokio::test(start_paused = true)] + async fn re_arming_later_defers_the_fire() { + let start = Instant::now(); + let mut deadline = Deadline::after(Duration::from_secs(1)); + + // Force the allocation so the re-arm goes through `Sleep::reset`. + assert!(poll_once(&mut deadline).is_pending()); + + let later = start + Duration::from_secs(10); + deadline.set(Some(later)); + deadline.wait().await; + + assert!(Instant::now() >= later, "fired at the original deadline"); + } + + #[tokio::test(start_paused = true)] + async fn disarming_a_live_countdown_stops_it() { + let mut deadline = Deadline::after(Duration::from_secs(1)); + assert!(poll_once(&mut deadline).is_pending()); + + deadline.set(None); + assert_eq!(deadline.deadline(), None); + + let res = tokio::time::timeout(Duration::from_secs(60), deadline.wait()).await; + assert!(res.is_err(), "a disarmed deadline fired"); + } + + #[tokio::test(start_paused = true)] + async fn re_arming_after_disarm_fires_again() { + let mut deadline = Deadline::after(Duration::from_secs(1)); + assert!(poll_once(&mut deadline).is_pending()); + deadline.set(None); + + let at = Instant::now() + Duration::from_secs(3); + deadline.set(Some(at)); + deadline.wait().await; + + assert!(Instant::now() >= at, "returned before the re-armed deadline"); + } +} diff --git a/rs/kio/src/tokio.rs b/rs/kio/src/tokio.rs index 21e701bb8b..6a34c715fc 100644 --- a/rs/kio/src/tokio.rs +++ b/rs/kio/src/tokio.rs @@ -13,10 +13,12 @@ use crate::Waiter; /// Construct it once for a deadline (`Sleep::new(tokio::time::sleep_until(deadline))`), then /// [`poll`](Self::poll) it against a [`Waiter`] each time your `poll_*` runs. Reading the /// clock through `tokio::time` means a `tokio::time::pause()` test advances it in step. +#[deprecated(note = "use kio::time::Deadline, which also works on wasm and can be re-armed in place")] pub struct Sleep { inner: Pin>, } +#[allow(deprecated)] impl Sleep { /// Wrap a tokio sleep future. pub fn new(sleep: ::tokio::time::Sleep) -> Self { @@ -35,6 +37,7 @@ impl Sleep { } #[cfg(test)] +#[allow(deprecated)] mod tests { use super::*; use ::tokio::time::Instant; diff --git a/rs/moq-net/Cargo.toml b/rs/moq-net/Cargo.toml index ed34e82ff5..c1be2db2c1 100644 --- a/rs/moq-net/Cargo.toml +++ b/rs/moq-net/Cargo.toml @@ -18,7 +18,7 @@ ignored = ["getrandom"] [dependencies] bytes = "1" futures = "0.3" -kio = { workspace = true } +kio = { workspace = true, features = ["time"] } num_enum = "0.7" rand = "0.10.1" serde = { workspace = true } diff --git a/rs/moq-net/src/ietf/control.rs b/rs/moq-net/src/ietf/control.rs index e6790cf035..91e136e6e2 100644 --- a/rs/moq-net/src/ietf/control.rs +++ b/rs/moq-net/src/ietf/control.rs @@ -30,7 +30,7 @@ impl Control { /// Allocate the next request_id, blocking until MAX_REQUEST_ID allows it. pub async fn next_request_id(&self) -> Result { - let mut timeout = std::pin::pin!(web_async::time::sleep(std::time::Duration::from_secs(10))); + let mut timeout = kio::time::Deadline::after(std::time::Duration::from_secs(10)); kio::wait(|waiter| { let allowed = self.state.poll(waiter, |state| { @@ -45,7 +45,7 @@ impl Control { return Poll::Ready(Ok(state.request_id_next.increment())); } - if waiter.poll_future(timeout.as_mut()).is_ready() { + if timeout.poll(waiter).is_ready() { tracing::warn!("timed out waiting for MAX_REQUEST_ID"); return Poll::Ready(Err(Error::Cancel)); } diff --git a/rs/moq-net/src/ietf/subscriber.rs b/rs/moq-net/src/ietf/subscriber.rs index b025e62bb0..61dc90e626 100644 --- a/rs/moq-net/src/ietf/subscriber.rs +++ b/rs/moq-net/src/ietf/subscriber.rs @@ -90,7 +90,7 @@ pub(super) struct Subscriber { } async fn resolve_track_alias(aliases: kio::Consumer>, alias: u64) -> Result { - let mut timeout = std::pin::pin!(web_async::time::sleep(TRACK_ALIAS_TIMEOUT)); + let mut timeout = kio::time::Deadline::after(TRACK_ALIAS_TIMEOUT); kio::wait(|waiter| { let resolved = aliases.poll(waiter, |aliases| match aliases.get(&alias) { Some(request_id) => Poll::Ready(*request_id), @@ -99,7 +99,7 @@ async fn resolve_track_alias(aliases: kio::Consumer>, al if let Poll::Ready(result) = resolved { return Poll::Ready(result.map_err(|_| Error::Dropped)); } - if waiter.poll_future(timeout.as_mut()).is_ready() { + if timeout.poll(waiter).is_ready() { return Poll::Ready(Err(Error::NotFound)); } Poll::Pending diff --git a/rs/moq-net/src/lite/publisher.rs b/rs/moq-net/src/lite/publisher.rs index 5e47911a9e..42c055eb14 100644 --- a/rs/moq-net/src/lite/publisher.rs +++ b/rs/moq-net/src/lite/publisher.rs @@ -136,12 +136,12 @@ impl Publisher { // Tick the probe interval, bailing as soon as the peer closes its side. let closed = { let mut closed = std::pin::pin!(stream.reader.closed()); - let mut tick = std::pin::pin!(interval.tick()); kio::wait(|waiter| { if let Poll::Ready(res) = waiter.poll_future(closed.as_mut()) { return Poll::Ready(Some(res)); } - waiter.poll_future(tick.as_mut()).map(|_| None) + let mut cx = std::task::Context::from_waker(waiter.waker()); + interval.poll_tick(&mut cx).map(|_| None) }) .await }; @@ -377,27 +377,18 @@ impl Publisher { // Send updates as they arrive. Closure wins the race so a dead peer can't // stall on a busy announce feed. + let mut linger = kio::time::Deadline::new(); loop { // The earliest deferred cost-restore, if any entry's linger is running. - let deadline = watched - .values() - .filter_map(|entry| entry.idle_at) - .min() - .map(|at| at + COST_LINGER); + linger.set( + watched + .values() + .filter_map(|entry| entry.idle_at) + .min() + .map(|at| at + COST_LINGER), + ); let op = { let mut closed = std::pin::pin!(stream.reader.closed()); - // Pending forever while no linger is running. Fused via `fired`: - // a completed future must not be polled again, and once it fires - // the turn always ends in a `Ready` below. - let mut linger = std::pin::pin!(async move { - match deadline { - Some(at) => { - web_async::time::sleep(at.saturating_duration_since(web_async::time::Instant::now())).await - } - None => std::future::pending().await, - } - }); - let mut fired: Option = None; kio::wait(|waiter| { if let Poll::Ready(res) = waiter.poll_future(closed.as_mut()) { return Poll::Ready(Err(res)); @@ -405,9 +396,9 @@ impl Publisher { if let Poll::Ready(next) = announced.poll_next(waiter) { return Poll::Ready(Ok(Op::Announce(next))); } - if fired.is_none() && waiter.poll_future(linger.as_mut()).is_ready() { - fired = Some(web_async::time::Instant::now()); - } + // Stamped per poll rather than kept: the turn always ends in a + // `Ready` below once it fires, so it never has to survive. + let fired = linger.poll(waiter).is_ready().then(web_async::time::Instant::now); // Poll every watched broadcast for a route change; each wake // rescans the map, which announce-control rates make fine. for (suffix, entry) in watched.iter_mut() { diff --git a/rs/moq-net/src/model/origin.rs b/rs/moq-net/src/model/origin.rs index 1568692d47..f6db925602 100644 --- a/rs/moq-net/src/model/origin.rs +++ b/rs/moq-net/src/model/origin.rs @@ -1413,33 +1413,22 @@ async fn run_front( // waiting for a replacement source. A graceful close never gets here (the // detach sets `closed` synchronously), so a running countdown always means a // reconnect is welcome. - let mut deadline: Option = None; + let mut deadline = kio::time::Deadline::new(); loop { let empty = { let s = state.read(); !s.closed && s.routes.is_empty() }; - deadline = match (empty, deadline) { + deadline.set(match (empty, deadline.deadline()) { // An unrepresentable deadline (e.g. `Duration::MAX`) lingers forever: // no timer, only a re-attach or teardown moves the front on. (true, None) => web_async::time::Instant::now().checked_add(linger), (true, at) => at, (false, _) => None, - }; + }); let step = { - // Pending forever while no countdown is running. Fused via `fired`: a - // completed future must not be polled again. - let mut sleep = std::pin::pin!(async { - match deadline { - Some(at) => { - web_async::time::sleep(at.saturating_duration_since(web_async::time::Instant::now())).await - } - None => std::future::pending().await, - } - }); - let mut fired = false; kio::wait(|waiter| { if let Poll::Ready((name, resume)) = broadcast.poll_spliced_assigned(waiter) { return Poll::Ready(Step::Serve(name, resume)); @@ -1459,13 +1448,7 @@ async fn run_front( Poll::Ready(Err(_)) => return Poll::Ready(Step::Closed), Poll::Pending => {} } - if deadline.is_some() && !fired && waiter.poll_future(sleep.as_mut()).is_ready() { - fired = true; - } - match fired { - true => Poll::Ready(Step::Expired), - false => Poll::Pending, - } + deadline.poll(waiter).map(|_| Step::Expired) }) .await }; @@ -1537,6 +1520,7 @@ async fn serve_track(state: kio::Producer, name: Arc, mut resum let mut dead: Option = None; // When the spliced segment stopped being read, starting the release countdown. let mut idle_since: Option = None; + let mut deadline = kio::time::Deadline::new(); loop { let serving_id = serving.as_ref().map(|(id, _)| *id); @@ -1549,22 +1533,9 @@ async fn serve_track(state: kio::Producer, name: Arc, mut resum (true, false) => idle_since.or_else(|| Some(web_async::time::Instant::now())), _ => None, }; - let deadline = idle_since.and_then(|at| at.checked_add(TRACK_IDLE_LINGER)); + deadline.set(idle_since.and_then(|at| at.checked_add(TRACK_IDLE_LINGER))); let step = { - // Pinned outside the closure: a future re-created on every poll would - // restart the countdown each time and never fire. Fused via `fired`: a - // completed future must not be polled again. - let mut sleep = std::pin::pin!(async { - match deadline { - Some(at) => { - web_async::time::sleep(at.saturating_duration_since(web_async::time::Instant::now())).await - } - None => std::future::pending().await, - } - }); - let mut fired = false; - kio::wait(|waiter| { // Watch the source table: the front closing, or the active source // moving away from the one currently spliced in (skipping one whose @@ -1620,13 +1591,7 @@ async fn serve_track(state: kio::Producer, name: Arc, mut resum }); } - if deadline.is_some() && !fired && waiter.poll_future(sleep.as_mut()).is_ready() { - fired = true; - } - if fired { - return Poll::Ready(Step::Idle); - } - Poll::Pending + deadline.poll(waiter).map(|_| Step::Idle) }) .await }; From 5cdbaf01d5c30cccbb57d7ce25dd750450109739 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Mon, 27 Jul 2026 13:39:40 -0700 Subject: [PATCH 2/2] fix(kio): address review feedback on Deadline - Document that `Duration`/`Instant` are re-exported from `web-async`, so a major bump there is a breaking change for them. - `Deadline::after` now uses checked arithmetic. A duration the clock cannot represent (e.g. `Duration::MAX`) leaves the deadline disarmed instead of panicking on the overflow, matching how `run_front` already treats an unrepresentable linger. - Add the item-level `#[doc(hidden)]` to the deprecated `kio::tokio::Sleep` alongside its `#[deprecated]`, per the repo's deprecation convention. Co-Authored-By: Claude Opus 5 --- rs/kio/src/time.rs | 17 ++++++++++++++++- rs/kio/src/tokio.rs | 1 + 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/rs/kio/src/time.rs b/rs/kio/src/time.rs index e10e20dc9b..e81a6776d7 100644 --- a/rs/kio/src/time.rs +++ b/rs/kio/src/time.rs @@ -10,6 +10,8 @@ use std::{pin::Pin, task::Poll}; +/// Re-exported from `web-async`, so a major bump of that crate is a breaking change +/// for these types. pub use web_async::time::{Duration, Instant}; use crate::Waiter; @@ -51,8 +53,14 @@ impl Deadline { } /// A deadline armed for `duration` from now. + /// + /// A duration the clock cannot represent (e.g. [`Duration::MAX`]) leaves the deadline + /// disarmed, so it never fires rather than panicking on the overflow. pub fn after(duration: Duration) -> Self { - Self::at(Instant::now() + duration) + Self { + at: Instant::now().checked_add(duration), + sleep: None, + } } /// Arm, re-arm, or disarm (`None`) the deadline. @@ -136,6 +144,13 @@ mod tests { assert!(Instant::now() >= at, "returned before the deadline"); } + #[tokio::test(start_paused = true)] + async fn an_unrepresentable_duration_disarms_instead_of_panicking() { + let mut deadline = Deadline::after(Duration::MAX); + assert_eq!(deadline.deadline(), None); + assert!(poll_once(&mut deadline).is_pending()); + } + #[tokio::test(start_paused = true)] async fn disarmed_never_fires() { let mut deadline = Deadline::new(); diff --git a/rs/kio/src/tokio.rs b/rs/kio/src/tokio.rs index 6a34c715fc..972a1c0f72 100644 --- a/rs/kio/src/tokio.rs +++ b/rs/kio/src/tokio.rs @@ -14,6 +14,7 @@ use crate::Waiter; /// [`poll`](Self::poll) it against a [`Waiter`] each time your `poll_*` runs. Reading the /// clock through `tokio::time` means a `tokio::time::pause()` test advances it in step. #[deprecated(note = "use kio::time::Deadline, which also works on wasm and can be re-armed in place")] +#[doc(hidden)] pub struct Sleep { inner: Pin>, }