From 419293bbe85e4250877702a682f7cfddec3c483d Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Tue, 2 Jun 2026 17:55:08 -0700 Subject: [PATCH 1/3] moq-lite-05: implement FETCH for past groups via TrackConsumer::fetch Add a one-shot `TrackConsumer::fetch(group, options) -> GroupConsumer` that retrieves a single past group without holding a subscription. A cached group is returned directly; otherwise the request bridges to a wire moq-lite FETCH, blocking on a new FETCH_OK (mirroring SUBSCRIBE/SUBSCRIBE_OK). - model: a dynamic group-request channel parallel to dynamic track requests (request_fetch / FetchPending / BroadcastDynamic::requested_group / GroupRequest). Cache hit returns the group; an aborted cached group is bypassed. - wire: new FetchOk message; Fetch.frame_start (lite-05+); the publisher honors frame_start by skipping earlier frames. - publisher: recv_fetch / run_fetch serve the group's frames on the fetch stream. - subscriber: serve group requests by issuing FETCH upstream and routing frames into the producer that resolves the fetcher. - relay web.rs: /fetch?group=N uses fetch() instead of subscribe(). Co-Authored-By: Claude Opus 4.8 (1M context) --- rs/moq-native/tests/broadcast.rs | 115 +++++++ rs/moq-net/src/lite/fetch.rs | 149 ++++++++- rs/moq-net/src/lite/publisher.rs | 148 ++++++++- rs/moq-net/src/lite/subscriber.rs | 141 ++++++++- rs/moq-net/src/model/broadcast.rs | 509 +++++++++++++++++++++++++++++- rs/moq-net/src/model/group.rs | 10 + rs/moq-relay/src/web.rs | 30 +- 7 files changed, 1063 insertions(+), 39 deletions(-) diff --git a/rs/moq-native/tests/broadcast.rs b/rs/moq-native/tests/broadcast.rs index 8c13306c47..6d57ff8e3c 100644 --- a/rs/moq-native/tests/broadcast.rs +++ b/rs/moq-native/tests/broadcast.rs @@ -227,6 +227,121 @@ async fn broadcast_moq_lite_05_timestamps_webtransport() { lite05_timestamp_roundtrip("https").await; } +/// Lite05 FETCH round-trip: retrieve a past group by sequence without holding a +/// subscription, exercising the FETCH / FETCH_OK control flow and per-frame +/// timestamp decoding on the fetch stream. +async fn lite05_fetch_roundtrip(scheme: &str) { + use moq_native::moq_net::{Timescale, Timestamp}; + + let pub_origin = Origin::random().produce(); + let mut broadcast = pub_origin.create_broadcast("test").expect("failed to create broadcast"); + let mut track = broadcast + .create_track(Track::new("video").with_timescale(Timescale::MICRO)) + .expect("failed to create track"); + + // A group with a few timestamped frames (middle PTS goes backwards, so the + // fetch stream carries a negative zigzag delta too). + let timestamps_us = [10_000u64, 30_000, 20_000]; + let mut group = track.append_group().expect("failed to append group"); // seq 0 + for &us in ×tamps_us { + let payload = format!("frame@{us}").into_bytes(); + let frame = moq_native::moq_net::Frame { + size: payload.len() as u64, + timestamp: Some(Timestamp::new(us, Timescale::MICRO).unwrap()), + }; + let mut writer = group.create_frame(frame).expect("failed to create frame"); + writer + .write(bytes::Bytes::from(payload)) + .expect("failed to write frame"); + writer.finish().expect("failed to finish frame"); + } + group.finish().expect("failed to finish group"); + + let mut server_config = moq_native::ServerConfig::default(); + server_config.bind = Some("[::]:0".to_string()); + server_config.tls.generate = vec!["localhost".into()]; + server_config.version = vec!["moq-lite-05-wip".parse().unwrap()]; + let mut server = server_config.init().expect("failed to init server"); + let addr = server.local_addr().expect("failed to get local addr"); + + let sub_origin = Origin::random().produce(); + let mut announcements = sub_origin.consume().announced(); + + let mut client_config = moq_native::ClientConfig::default(); + client_config.tls.disable_verify = Some(true); + client_config.version = vec!["moq-lite-05-wip".parse().unwrap()]; + let client = client_config.init().expect("failed to init client"); + let url: url::Url = format!("{scheme}://localhost:{}", addr.port()).parse().unwrap(); + + let server_handle = tokio::spawn(async move { + let request = server.accept().await.expect("no incoming connection"); + let session = request.with_publisher(pub_origin.clone()).ok().await?; + let _broadcast = broadcast; + let _track = track; + let _ = session.closed().await; + Ok::<_, anyhow::Error>(()) + }); + + let client = client.with_consumer(sub_origin); + let session = tokio::time::timeout(TIMEOUT, client.connect(url)) + .await + .expect("client connect timed out") + .expect("client connect failed"); + + let (path, bc) = tokio::time::timeout(TIMEOUT, announcements.next()) + .await + .expect("announce timed out") + .expect("origin closed"); + assert_eq!(path.as_str(), "test"); + let bc = bc.broadcast().expect("expected announce, got unannounce"); + + // Fetch group 0 directly, without subscribing. No live producer holds the + // group on the client, so this issues a wire FETCH upstream. + let mut group_sub = tokio::time::timeout(TIMEOUT, bc.consume_track("video").fetch(0, None)) + .await + .expect("fetch timed out") + .expect("fetch failed"); + assert_eq!(group_sub.sequence, 0); + + for &expected_us in ×tamps_us { + let mut frame_sub = tokio::time::timeout(TIMEOUT, group_sub.next_frame()) + .await + .expect("next_frame timed out") + .expect("next_frame failed") + .expect("group closed prematurely"); + + let ts = frame_sub + .timestamp + .expect("Lite05 fetch must carry per-frame timestamps"); + assert_eq!(ts.scale(), Timescale::MICRO); + assert_eq!(ts.value(), expected_us); + + let payload = frame_sub.read_all().await.expect("failed to read frame"); + assert_eq!(payload, bytes::Bytes::from(format!("frame@{expected_us}"))); + } + + // The fetched group ends cleanly (stream FIN → no more frames). + let end = tokio::time::timeout(TIMEOUT, group_sub.next_frame()) + .await + .expect("next_frame timed out") + .expect("next_frame failed"); + assert!(end.is_none(), "group should finish after its frames"); + + drop(session); + server_handle + .await + .expect("server task panicked") + .expect("server task failed"); +} + +#[tracing_test::traced_test] +#[tokio::test] +async fn broadcast_moq_lite_05_fetch_webtransport() { + // WebTransport only: Lite05Wip isn't advertised over ALPN, so raw QUIC (moqt://) + // can't negotiate it (same reason the other Lite05 tests are https-only). + lite05_fetch_roundtrip("https").await; +} + /// On Lite05 a publisher that doesn't advertise a timescale still works: /// SUBSCRIBE_OK carries `timescale = 0` and neither side encodes a /// per-frame timestamp byte. Subscribers receive `frame.timestamp = None`. diff --git a/rs/moq-net/src/lite/fetch.rs b/rs/moq-net/src/lite/fetch.rs index 04281e65f9..3507383a12 100644 --- a/rs/moq-net/src/lite/fetch.rs +++ b/rs/moq-net/src/lite/fetch.rs @@ -1,7 +1,7 @@ use std::borrow::Cow; use crate::{ - Path, + Compression, Path, Timescale, coding::{Decode, DecodeError, Encode, EncodeError}, }; @@ -10,13 +10,16 @@ use super::{Message, Version}; /// Sent by the subscriber to fetch a specific group from a track. /// /// Lite03+ only. -#[allow(dead_code)] #[derive(Clone, Debug)] pub struct Fetch<'a> { pub broadcast: Path<'a>, pub track: Cow<'a, str>, pub priority: u8, pub group: u64, + /// 0-based index of the first frame to return. The publisher skips earlier + /// frames so a subscriber can resume partway through a partially-received + /// group; `0` returns the whole group. Lite05+ only; older drafts always 0. + pub frame_start: u64, } impl Message for Fetch<'_> { @@ -32,12 +35,17 @@ impl Message for Fetch<'_> { let track = Cow::::decode(r, version)?; let priority = u8::decode(r, version)?; let group = u64::decode(r, version)?; + let frame_start = match version { + Version::Lite03 | Version::Lite04 => 0, + _ => u64::decode(r, version)?, + }; Ok(Self { broadcast, track, priority, group, + frame_start, }) } @@ -53,6 +61,143 @@ impl Message for Fetch<'_> { self.track.encode(w, version)?; self.priority.encode(w, version)?; self.group.encode(w, version)?; + match version { + Version::Lite03 | Version::Lite04 => {} + _ => self.frame_start.encode(w, version)?, + } Ok(()) } } + +/// Publisher's response to a [`Fetch`], sent on the same bidirectional stream +/// immediately before the group's frames. +/// +/// Mirrors the codec/timescale negotiation in [`super::SubscribeOk`] so the +/// subscriber can decode the frames that follow. There is no error variant: a +/// failed fetch resets the stream instead. Lite05+ only. +#[derive(Clone, Debug)] +pub struct FetchOk { + /// Echo of the requested group sequence, for a sanity check. + pub group: u64, + /// Codec the publisher used for every frame in this group. + pub compression: Compression, + /// Per-frame timestamp scale, or `None` if the frames carry no timestamps. + /// On the wire `None` is `0` and `Some(n)` is `n`. + pub timescale: Option, +} + +impl Message for FetchOk { + fn decode_msg(r: &mut R, version: Version) -> Result { + match version { + Version::Lite01 | Version::Lite02 | Version::Lite03 | Version::Lite04 => { + return Err(DecodeError::Version); + } + _ => {} + } + + let group = u64::decode(r, version)?; + let timescale = Timescale::new(u64::decode(r, version)?).ok(); + let compression = Compression::from_code(u64::decode(r, version)?).map_err(|_| DecodeError::InvalidValue)?; + + Ok(Self { + group, + compression, + timescale, + }) + } + + fn encode_msg(&self, w: &mut W, version: Version) -> Result<(), EncodeError> { + match version { + Version::Lite01 | Version::Lite02 | Version::Lite03 | Version::Lite04 => { + return Err(EncodeError::Version); + } + _ => {} + } + + self.group.encode(w, version)?; + self.timescale.map(u64::from).unwrap_or(0).encode(w, version)?; + self.compression.to_code().encode(w, version)?; + Ok(()) + } +} + +#[cfg(test)] +mod test { + use super::*; + + fn sample() -> FetchOk { + FetchOk { + group: 42, + compression: Compression::Deflate, + timescale: Some(Timescale::MICRO), + } + } + + fn roundtrip(version: Version, ok: &FetchOk) -> FetchOk { + let mut buf = Vec::new(); + ok.encode_msg(&mut buf, version).unwrap(); + let mut slice = buf.as_slice(); + FetchOk::decode_msg(&mut slice, version).unwrap() + } + + #[test] + fn fetch_ok_roundtrips_on_lite05() { + let got = roundtrip(Version::Lite05Wip, &sample()); + assert_eq!(got.group, 42); + assert_eq!(got.compression, Compression::Deflate); + assert_eq!(got.timescale, Some(Timescale::MICRO)); + } + + #[test] + fn fetch_ok_timescale_none_roundtrips() { + let mut ok = sample(); + ok.timescale = None; + assert_eq!(roundtrip(Version::Lite05Wip, &ok).timescale, None); + } + + #[test] + fn fetch_ok_errors_before_lite05() { + let mut buf = Vec::new(); + assert!(sample().encode_msg(&mut buf, Version::Lite04).is_err()); + } + + fn fetch_sample() -> Fetch<'static> { + Fetch { + broadcast: Path::new("room").to_owned(), + track: Cow::Borrowed("video"), + priority: 3, + group: 7, + frame_start: 5, + } + } + + fn fetch_roundtrip(version: Version, msg: &Fetch<'_>) -> Fetch<'static> { + let mut buf = Vec::new(); + msg.encode_msg(&mut buf, version).unwrap(); + let mut slice = buf.as_slice(); + Fetch::decode_msg(&mut slice, version).unwrap() + } + + #[test] + fn fetch_frame_start_roundtrips_on_lite05() { + assert_eq!(fetch_roundtrip(Version::Lite05Wip, &fetch_sample()).frame_start, 5); + } + + #[test] + fn fetch_frame_start_absent_before_lite05() { + let msg = fetch_sample(); + + // The frame_start varint only exists on lite-05+, so the older encoding is + // strictly shorter and always decodes back as 0. + let mut buf04 = Vec::new(); + msg.encode_msg(&mut buf04, Version::Lite04).unwrap(); + let mut buf05 = Vec::new(); + msg.encode_msg(&mut buf05, Version::Lite05Wip).unwrap(); + assert!( + buf05.len() > buf04.len(), + "lite-05 carries the extra frame_start varint" + ); + + assert_eq!(fetch_roundtrip(Version::Lite04, &msg).frame_start, 0); + } +} diff --git a/rs/moq-net/src/lite/publisher.rs b/rs/moq-net/src/lite/publisher.rs index 547b15e120..09176e71e0 100644 --- a/rs/moq-net/src/lite/publisher.rs +++ b/rs/moq-net/src/lite/publisher.rs @@ -62,6 +62,7 @@ impl Publisher { if let Err(err) = match kind { lite::ControlType::Announce => self.recv_announce(stream).await, lite::ControlType::Subscribe => self.recv_subscribe(stream).await, + lite::ControlType::Fetch => self.recv_fetch(stream).await, lite::ControlType::Probe => { self.recv_probe(stream); Ok(()) @@ -70,7 +71,7 @@ impl Publisher { tracing::info!("received goaway stream"); Ok(()) } - lite::ControlType::Session | lite::ControlType::Fetch => Err(Error::UnexpectedStream), + lite::ControlType::Session => Err(Error::UnexpectedStream), } { tracing::warn!(%err, "control stream error"); } @@ -528,6 +529,151 @@ impl Publisher { stream.writer.finish()?; stream.writer.closed().await } + + pub async fn recv_fetch(&mut self, mut stream: Stream) -> Result<(), Error> { + // FETCH is lite-05+ only; older drafts have no FETCH_OK / frame format. + if !self.version.has_timestamps() { + return Err(Error::UnexpectedStream); + } + + let fetch = stream.reader.decode::().await?; + + let track = fetch.track.clone(); + let group = fetch.group; + let absolute = self.origin.absolute(&fetch.broadcast).to_owned(); + + tracing::info!(broadcast = %absolute, %track, %group, "fetch started"); + + // The peer fetched this exact path, so it has already seen an announcement + // for it; a synchronous lookup is appropriate (as in recv_subscribe). + let broadcast = self.origin.get_broadcast(&fetch.broadcast); + let version = self.version; + let track_stats = self.stats.broadcast(&absolute).publisher_track(&track); + + web_async::spawn(async move { + if let Err(err) = Self::run_fetch(&mut stream, &fetch, broadcast, track_stats, version).await { + match &err { + Error::Cancel | Error::Transport(_) => { + tracing::info!(broadcast = %absolute, %track, %group, "fetch cancelled") + } + err => tracing::warn!(broadcast = %absolute, %track, %group, %err, "fetch error"), + } + stream.writer.abort(&err); + } else { + tracing::info!(broadcast = %absolute, %track, %group, "fetch complete"); + } + }); + + Ok(()) + } + + async fn run_fetch( + stream: &mut Stream, + fetch: &lite::Fetch<'_>, + consumer: Option, + track_stats: crate::PublisherTrack, + version: Version, + ) -> Result<(), Error> { + let broadcast = consumer.ok_or(Error::NotFound)?; + + let group = broadcast + .consume_track(&fetch.track) + .fetch( + fetch.group, + crate::Fetch { + priority: fetch.priority, + }, + ) + .await?; + + // FETCH is gated to lite-05+, which always carries timestamps when the track + // advertised a timescale; `None` means an untimed track (frames omit them). + let timescale = if version.has_timestamps() { + group.timescale() + } else { + None + }; + // v1: fetched groups stream uncompressed (the model carries no per-track compress hint). + let compression = Compression::None; + + stream + .writer + .encode(&lite::FetchOk { + group: fetch.group, + compression, + timescale, + }) + .await?; + track_stats.group(); + + // Honor frame_start: skip earlier frames, then stream the rest in order. The + // delta-timestamp baseline resets to 0, so the first served frame's delta is + // its absolute timestamp (the subscriber decodes against the same baseline). + let mut index = fetch.frame_start as usize; + let mut prev_ts: u64 = 0; + while let Some(mut frame) = group.get_frame(index).await? { + write_fetch_frame( + &mut stream.writer, + &mut frame, + compression, + timescale, + &mut prev_ts, + &track_stats, + ) + .await?; + index += 1; + } + + stream.writer.finish()?; + stream.writer.closed().await + } +} + +/// Write one frame to a fetch stream in the lite wire format: an optional +/// zigzag-delta timestamp, the size, then the payload. Mirrors the per-frame +/// encoding in [`Subscription::serve_frame`] without the priority machinery, since +/// a one-shot fetch carries a single static priority set on the stream up front. +async fn write_fetch_frame( + writer: &mut Writer, + frame: &mut FrameConsumer, + compression: Compression, + timescale: Option, + prev_ts: &mut u64, + track_stats: &crate::PublisherTrack, +) -> Result<(), Error> { + if timescale.is_some() { + let ts = frame.timestamp.expect("model layer validated timestamp presence"); + let curr = ts.value(); + let delta: i64 = (curr as i128 - *prev_ts as i128) + .try_into() + .map_err(|_| Error::BoundsExceeded(crate::coding::BoundsExceeded))?; + let zz = crate::coding::VarInt::from_zigzag(delta).map_err(crate::coding::EncodeError::from)?; + writer.encode(&zz).await?; + *prev_ts = curr; + } + + match compression { + Compression::None => { + writer.encode(&frame.size).await?; + track_stats.frame(); + while let Some(mut chunk) = frame.read_chunk().await? { + let n = chunk.len() as u64; + writer.write_all(&mut chunk).await?; + track_stats.bytes(n); + } + } + compression => { + let payload = frame.read_all().await?; + let mut chunk = bytes::Bytes::from(compression.compress(&payload)); + let n = chunk.len() as u64; + writer.encode(&n).await?; + track_stats.frame(); + writer.write_all(&mut chunk).await?; + track_stats.bytes(n); + } + } + + Ok(()) } /// Shared per-subscription state for the publisher side. Cloned (cheaply — every diff --git a/rs/moq-net/src/lite/subscriber.rs b/rs/moq-net/src/lite/subscriber.rs index 1aaa124397..c532de4be1 100644 --- a/rs/moq-net/src/lite/subscriber.rs +++ b/rs/moq-net/src/lite/subscriber.rs @@ -9,8 +9,8 @@ use futures::{StreamExt, stream::FuturesUnordered}; use crate::{ AsPath, BandwidthProducer, Broadcast, BroadcastDynamic, Compression, Error, Frame, FrameProducer, Group, - GroupProducer, MAX_FRAME_SIZE, OriginProducer, Path, PathOwned, StatsHandle, SubscriberStats, SubscriberTrack, - Timescale, Timestamp, Track, TrackProducer, TrackRequest, + GroupProducer, GroupRequest, MAX_FRAME_SIZE, OriginProducer, Path, PathOwned, StatsHandle, SubscriberStats, + SubscriberTrack, Timescale, Timestamp, Track, TrackProducer, TrackRequest, coding::{Reader, Stream}, lite, model::BroadcastProducer, @@ -446,27 +446,134 @@ impl Subscriber { } async fn run_broadcast(self, path: PathOwned, mut broadcast: BroadcastDynamic) { - // Actually start serving subscriptions. + // A second handle (sharing the same state) so track and group requests can be + // awaited in the same select without borrowing `broadcast` mutably twice. + let mut fetches = broadcast.clone(); + + // Actually start serving subscriptions and fetches. loop { // Keep serving requests until there are no more consumers. // This way we'll clean up the task when the broadcast is no longer needed. - let request = tokio::select! { - request = broadcast.requested_track() => match request { - Ok(request) => request, - Err(err) => { - tracing::debug!(%err, "broadcast closed"); - break; + tokio::select! { + request = broadcast.requested_track() => { + let request = match request { + Ok(request) => request, + Err(err) => { + tracing::debug!(%err, "broadcast closed"); + break; + } + }; + + let mut this = self.clone(); + let path = path.clone(); + let broadcast = broadcast.clone(); + web_async::spawn(async move { + this.run_subscribe(path, broadcast, request).await; + }); + } + request = fetches.requested_group() => { + let request = match request { + Ok(request) => request, + Err(err) => { + tracing::debug!(%err, "broadcast closed"); + break; + } + }; + + // FETCH only exists upstream on lite-05+. Below that, fail the fetch + // fast rather than open a stream the peer would reject. + if !self.version.has_timestamps() { + request.deny(Error::UnexpectedStream); + continue; } - }, + + let mut this = self.clone(); + let path = path.clone(); + web_async::spawn(async move { + this.run_fetch(path, request).await; + }); + } _ = self.session.closed() => break, - }; + } + } + } - let mut this = self.clone(); - let path = path.clone(); - let broadcast = broadcast.clone(); - web_async::spawn(async move { - this.run_subscribe(path, broadcast, request).await; - }); + /// Serve one downstream fetch by issuing a wire FETCH upstream: open a bidi + /// stream, send FETCH, block on FETCH_OK, then route the group's frames into + /// the producer that resolves the waiting fetcher. + async fn run_fetch(&mut self, path: PathOwned, request: GroupRequest) { + let name = request.name().to_string(); + let group = request.group(); + let priority = request.priority(); + let frame_start = request.frame_start(); + + let abs = self.origin.absolute(&path); + let track_stats = Arc::new(self.stats.broadcast(&abs).subscriber_track(&name)); + + tracing::info!(broadcast = %self.log_path(&path), track = %name, %group, "fetch started"); + + let mut stream = match Stream::open(&self.session, self.version).await { + Ok(stream) => stream, + Err(err) => { + request.deny(err); + return; + } + }; + + if let Err(err) = stream.writer.encode(&lite::ControlType::Fetch).await { + stream.writer.abort(&err); + request.deny(err); + return; + } + + let msg = lite::Fetch { + broadcast: path.as_path(), + track: (&name).into(), + priority, + group, + frame_start, + }; + if let Err(err) = stream.writer.encode(&msg).await { + stream.writer.abort(&err); + request.deny(err); + return; + } + + // The first response MUST be a FETCH_OK; it carries the codec/timescale + // needed to decode the frames that follow on the same stream. + let info = match stream.reader.decode::().await { + Ok(info) => info, + Err(err) => { + stream.writer.abort(&err); + request.deny(err); + return; + } + }; + + // The publisher confirmed: build the producer (unblocking the fetcher) and + // fill it from the wire, reusing the live-subscription frame decode path. + let mut producer = match request.accept(info.timescale) { + Ok(producer) => producer, + Err(err) => { + stream.writer.abort(&err); + return; + } + }; + + let res = tokio::select! { + err = producer.closed() => Err(err), + res = self.run_group(&mut stream.reader, producer.clone(), track_stats, info.compression, info.timescale) => res, + }; + + match res { + Ok(()) => { + let _ = producer.finish(); + tracing::info!(broadcast = %self.log_path(&path), track = %name, %group, "fetch complete"); + } + Err(err) => { + let _ = producer.abort(err.clone()); + tracing::warn!(broadcast = %self.log_path(&path), track = %name, %group, %err, "fetch error"); + } } } diff --git a/rs/moq-net/src/model/broadcast.rs b/rs/moq-net/src/model/broadcast.rs index f8bf814753..d06095fc21 100644 --- a/rs/moq-net/src/model/broadcast.rs +++ b/rs/moq-net/src/model/broadcast.rs @@ -4,7 +4,10 @@ use std::{ task::{Poll, ready}, }; -use crate::{Error, Subscription, TrackProducer, TrackSubscriber, model::track::TrackWeak}; +use crate::{ + Error, Group, GroupConsumer, GroupProducer, Subscription, Timescale, TrackProducer, TrackSubscriber, + model::track::TrackWeak, +}; use super::{OriginList, Track}; @@ -61,6 +64,33 @@ fn fail_resolvers(resolvers: Vec, err: &Error) { } } +/// The slot a pending fetch resolves into: `None` until the dynamic handler +/// accepts (delivering the group) or denies (delivering an error). One-shot +/// analogue of [`PendingSlot`]. +type GroupSlot = Option>; + +/// One waiting fetcher: its priority and the producer side of its resolver channel. +type GroupResolver = (u8, kio::Producer); + +/// A group that has been fetched but not yet served by the dynamic handler. +/// +/// Concurrent fetches for the same `(track, group)` coalesce into one pending +/// request, each adding a resolver so they all receive a consumer for the same +/// group once it is accepted. +#[derive(Default)] +struct GroupPendingRequest { + resolvers: Vec, +} + +/// Resolve every waiting fetcher with `err`. +fn fail_group_resolvers(resolvers: Vec, err: &Error) { + for (_, slot) in resolvers { + if let Ok(mut slot) = slot.write() { + *slot = Some(Err(err.clone())); + } + } +} + #[derive(Default)] struct State { // Weak references for deduplication. Doesn't prevent track auto-close. @@ -74,6 +104,13 @@ struct State { // stays in `requests` (but not here) once handed out as a `TrackRequest`. request_order: VecDeque, + // Pending fetches keyed by (track name, group sequence), waiting for the + // dynamic handler to accept or deny them. + group_requests: HashMap<(String, u64), GroupPendingRequest>, + + // Requested (track, group) pairs in FIFO order for the dynamic handler to drain. + group_request_order: VecDeque<(String, u64)>, + // The current number of dynamic producers. // If this is 0, requests must be empty. dynamic: usize, @@ -99,12 +136,16 @@ impl State { Ok(()) } - /// Drop every pending request, notifying all waiting subscribers with `err`. + /// Drop every pending request, notifying all waiting subscribers and fetchers with `err`. fn abort_requests(&mut self, err: &Error) { self.request_order.clear(); for (_, pending) in self.requests.drain() { fail_resolvers(pending.resolvers, err); } + self.group_request_order.clear(); + for (_, pending) in self.group_requests.drain() { + fail_group_resolvers(pending.resolvers, err); + } } /// Drop a single named pending request, notifying its subscribers with `err`. @@ -114,6 +155,14 @@ impl State { fail_resolvers(pending.resolvers, &err); } } + + /// Drop a single pending fetch, notifying its fetchers with `err`. + fn deny_group_request(&mut self, key: &(String, u64), err: Error) { + self.group_request_order.retain(|k| k != key); + if let Some(pending) = self.group_requests.remove(key) { + fail_group_resolvers(pending.resolvers, &err); + } + } } /// Manages tracks within a broadcast. @@ -352,6 +401,96 @@ impl Drop for TrackRequest { } } +/// A fetch waiting to be served, handed out by [`BroadcastDynamic::requested_group`]. +/// +/// The publisher inspects the requested [`Self::name`] / [`Self::group`], then +/// either [`Self::accept`]s it (creating a [`GroupProducer`] to fill, which +/// resolves all waiting fetchers) or [`Self::deny`]s it. Dropping without doing +/// either denies with [`Error::Cancel`]. +pub struct GroupRequest { + name: String, + group: u64, + priority: u8, + /// First frame the fetcher wants. Always 0 today; the field carries the wire + /// FETCH `frame_start` so the serving side honors it once partial fetches land. + frame_start: u64, + state: kio::Weak, + /// Set once accepted or denied so [`Drop`] doesn't deny a second time. + completed: bool, +} + +impl GroupRequest { + /// The requested track name. + pub fn name(&self) -> &str { + &self.name + } + + /// The requested group sequence. + pub fn group(&self) -> u64 { + self.group + } + + /// The first waiting fetcher's priority, as a hint for the serving stream. + pub fn priority(&self) -> u8 { + self.priority + } + + /// The first frame to serve (0 returns the whole group). + pub fn frame_start(&self) -> u64 { + self.frame_start + } + + /// Serve the request, returning a [`GroupProducer`] to fill and resolving + /// every waiting fetcher with a consumer for it. + /// + /// `timescale` is the parent track's scale (from the wire FETCH_OK), so the + /// producer accepts the timestamped frames that follow. Returns [`Error::Cancel`] + /// if the fetch was cancelled or the broadcast closed while pending. + pub fn accept(mut self, timescale: impl Into>) -> Result { + self.completed = true; + + let mut state = self + .state + .write() + .map_err(|r| r.abort.clone().unwrap_or(Error::Cancel))?; + + let key = (self.name.clone(), self.group); + let pending = state.group_requests.remove(&key).ok_or(Error::Cancel)?; + state.group_request_order.retain(|k| k != &key); + + let producer = GroupProducer::new_with_timescale(Group { sequence: self.group }, timescale.into()); + + // Build each fetcher's consumer now (not when it polls) so it counts toward + // the producer immediately, mirroring TrackRequest::accept. + for (_, slot) in pending.resolvers { + if let Ok(mut slot) = slot.write() { + *slot = Some(Ok(producer.consume())); + } + } + + drop(state); + Ok(producer) + } + + /// Reject the fetch, waking all waiting fetchers with `err`. + pub fn deny(mut self, err: Error) { + self.completed = true; + if let Ok(mut state) = self.state.write() { + state.deny_group_request(&(self.name.clone(), self.group), err); + } + } +} + +impl Drop for GroupRequest { + fn drop(&mut self) { + if !self.completed + && let Ok(mut state) = self.state.write() + { + state.deny_group_request(&(self.name.clone(), self.group), Error::Cancel); + } + } +} + /// A pending subscription returned by [`TrackConsumer::subscribe`]. /// /// The subscription isn't live until the publisher accepts it (for the wire, @@ -419,6 +558,69 @@ impl std::future::Future for TrackPending { } } +/// A pending fetch returned by [`TrackConsumer::fetch`]. +/// +/// One-shot analogue of [`TrackPending`]: resolves to the [`GroupConsumer`] once +/// the group is available (already cached, or FETCH_OK on the wire), or an error. +pub struct FetchPending { + inner: FetchPendingInner, + /// Kept alive between `Future::poll` calls so its registration in the resolver + /// channel stays valid until the next poll replaces it. + waiter: Option, +} + +enum FetchPendingInner { + /// Resolved synchronously: the group was cached, or it failed immediately. + Ready(Result), + /// Waiting for the dynamic handler to accept or deny via the wire. + Waiting(kio::Consumer), +} + +impl FetchPending { + fn ready(result: Result) -> Self { + Self { + inner: FetchPendingInner::Ready(result), + waiter: None, + } + } + + fn waiting(consumer: kio::Consumer) -> Self { + Self { + inner: FetchPendingInner::Waiting(consumer), + waiter: None, + } + } + + /// Poll for the resolved [`GroupConsumer`], without blocking. + pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll> { + match &self.inner { + FetchPendingInner::Ready(result) => Poll::Ready(result.clone()), + FetchPendingInner::Waiting(consumer) => match consumer.poll(waiter, |slot| match &**slot { + Some(result) => Poll::Ready(result.clone()), + None => Poll::Pending, + }) { + Poll::Ready(Ok(result)) => Poll::Ready(result), + // Channel closed: the resolver may have left the final result behind. + Poll::Ready(Err(closed)) => Poll::Ready(match &*closed { + Some(result) => result.clone(), + None => Err(Error::Cancel), + }), + Poll::Pending => Poll::Pending, + }, + } + } +} + +impl std::future::Future for FetchPending { + type Output = Result; + + fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { + let this = self.get_mut(); + this.waiter = Some(kio::Waiter::new(cx.waker().clone())); + this.poll_ok(this.waiter.as_ref().unwrap()) + } +} + /// Handles on-demand track creation for a broadcast. /// /// When a consumer requests a track that doesn't exist, the dynamic producer @@ -505,6 +707,37 @@ impl BroadcastDynamic { kio::wait(|waiter| self.poll_requested_track(waiter)).await } + /// Poll for the next consumer-requested group (fetch), without blocking. + pub fn poll_requested_group(&mut self, waiter: &kio::Waiter) -> Poll> { + let weak = self.state.weak(); + self.poll(waiter, |state| { + let Some(key) = state.group_request_order.pop_front() else { + return Poll::Pending; + }; + // The key stays in `group_requests` so concurrent fetchers can still + // coalesce onto it until the publisher accepts or denies. + let pending = state.group_requests.get(&key).expect("group_request_order out of sync"); + let priority = pending.resolvers.first().map(|(p, _)| *p).unwrap_or(0); + Poll::Ready((key, priority)) + }) + .map(|res| { + res.map(|((name, group), priority)| GroupRequest { + name, + group, + priority, + // Always 0 today; partial-fetch resume would set this from the request. + frame_start: 0, + state: weak, + completed: false, + }) + }) + } + + /// Block until a consumer fetches a group, returning a [`GroupRequest`] to serve. + pub async fn requested_group(&mut self) -> Result { + kio::wait(|waiter| self.poll_requested_group(waiter)).await + } + /// Create a consumer that can subscribe to tracks in this broadcast. pub fn consume(&self) -> BroadcastConsumer { BroadcastConsumer { @@ -572,6 +805,17 @@ impl BroadcastDynamic { pub fn assert_no_request(&mut self) { assert!(self.requested_track().now_or_never().is_none(), "should have blocked"); } + + pub fn assert_group_request(&mut self) -> GroupRequest { + self.requested_group() + .now_or_never() + .expect("should not have blocked") + .expect("should not have errored") + } + + pub fn assert_no_group_request(&mut self) { + assert!(self.requested_group().now_or_never().is_none(), "should have blocked"); + } } /// Subscribe to arbitrary broadcast/tracks. @@ -652,6 +896,70 @@ impl BroadcastConsumer { TrackPending::waiting(consumer) } + /// Register a one-shot fetch for `group` on `name` and return a [`FetchPending`] + /// that resolves once the group is available. + /// + /// Returns a cached group directly if a live producer already holds it; bypasses + /// the cache and queues a wire request if the cached group is aborted or evicted. + /// Resolves to [`Error::NotFound`] if the group can never exist (past the track's + /// final sequence) or no dynamic producer exists to fetch it. + fn request_fetch(&self, name: &str, group: u64, fetch: Fetch) -> FetchPending { + // Upgrade to a temporary producer so we can modify the state. + let Some(producer) = self.state.produce() else { + let err = self.state.read().abort.clone().unwrap_or(Error::Dropped); + return FetchPending::ready(Err(err)); + }; + let mut state = match modify(&producer) { + Ok(state) => state, + Err(err) => return FetchPending::ready(Err(err)), + }; + + // Return a cached group if a live producer is publishing the track. + if let Some(weak) = state.tracks.get(name) { + if weak.is_closed() { + // Drop the stale entry and fall through to a fresh request. + state.tracks.remove(name); + } else { + let sub = weak.subscribe(Subscription::default()); + match sub.poll_get_group(&kio::Waiter::noop(), group) { + Poll::Ready(Ok(Some(mut cached))) => { + // Bypass an aborted cached group with a fresh fetch; otherwise return it + // (complete, or in-progress with frames still streaming in). + if !matches!(cached.poll_finished(&kio::Waiter::noop()), Poll::Ready(Err(_))) { + return FetchPending::ready(Ok(cached)); + } + } + // At or past the track's final sequence: the group can never exist. + Poll::Ready(Ok(None)) => return FetchPending::ready(Err(Error::NotFound)), + Poll::Ready(Err(err)) => return FetchPending::ready(Err(err)), + // Not cached (old / evicted): fall through to a wire request. + Poll::Pending => {} + } + } + } + + let slot = kio::Producer::new(None); + let consumer = slot.consume(); + let key = (name.to_string(), group); + + if let Some(pending) = state.group_requests.get_mut(&key) { + // Coalesce onto an in-flight fetch for the same track + group. + pending.resolvers.push((fetch.priority, slot)); + } else if state.dynamic == 0 { + return FetchPending::ready(Err(Error::NotFound)); + } else { + state.group_requests.insert( + key.clone(), + GroupPendingRequest { + resolvers: vec![(fetch.priority, slot)], + }, + ); + state.group_request_order.push_back(key); + } + + FetchPending::waiting(consumer) + } + /// Block until the broadcast is closed and return the cause. /// /// Returns [`Error::Dropped`] if every producer was dropped without an @@ -681,13 +989,27 @@ impl BroadcastConsumer { } } +/// Options for a one-shot [`TrackConsumer::fetch`] of a past group. +#[derive(Clone, Debug, Default)] +pub struct Fetch { + /// Delivery priority for the fetched group's stream. Defaults to 0. + pub priority: u8, +} + +impl Fetch { + /// Set the delivery priority, returning `self` for chaining. + pub fn with_priority(mut self, priority: u8) -> Self { + self.priority = priority; + self + } +} + /// A handle to a single track within a broadcast. /// /// Obtained from [`BroadcastConsumer::consume_track`]. Holding it sends nothing /// to the publisher; it just names a track you can [`subscribe`](Self::subscribe) -/// to (a live, ongoing stream of groups) later. The same handle can be subscribed -/// to multiple times, and clones are cheap. -// TODO: add `fetch` for one-shot retrieval of a past group range. +/// to (a live, ongoing stream of groups) or [`fetch`](Self::fetch) a single past +/// group from. The same handle can be reused, and clones are cheap. #[derive(Clone)] pub struct TrackConsumer { broadcast: BroadcastConsumer, @@ -713,6 +1035,20 @@ impl TrackConsumer { pub fn subscribe(&self, subscription: Subscription) -> TrackPending { self.broadcast.request_subscribe(&self.name, subscription) } + + /// Fetch a single past group, without holding a live subscription. + /// + /// Returns the [`GroupConsumer`] once the group is available: immediately if a + /// live producer already has it cached, otherwise after a wire FETCH completes + /// (blocking on FETCH_OK, like [`subscribe`](Self::subscribe) blocks on + /// SUBSCRIBE_OK). `options` accepts `None`, a [`Fetch`], or `Fetch::default()`. + /// + /// Errors with [`Error::NotFound`] if the group can never exist (past the + /// track's final sequence) or no publisher handles fetch. + pub async fn fetch(&self, group: u64, options: impl Into>) -> Result { + let options = options.into().unwrap_or_default(); + self.broadcast.request_fetch(&self.name, group, options).await + } } #[cfg(test)] @@ -743,6 +1079,19 @@ mod test { }}; } + /// Fetch and assert the result hasn't resolved yet (it stays pending until a + /// publisher accepts). Returns the [`FetchPending`] to resolve after accepting. + macro_rules! fetch_pending { + ($consumer:expr, $name:expr, $group:expr) => {{ + let pending = $consumer.request_fetch($name, $group, Fetch::default()); + assert!( + pending.poll_ok(&kio::Waiter::noop()).is_pending(), + "fetch should stay pending until the request is accepted" + ); + pending + }}; + } + #[tokio::test] async fn insert() { let mut producer = Broadcast::new().produce(); @@ -961,4 +1310,154 @@ mod test { // instead of failing with NotFound. let _fut = subscribe_pending!(consumer, "track1"); } + + #[tokio::test] + async fn fetch_cache_hit() { + let mut producer = Broadcast::new().produce(); + let mut track = Track::new("track1").produce(); + producer.assert_insert_track(&track); + + // Produce a cached group. + let mut group = track.append_group().unwrap(); // seq 0 + group.write_frame(bytes::Bytes::from_static(b"hello")).unwrap(); + group.finish().unwrap(); + + // A cached group is returned directly, no dynamic handler needed. + let consumer = producer.consume(); + let mut g = consumer.consume_track("track1").fetch(0, None).await.unwrap(); + assert_eq!(g.sequence, 0); + assert_eq!(&g.read_frame().await.unwrap().unwrap()[..], b"hello"); + } + + #[tokio::test] + async fn fetch_miss_queues_request() { + let mut broadcast = Broadcast::new().produce().dynamic(); + let consumer = broadcast.consume(); + + // No live track has the group, so the fetch queues a dynamic request. + let fetch_fut = fetch_pending!(consumer, "track1", 5); + + let request = broadcast.assert_group_request(); + broadcast.assert_no_group_request(); + assert_eq!(request.name(), "track1"); + assert_eq!(request.group(), 5); + assert_eq!(request.frame_start(), 0); + + let mut group = request.accept(None).unwrap(); + let mut g = fetch_fut.await.unwrap(); + assert_eq!(g.sequence, 5); + + group.write_frame(bytes::Bytes::from_static(b"hi")).unwrap(); + group.finish().unwrap(); + assert_eq!(&g.read_frame().await.unwrap().unwrap()[..], b"hi"); + } + + #[tokio::test] + async fn fetch_past_final_not_found() { + let mut producer = Broadcast::new().produce(); + let mut track = Track::new("track1").produce(); + producer.assert_insert_track(&track); + track.append_group().unwrap(); // seq 0 + track.finish().unwrap(); // final_sequence = 1 + + // A group at or past the final sequence can never exist. + let consumer = producer.consume(); + assert!(matches!( + consumer.consume_track("track1").fetch(5, None).await, + Err(Error::NotFound) + )); + } + + #[tokio::test] + async fn fetch_errored_cache_bypass() { + let mut broadcast = Broadcast::new().produce(); + let mut dynamic = broadcast.dynamic(); + let mut track = Track::new("track1").produce(); + broadcast.assert_insert_track(&track); + + // A cached but aborted group must be bypassed, not returned. + let mut group = track.append_group().unwrap(); // seq 0 + group.abort(Error::Cancel).unwrap(); + + let consumer = broadcast.consume(); + let fetch_fut = fetch_pending!(consumer, "track1", 0); + + // The aborted cache was skipped: a fresh dynamic request is queued instead. + let request = dynamic.assert_group_request(); + assert_eq!(request.group(), 0); + let mut fresh = request.accept(None).unwrap(); + let mut g = fetch_fut.await.unwrap(); + + fresh.write_frame(bytes::Bytes::from_static(b"ok")).unwrap(); + fresh.finish().unwrap(); + assert_eq!(&g.read_frame().await.unwrap().unwrap()[..], b"ok"); + } + + #[tokio::test] + async fn fetch_coalesces() { + let mut broadcast = Broadcast::new().produce().dynamic(); + let consumer = broadcast.consume(); + + // Two fetches for the same track + group coalesce into one request. + let f1 = fetch_pending!(consumer, "track1", 3); + let f2 = fetch_pending!(consumer, "track1", 3); + + let request = broadcast.assert_group_request(); + broadcast.assert_no_group_request(); + + let mut group = request.accept(None).unwrap(); + let mut g1 = f1.await.unwrap(); + let mut g2 = f2.await.unwrap(); + assert_eq!(g1.sequence, 3); + assert_eq!(g2.sequence, 3); + + // Both resolve to consumers of the same group. + group.write_frame(bytes::Bytes::from_static(b"x")).unwrap(); + group.finish().unwrap(); + assert_eq!(&g1.read_frame().await.unwrap().unwrap()[..], b"x"); + assert_eq!(&g2.read_frame().await.unwrap().unwrap()[..], b"x"); + } + + #[tokio::test] + async fn fetch_not_found_without_dynamic() { + // No dynamic handler and no cached track: the fetch fails outright. + let broadcast = Broadcast::new().produce(); + let consumer = broadcast.consume(); + assert!(matches!( + consumer.consume_track("track1").fetch(0, None).await, + Err(Error::NotFound) + )); + } + + #[tokio::test] + async fn fetch_request_deny() { + let mut broadcast = Broadcast::new().produce().dynamic(); + let consumer = broadcast.consume(); + + let f = fetch_pending!(consumer, "track1", 0); + broadcast.assert_group_request().deny(Error::NotFound); + assert!(f.await.is_err()); + } + + #[tokio::test] + async fn fetch_request_drop_cancels() { + let mut broadcast = Broadcast::new().produce().dynamic(); + let consumer = broadcast.consume(); + + let f = fetch_pending!(consumer, "track1", 0); + // Dropping the request without accepting/denying cancels the fetch. + drop(broadcast.assert_group_request()); + assert!(f.await.is_err()); + } + + #[tokio::test] + async fn fetch_dynamic_drop_aborts() { + let broadcast = Broadcast::new().produce().dynamic(); + let consumer = broadcast.consume(); + + let f = fetch_pending!(consumer, "track1", 0); + // No dynamic producer remains to serve the pending fetch. + drop(broadcast); + assert!(f.await.is_err()); + } } diff --git a/rs/moq-net/src/model/group.rs b/rs/moq-net/src/model/group.rs index 247a43accc..cb16f7fda7 100644 --- a/rs/moq-net/src/model/group.rs +++ b/rs/moq-net/src/model/group.rs @@ -261,6 +261,7 @@ impl GroupProducer { GroupConsumer { info: self.info.clone(), state: self.state.consume(), + timescale: self.timescale, index: 0, } } @@ -305,6 +306,10 @@ pub struct GroupConsumer { // Immutable stream state. info: Group, + // Parent track's negotiated timescale, copied from the producer. Lets the + // wire publisher decide whether to emit per-frame timestamps for a fetched group. + timescale: Option, + // The number of frames we've read. // NOTE: Cloned readers inherit this offset, but then run in parallel. index: usize, @@ -319,6 +324,11 @@ impl std::ops::Deref for GroupConsumer { } impl GroupConsumer { + /// The parent track's negotiated timescale, or `None` for untimed tracks. + pub fn timescale(&self) -> Option { + self.timescale + } + // A helper to automatically apply Dropped if the state is closed without an error. fn poll(&self, waiter: &kio::Waiter, f: F) -> Poll> where diff --git a/rs/moq-relay/src/web.rs b/rs/moq-relay/src/web.rs index 3836b6c8f0..1a313b43f2 100644 --- a/rs/moq-relay/src/web.rs +++ b/rs/moq-relay/src/web.rs @@ -519,29 +519,31 @@ async fn serve_fetch( .announced_broadcast("") .await .ok_or(StatusCode::NOT_FOUND)?; - let mut track = broadcast - .consume_track(&track) - .subscribe(moq_net::Subscription::default()) - .await - .map_err(|err| match err { - moq_net::Error::NotFound => StatusCode::NOT_FOUND, - _ => StatusCode::INTERNAL_SERVER_ERROR, - })?; let group = match params.group { - FetchGroup::Latest => match track.latest() { - Some(sequence) => track.get_group(sequence).await, - None => track.recv_group().await, + // "latest" needs a live subscription to learn the newest sequence; + // fetch only retrieves a specified past group. + FetchGroup::Latest => match broadcast + .consume_track(&track) + .subscribe(moq_net::Subscription::default()) + .await + { + Ok(mut sub) => match sub.latest() { + Some(sequence) => sub.get_group(sequence).await, + None => sub.recv_group().await, + }, + Err(err) => Err(err), }, - FetchGroup::Num(sequence) => track.get_group(sequence).await, + // A one-shot fetch, no subscription required. + FetchGroup::Num(sequence) => broadcast.consume_track(&track).fetch(sequence, None).await.map(Some), }; let group = match group { Ok(Some(group)) => group, - Ok(None) => return Err(StatusCode::NOT_FOUND), + Ok(None) | Err(moq_net::Error::NotFound) => return Err(StatusCode::NOT_FOUND), Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR), }; - tracing::info!(track = %track.name, group = %group.sequence, "serving group"); + tracing::info!(%track, group = %group.sequence, "serving group"); match params.frame { FetchFrame::Num(index) => match group.get_frame(index).await { From f5cdf3e2896f0fdd58890638804c31fa0cac1620 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 4 Jun 2026 21:06:04 -0700 Subject: [PATCH 2/3] kio: add a Future trait + Pending adapter for named poll-based futures Implementing std Future on a kio-backed type means stashing the strong Waiter in an Option field and replacing it every poll, because the channel's WaiterList only holds a Weak (a dropped waiter loses its wakeup). kio::wait() already hides this for closures, but named futures returned from sync methods had to repeat the dance by hand. Add `kio::Future` (a poll-based trait, method `poll(&mut self, waiter)`) and `kio::Pending`, which carries the retained Waiter and provides the std Future impl once. Pending derefs to the inner value, so any inherent methods you put on it are reachable through the awaitable handle. Convert `TrackFetchPending` to `kio::Pending` as the first user: the hand-written waiter field and std Future impl are gone; `TrackFetch` writes only `kio::Future::poll`. Co-Authored-By: Claude Opus 4.8 (1M context) --- rs/kio/src/future.rs | 142 ++++++++++++++++++++++++++++++++++ rs/kio/src/lib.rs | 2 + rs/moq-net/src/model/track.rs | 59 ++++++-------- 3 files changed, 169 insertions(+), 34 deletions(-) create mode 100644 rs/kio/src/future.rs diff --git a/rs/kio/src/future.rs b/rs/kio/src/future.rs new file mode 100644 index 0000000000..510a1d76ef --- /dev/null +++ b/rs/kio/src/future.rs @@ -0,0 +1,142 @@ +use std::{ + ops::{Deref, DerefMut}, + pin::Pin, + task::{Context, Poll}, +}; + +use crate::Waiter; + +/// A pollable computation backed by kio channels. +/// +/// Implementors write only [`Self::poll`], registering the [`Waiter`] with the +/// channels they read. Wrap the value in [`Pending`] to get a real +/// [`std::future::Future`]. +/// +/// This exists because a kio [`Waiter`] holds the strong `Arc` while the +/// channel's [`crate::WaiterList`] keeps only a `Weak`. A bare +/// [`std::future::Future`] would have to stash the strong `Waiter` in a field and +/// replace it every poll (or lose its wakeup); [`Pending`] does that once so each +/// implementor doesn't have to. +pub trait Future: Unpin { + type Output; + + /// Poll for the output, registering `waiter` with the relevant channels if not + /// yet ready. + fn poll(&mut self, waiter: &Waiter) -> Poll; +} + +/// Adapts a kio [`Future`] into a [`std::future::Future`], retaining the strong +/// [`Waiter`] between polls so its weak registration stays live. +/// +/// Derefs to the inner value, so any inherent methods you define on it are +/// reachable through the pending handle (e.g. a non-blocking `poll`, or an +/// `update`). +pub struct Pending { + inner: F, + // Retain the previous waiter so its Weak registration survives until the next + // poll replaces it (see [`crate::WaiterList`]). + waiter: Option, +} + +impl Pending { + /// Wrap a [`Future`] so it can be `.await`ed. + pub fn new(inner: F) -> Self { + Self { inner, waiter: None } + } + + /// Consume the wrapper, returning the inner value. + pub fn into_inner(self) -> F { + self.inner + } +} + +impl Deref for Pending { + type Target = F; + + fn deref(&self) -> &F { + &self.inner + } +} + +impl DerefMut for Pending { + fn deref_mut(&mut self) -> &mut F { + &mut self.inner + } +} + +impl std::future::Future for Pending { + type Output = F::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // Replacing drops the previous waiter, killing its Weak ref in the list so + // the inner poll's register call can recycle the slot (see `WaiterList`). + // `Pending` is `Unpin` (F is, via the trait bound), so this deref is sound. + let this = &mut *self; + this.waiter = Some(Waiter::new(cx.waker().clone())); + Future::poll(&mut this.inner, this.waiter.as_ref().unwrap()) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::Producer; + + /// A pollable that waits for the channel value to reach a threshold, with an + /// inherent method reachable through `Pending`'s `DerefMut`. + struct AtLeast { + consumer: crate::Consumer, + threshold: u64, + } + + impl AtLeast { + fn bump_threshold(&mut self) { + self.threshold += 1; + } + } + + impl Future for AtLeast { + type Output = u64; + + fn poll(&mut self, waiter: &Waiter) -> Poll { + let threshold = self.threshold; + match self.consumer.poll(waiter, |v| { + let current = **v; + if current >= threshold { + Poll::Ready(current) + } else { + Poll::Pending + } + }) { + Poll::Ready(Ok(v)) => Poll::Ready(v), + _ => Poll::Pending, + } + } + } + + #[test] + fn pending_derefs_and_drives() { + use std::task::Waker; + + let producer = Producer::new(0u64); + let mut pending = Pending::new(AtLeast { + consumer: producer.consume(), + threshold: 5, + }); + + // Inherent method on the inner reached via DerefMut. + pending.bump_threshold(); // threshold now 6 + + // The kio-level poll (reached through Deref) is pending until the value catches up. + assert!(Future::poll(&mut *pending, &Waiter::noop()).is_pending()); + + if let Ok(mut v) = producer.write() { + *v = 6; + } + + // The std Future resolves once the threshold is met. + let mut cx = Context::from_waker(Waker::noop()); + let mut pending = std::pin::pin!(pending); + assert_eq!(std::future::Future::poll(pending.as_mut(), &mut cx), Poll::Ready(6)); + } +} diff --git a/rs/kio/src/lib.rs b/rs/kio/src/lib.rs index e669b18370..5493a91ee1 100644 --- a/rs/kio/src/lib.rs +++ b/rs/kio/src/lib.rs @@ -14,10 +14,12 @@ mod lock; mod waiter; mod consumer; +mod future; mod producer; mod weak; pub use consumer::Consumer; +pub use future::{Future, Pending}; pub use producer::{Mut, Producer, Ref}; pub use waiter::{Waiter, WaiterList, wait}; pub use weak::Weak; diff --git a/rs/moq-net/src/model/track.rs b/rs/moq-net/src/model/track.rs index 508db0e398..0e6126907d 100644 --- a/rs/moq-net/src/model/track.rs +++ b/rs/moq-net/src/model/track.rs @@ -658,7 +658,7 @@ impl TrackProducer { } pub fn poll_requested_fetch(&mut self, waiter: &kio::Waiter) -> Poll> { - let res = self.state.poll(waiter, |state| { + match self.state.poll(waiter, |state| { // Only take the `&mut` (which flags the state modified) once we actually // have a request to pop, so idle polls don't wake unrelated waiters. if !state.fetches.is_empty() { @@ -668,8 +668,7 @@ impl TrackProducer { Some(err) => Poll::Ready(Err(err.clone())), None => Poll::Pending, } - }); - match res { + }) { Poll::Ready(Ok(res)) => Poll::Ready(res), Poll::Ready(Err(state)) => Poll::Ready(Err(state.abort.clone().unwrap_or(Error::Dropped))), Poll::Pending => Poll::Pending, @@ -800,6 +799,7 @@ impl TrackConsumer { // Only signal the producer when the group isn't already determined by // the cache (cached, evicted-past-final, or aborted); otherwise the // pending resolves straight from `poll_get_group` with no wire fetch. + // TODO: This is a tiny bit racey if state.poll_get_group(sequence).is_pending() { state.fetches.push_back(FetchRequest { sequence, @@ -810,11 +810,10 @@ impl TrackConsumer { Err(state) => return Err(state.abort.clone().unwrap_or(Error::Dropped)), }; - Ok(TrackFetchPending { + Ok(kio::Pending::new(TrackFetch { state: self.state.clone(), sequence, - waiter: None, - }) + })) } pub fn info(&self) -> TrackInfoPending { @@ -925,41 +924,32 @@ pub struct FetchRequest { pub priority: u8, } -/// A pending fetch returned by [`TrackConsumer::fetch`]. +/// The pollable state of a [`TrackConsumer::fetch`]. /// -/// Resolves to the [`GroupConsumer`] once the group lands in the track's cache -/// (already present, or produced after a wire FETCH), or an error. -pub struct TrackFetchPending { +/// Awaited via the [`TrackFetchPending`] wrapper; resolves to the +/// [`GroupConsumer`] once the group lands in the track's cache (already present, +/// or produced after a wire FETCH), or [`Error::NotFound`] if it can never exist. +pub struct TrackFetch { state: kio::Consumer, sequence: u64, - // Kept alive between `Future::poll` calls so the registered weak waker stays - // upgradeable until the next poll replaces it (see [`TrackSubscriberPending`]). - waiter: Option, } -impl TrackFetchPending { - pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll> { +impl kio::Future for TrackFetch { + type Output = Result; + + fn poll(&mut self, waiter: &kio::Waiter) -> Poll { let group = ready!(self.state.poll(waiter, |state| state.poll_get_group(self.sequence))) .map_err(|e| e.abort.clone().unwrap_or(Error::Dropped))??; - match group { - Some(group) => Poll::Ready(Ok(group)), - // `poll_get_group` only returns `None` once the group is at/past the - // final sequence, so it can never be produced. - None => Poll::Ready(Err(Error::NotFound)), - } + // `poll_get_group` only returns `None` once the group is at/past the + // final sequence, so it can never be produced. + Poll::Ready(group.ok_or(Error::NotFound)) } } -impl Future for TrackFetchPending { - type Output = Result; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - this.waiter = Some(kio::Waiter::new(cx.waker().clone())); - this.poll_ok(this.waiter.as_ref().unwrap()) - } -} +/// A pending fetch returned by [`TrackConsumer::fetch`]. `.await` it for the +/// [`GroupConsumer`]. +pub type TrackFetchPending = kio::Pending; /// A live subscription to a track, used to read its groups. /// @@ -2087,8 +2077,9 @@ mod test { let consumer = producer.consume(); // The group isn't cached yet, so the fetch stays pending and queues a request. - let pending = consumer.fetch(5, Fetch::default().with_priority(7)).unwrap(); - assert!(pending.poll_ok(&kio::Waiter::noop()).is_pending()); + // `*pending` derefs the wrapper to the inner `TrackFetch` (a `kio::Future`). + let mut pending = consumer.fetch(5, Fetch::default().with_priority(7)).unwrap(); + assert!(kio::Future::poll(&mut *pending, &kio::Waiter::noop()).is_pending()); let req = producer .requested_fetch() @@ -2132,8 +2123,8 @@ mod test { let mut producer = TrackProducer::new("test", None); let consumer = producer.consume(); - let pending = consumer.fetch(3, None).unwrap(); - assert!(pending.poll_ok(&kio::Waiter::noop()).is_pending()); + let mut pending = consumer.fetch(3, None).unwrap(); + assert!(kio::Future::poll(&mut *pending, &kio::Waiter::noop()).is_pending()); producer.abort(Error::Cancel).unwrap(); assert!(pending.await.is_err()); From 56063222e3cd17aca00100a3dbc2c8433a5fdb0b Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 4 Jun 2026 21:16:42 -0700 Subject: [PATCH 3/3] moq-net: convert the remaining track futures to kio::Pending Drop the hand-written `waiter: Option` field and `impl std::future::Future` from `TrackSubscriberPending` and `TrackInfoPending`, mirroring `TrackFetchPending`: each is now `kio::Pending` where `Inner` implements `kio::Future`. `kio::Future::poll` takes `&self` (kio channels poll immutably), so a pollable can be driven through a shared borrow. `Pending` derefs to the inner, so the existing `poll_ok` / `update` surface keeps working through the wrapper with no caller changes (moq-mux polls `poll_ok` through an `&self`-borrowed enum, and the broadcast test macro calls it on a shared handle). No hand-written future plumbing or retained-waiter fields remain in moq-net. Co-Authored-By: Claude Opus 4.8 (1M context) --- rs/kio/src/future.rs | 12 ++++-- rs/moq-net/src/model/track.rs | 69 +++++++++++++++++------------------ 2 files changed, 41 insertions(+), 40 deletions(-) diff --git a/rs/kio/src/future.rs b/rs/kio/src/future.rs index 510a1d76ef..d05cc65f96 100644 --- a/rs/kio/src/future.rs +++ b/rs/kio/src/future.rs @@ -22,7 +22,11 @@ pub trait Future: Unpin { /// Poll for the output, registering `waiter` with the relevant channels if not /// yet ready. - fn poll(&mut self, waiter: &Waiter) -> Poll; + /// + /// Takes `&self`: kio channels poll immutably, so a pollable can be driven + /// through a shared borrow (e.g. while it lives inside an `&self`-borrowed enum). + /// Carry any per-poll mutable state in a kio channel or a [`std::cell`] type. + fn poll(&self, waiter: &Waiter) -> Poll; } /// Adapts a kio [`Future`] into a [`std::future::Future`], retaining the strong @@ -73,7 +77,7 @@ impl std::future::Future for Pending { // `Pending` is `Unpin` (F is, via the trait bound), so this deref is sound. let this = &mut *self; this.waiter = Some(Waiter::new(cx.waker().clone())); - Future::poll(&mut this.inner, this.waiter.as_ref().unwrap()) + Future::poll(&this.inner, this.waiter.as_ref().unwrap()) } } @@ -98,7 +102,7 @@ mod test { impl Future for AtLeast { type Output = u64; - fn poll(&mut self, waiter: &Waiter) -> Poll { + fn poll(&self, waiter: &Waiter) -> Poll { let threshold = self.threshold; match self.consumer.poll(waiter, |v| { let current = **v; @@ -128,7 +132,7 @@ mod test { pending.bump_threshold(); // threshold now 6 // The kio-level poll (reached through Deref) is pending until the value catches up. - assert!(Future::poll(&mut *pending, &Waiter::noop()).is_pending()); + assert!(Future::poll(&*pending, &Waiter::noop()).is_pending()); if let Ok(mut v) = producer.write() { *v = 6; diff --git a/rs/moq-net/src/model/track.rs b/rs/moq-net/src/model/track.rs index 0e6126907d..a8b8afcefa 100644 --- a/rs/moq-net/src/model/track.rs +++ b/rs/moq-net/src/model/track.rs @@ -19,9 +19,8 @@ use super::{Group, GroupConsumer, GroupProducer}; use std::{ collections::{HashSet, VecDeque}, - pin::Pin, sync::Arc, - task::{Context, Poll, ready}, + task::{Poll, ready}, time::Duration, }; @@ -776,12 +775,11 @@ impl TrackConsumer { Err(state) => return Err(state.abort.clone().unwrap_or(Error::Dropped)), }; - Ok(TrackSubscriberPending { + Ok(kio::Pending::new(TrackSubscribe { name: self.name.clone(), state: self.state.clone(), subscription, - waiter: None, - }) + })) } /// Fetch a single past group, without holding a live subscription. @@ -817,24 +815,21 @@ impl TrackConsumer { } pub fn info(&self) -> TrackInfoPending { - TrackInfoPending { + kio::Pending::new(TrackInfoQuery { state: self.state.clone(), - waiter: None, - } + }) } } -pub struct TrackSubscriberPending { +/// The pollable state of a [`TrackConsumer::subscribe`]; awaited via the +/// [`TrackSubscriberPending`] wrapper, whose `DerefMut` exposes [`Self::update`]. +pub struct TrackSubscribe { name: Arc, state: kio::Consumer, subscription: kio::Producer, - // Kept alive between `Future::poll` calls so the weak waker kio registered - // stays upgradeable until the next poll replaces it. A temporary would drop - // after poll returns, leaving a dead weak ref and a lost wakeup on accept. - waiter: Option, } -impl TrackSubscriberPending { +impl TrackSubscribe { pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll> { // Wait until the track info is available let info = ready!(self.state.poll(waiter, |state| state.poll_info())) @@ -852,6 +847,7 @@ impl TrackSubscriberPending { })) } + /// Change the subscription preferences before (or after) it resolves. pub fn update(&mut self, subscription: Subscription) { if let Ok(mut state) = self.subscription.write() { *state = subscription; @@ -861,26 +857,26 @@ impl TrackSubscriberPending { } } -impl Future for TrackSubscriberPending { +impl kio::Future for TrackSubscribe { type Output = Result; - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - // Replacing drops the previous waiter, freeing its slot so the register - // call below can recycle it (see kio's weak-waker GC). - this.waiter = Some(kio::Waiter::new(cx.waker().clone())); - this.poll_ok(this.waiter.as_ref().unwrap()) + fn poll(&self, waiter: &kio::Waiter) -> Poll { + self.poll_ok(waiter) } } -pub struct TrackInfoPending { +/// A pending subscription returned by [`TrackConsumer::subscribe`]. `.await` it for +/// the [`TrackSubscriber`], or call [`TrackSubscribe::update`] / [`TrackSubscribe::poll_ok`] +/// through its `Deref`. +pub type TrackSubscriberPending = kio::Pending; + +/// The pollable state of a [`TrackConsumer::info`]; awaited via the +/// [`TrackInfoPending`] wrapper. +pub struct TrackInfoQuery { state: kio::Consumer, - // See [`TrackSubscriberPending::waiter`]: kept alive so the registered weak - // waker stays upgradeable between polls. - waiter: Option, } -impl TrackInfoPending { +impl TrackInfoQuery { pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll> { // Wait until the track info is available let info = ready!(self.state.poll(waiter, |state| state.poll_info())) @@ -889,16 +885,17 @@ impl TrackInfoPending { } } -impl Future for TrackInfoPending { +impl kio::Future for TrackInfoQuery { type Output = Result; - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - this.waiter = Some(kio::Waiter::new(cx.waker().clone())); - this.poll_ok(this.waiter.as_ref().unwrap()) + fn poll(&self, waiter: &kio::Waiter) -> Poll { + self.poll_ok(waiter) } } +/// A pending [`TrackInfo`] lookup returned by [`TrackConsumer::info`]. `.await` it. +pub type TrackInfoPending = kio::Pending; + /// Options for a one-shot [`TrackConsumer::fetch`] of a past group. #[derive(Clone, Debug, Default)] pub struct Fetch { @@ -937,7 +934,7 @@ pub struct TrackFetch { impl kio::Future for TrackFetch { type Output = Result; - fn poll(&mut self, waiter: &kio::Waiter) -> Poll { + fn poll(&self, waiter: &kio::Waiter) -> Poll { let group = ready!(self.state.poll(waiter, |state| state.poll_get_group(self.sequence))) .map_err(|e| e.abort.clone().unwrap_or(Error::Dropped))??; @@ -2078,8 +2075,8 @@ mod test { // The group isn't cached yet, so the fetch stays pending and queues a request. // `*pending` derefs the wrapper to the inner `TrackFetch` (a `kio::Future`). - let mut pending = consumer.fetch(5, Fetch::default().with_priority(7)).unwrap(); - assert!(kio::Future::poll(&mut *pending, &kio::Waiter::noop()).is_pending()); + let pending = consumer.fetch(5, Fetch::default().with_priority(7)).unwrap(); + assert!(kio::Future::poll(&*pending, &kio::Waiter::noop()).is_pending()); let req = producer .requested_fetch() @@ -2123,8 +2120,8 @@ mod test { let mut producer = TrackProducer::new("test", None); let consumer = producer.consume(); - let mut pending = consumer.fetch(3, None).unwrap(); - assert!(kio::Future::poll(&mut *pending, &kio::Waiter::noop()).is_pending()); + let pending = consumer.fetch(3, None).unwrap(); + assert!(kio::Future::poll(&*pending, &kio::Waiter::noop()).is_pending()); producer.abort(Error::Cancel).unwrap(); assert!(pending.await.is_err());