feat(kio): add a poll-native Deadline and adopt it in moq-net

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"] }

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)]

rs/kio/src/time.rs

@@ -0,0 +1,224 @@
+//! 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};
+
+/// 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;
+
+/// 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<kio::time::Instant>) {
+/// let mut deadline = kio::time::Deadline::new();
+/// deadline.set(next_expiry);
+/// kio::wait(|waiter| deadline.poll(waiter)).await;
+/// # }
+/// ```
+pub struct Deadline {
+	at: Option<Instant>,
+
+	// 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<Pin<Box<web_async::time::Sleep>>>,
+}
+
+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.
+	///
+	/// 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().checked_add(duration),
+			sleep: None,
+		}
+	}
+
+	/// 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<Instant>) {
+		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<Instant> {
+		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 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();
+		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");
+	}
+}

rs/kio/src/tokio.rs

@@ -13,10 +13,13 @@ 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")]
+#[doc(hidden)]
 pub struct Sleep {
 	inner: Pin<Box<::tokio::time::Sleep>>,
 }
 
+#[allow(deprecated)]
 impl Sleep {
 	/// Wrap a tokio sleep future.
 	pub fn new(sleep: ::tokio::time::Sleep) -> Self {
@@ -35,6 +38,7 @@ impl Sleep {
 }
 
 #[cfg(test)]
+#[allow(deprecated)]
 mod tests {
 	use super::*;
 	use ::tokio::time::Instant;

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 }

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<RequestId, Error> {
-		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));
 			}

rs/moq-net/src/ietf/subscriber.rs

@@ -90,7 +90,7 @@ pub(super) struct Subscriber<S: web_transport_trait::Session> {
 }
 
 async fn resolve_track_alias(aliases: kio::Consumer<HashMap<u64, RequestId>>, alias: u64) -> Result<RequestId, Error> {
-	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<HashMap<u64, RequestId>>, 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

rs/moq-net/src/lite/publisher.rs

@@ -136,12 +136,12 @@ impl<S: web_transport_trait::Session> Publisher<S> {
 			// 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,37 +377,28 @@ impl<S: web_transport_trait::Session> Publisher<S> {
 
 		// 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<web_async::time::Instant> = None;
 				kio::wait(|waiter| {
 					if let Poll::Ready(res) = waiter.poll_future(closed.as_mut()) {
 						return Poll::Ready(Err(res));
 					}
 					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() {