diff --git a/rs/kio/src/future.rs b/rs/kio/src/future.rs new file mode 100644 index 0000000000..d05cc65f96 --- /dev/null +++ b/rs/kio/src/future.rs @@ -0,0 +1,146 @@ +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. + /// + /// 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 +/// [`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(&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(&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(&*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/kio/src/producer.rs b/rs/kio/src/producer.rs index d2389a4a76..1d41ca7ecd 100644 --- a/rs/kio/src/producer.rs +++ b/rs/kio/src/producer.rs @@ -148,7 +148,9 @@ impl Producer { } } - fn poll_unused(&self, waiter: &Waiter) -> Poll> { + /// Poll-based variant of [`Self::unused`]: `Ready(Some(()))` when no consumers + /// remain, `Ready(None)` if the channel closed first, else `Pending`. + pub fn poll_unused(&self, waiter: &Waiter) -> Poll> { let mut state = self.state.lock(); if state.closed { return Poll::Ready(None); diff --git a/rs/moq-native/tests/broadcast.rs b/rs/moq-native/tests/broadcast.rs index 32abfb1a6f..afea09fbe8 100644 --- a/rs/moq-native/tests/broadcast.rs +++ b/rs/moq-native/tests/broadcast.rs @@ -229,6 +229,123 @@ 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("video", moq_net::TrackInfo::default().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, async { + bc.track("video").unwrap().fetch(0, None).unwrap().await + }) + .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 4af8229b6e..47d19ea54d 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,7 +10,6 @@ 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>, @@ -36,7 +35,6 @@ 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)?, @@ -63,56 +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 tests { +mod test { use super::*; - fn sample() -> Fetch<'static> { - Fetch { - broadcast: Path::new("room/1"), - track: Cow::Borrowed("video"), - priority: 3, + fn sample() -> FetchOk { + FetchOk { group: 42, - frame_start: 7, + compression: Compression::Deflate, + timescale: Some(Timescale::MICRO), } } - fn roundtrip(version: Version, fetch: &Fetch<'_>) -> Fetch<'static> { + fn roundtrip(version: Version, ok: &FetchOk) -> FetchOk { let mut buf = Vec::new(); - fetch.encode_msg(&mut buf, version).unwrap(); - let mut r = buf.as_slice(); - Fetch::decode_msg(&mut r, version).unwrap() + ok.encode_msg(&mut buf, version).unwrap(); + let mut slice = buf.as_slice(); + FetchOk::decode_msg(&mut slice, version).unwrap() } #[test] - fn frame_start_roundtrips_on_lite05() { + fn fetch_ok_roundtrips_on_lite05() { let got = roundtrip(Version::Lite05Wip, &sample()); assert_eq!(got.group, 42); - assert_eq!(got.frame_start, 7); + assert_eq!(got.compression, Compression::Deflate); + assert_eq!(got.timescale, Some(Timescale::MICRO)); } #[test] - fn frame_start_absent_before_lite05() { - // Lite03/Lite04 don't carry the frame start varint, so it always decodes as 0. - let got = roundtrip(Version::Lite04, &sample()); - assert_eq!(got.group, 42); - assert_eq!(got.frame_start, 0); + 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); + } - // The lite-04 encoding is strictly shorter (no trailing frame start varint). + #[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(); - sample().encode_msg(&mut buf04, Version::Lite04).unwrap(); + msg.encode_msg(&mut buf04, Version::Lite04).unwrap(); let mut buf05 = Vec::new(); - sample().encode_msg(&mut buf05, Version::Lite05Wip).unwrap(); - assert!(buf05.len() > buf04.len(), "lite-05 carries an extra frame start varint"); + 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 5d2b24b84a..62633ec16b 100644 --- a/rs/moq-net/src/lite/publisher.rs +++ b/rs/moq-net/src/lite/publisher.rs @@ -80,6 +80,7 @@ impl Publisher { 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).await; Ok(()) @@ -88,7 +89,7 @@ impl Publisher { tracing::info!("received goaway stream"); Ok(()) } - lite::ControlType::Session | lite::ControlType::Fetch => Err(Error::UnexpectedStream), + lite::ControlType::Session => Err(Error::UnexpectedStream), } } @@ -535,6 +536,148 @@ impl Publisher { stream.writer.finish()?; stream.writer.closed().await } + + pub async fn recv_fetch(&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 track_stats = self.stats.broadcast(&absolute).publisher_track(&track); + + if let Err(err) = Self::run_fetch(&mut stream, &fetch, broadcast, track_stats, self.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 + .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 069f363095..014b8e57be 100644 --- a/rs/moq-net/src/lite/subscriber.rs +++ b/rs/moq-net/src/lite/subscriber.rs @@ -81,6 +81,14 @@ enum SessionOutcome { Error(Error), } +/// The next thing to serve on a track request: a downstream subscription, a +/// one-shot fetch, or nothing left (every consumer dropped). +enum Demand { + Subscription(Option), + Fetch(Result), + Done, +} + impl Subscriber { pub fn new(config: SubscriberConfig) -> Self { // Identity for incoming-hop loop detection. Derived from the local @@ -493,19 +501,70 @@ impl Subscriber { let name = request.name().to_string(); let abs = self.origin.absolute(&path); let track_stats = Arc::new(self.stats.broadcast(&abs).subscriber_track(&name)); + + // Serve whatever demand arrives on this track. A subscription drives the full + // subscribe lifecycle (and consumes the request); fetches are one-shot on + // their own FETCH streams and don't consume the request, so we keep looping. + loop { + let mut waiter = None; + let demand = std::future::poll_fn(|cx| { + let w = waiter.insert(kio::Waiter::new(cx.waker().clone())); + // A fetch is cheap and one-shot, so serve it ahead of subscriptions. + if let Poll::Ready(res) = request.poll_requested_fetch(w) { + return Poll::Ready(Demand::Fetch(res)); + } + if let Poll::Ready(sub) = request.poll_subscription_changed(w) { + return Poll::Ready(Demand::Subscription(sub)); + } + // No demand and no consumers left: stop serving this track. + if request.poll_unused(w).is_ready() { + return Poll::Ready(Demand::Done); + } + Poll::Pending + }) + .await; + + match demand { + // `None` would mean the last subscriber dropped; only a real + // subscription (Some) opens an upstream subscribe. + Demand::Subscription(Some(subscription)) => { + self.run_subscribe_track(path, broadcast, request, track_stats, &name, subscription) + .await; + return; + } + Demand::Subscription(None) => return, + Demand::Fetch(Ok(req)) => { + self.serve_fetch(&path, &request, &track_stats, req).await; + } + Demand::Fetch(Err(err)) => { + tracing::debug!(track = %name, %err, "fetch request aborted"); + return; + } + Demand::Done => return, + } + } + } + + /// Open one upstream SUBSCRIBE for a downstream subscription and run its lifecycle. + #[allow(clippy::too_many_arguments)] + async fn run_subscribe_track( + &mut self, + path: PathOwned, + broadcast: BroadcastDynamic, + request: TrackRequest, + track_stats: Arc, + name: &str, + subscription: crate::Subscription, + ) { // The per-(session, broadcast) `broadcasts` sentinel is taken later, once // the upstream confirms with SUBSCRIBE_OK (see `run_subscribe_session`), so a // sub cancelled before then isn't counted as a feeding session. - let id = self.next_id.fetch_add(1, atomic::Ordering::Relaxed); - // Forward the aggregate of every downstream subscriber's preferences upstream. - // Blocks until the first downstream subscriber registers its preferences. - let subscription = request.subscription_changed().await.unwrap_or_default(); let msg = lite::Subscribe { id, broadcast: path.as_path(), - track: (&name).into(), + track: name.into(), priority: subscription.priority, ordered: subscription.ordered, max_latency: subscription.stale, @@ -537,6 +596,89 @@ impl Subscriber { } } + /// Serve one downstream fetch by issuing a wire FETCH upstream: open a bidi + /// stream, send FETCH, read FETCH_OK, then fill the group (spawned) so the loop + /// can keep serving more demand. Doesn't consume the request. + async fn serve_fetch( + &self, + path: &PathOwned, + request: &TrackRequest, + track_stats: &Arc, + req: crate::FetchRequest, + ) { + let name = request.name().to_string(); + tracing::info!(broadcast = %self.log_path(path), track = %name, group = req.sequence, "fetch started"); + + let mut stream = match Stream::open(&self.session, self.version).await { + Ok(stream) => stream, + Err(err) => { + tracing::warn!(track = %name, %err, "fetch stream open failed"); + return; + } + }; + + let msg = lite::Fetch { + broadcast: path.as_path(), + track: (&name).into(), + priority: req.priority, + group: req.sequence, + frame_start: 0, + }; + if let Err(err) = async { + stream.writer.encode(&lite::ControlType::Fetch).await?; + stream.writer.encode(&msg).await + } + .await + { + stream.writer.abort(&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); + return; + } + }; + + // Make the group available (resolving the downstream fetch) and fill it. + let producer = match request.serve_fetch(req.sequence, info.timescale) { + Ok(producer) => producer, + Err(err) => { + // Already served (a concurrent fetch) or the track closed. + tracing::debug!(track = %name, group = req.sequence, %err, "fetch not served"); + stream.writer.abort(&err); + return; + } + }; + + let mut this = self.clone(); + let track_stats = track_stats.clone(); + web_async::spawn(async move { + let mut producer = producer; + let res = this + .run_group( + &mut stream.reader, + producer.clone(), + track_stats, + info.compression, + info.timescale, + ) + .await; + match res { + Ok(()) => { + let _ = producer.finish(); + } + Err(err) => { + let _ = producer.abort(err); + } + } + }); + } + /// Open the upstream subscribe stream, wait for SUBSCRIBE_OK, then accept the /// pending request (unblocking the downstream subscriber) and run the linger /// lifecycle. The producer is created only after SUBSCRIBE_OK, so a downstream 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-net/src/model/track.rs b/rs/moq-net/src/model/track.rs index 8e0bb676e1..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, }; @@ -156,6 +155,10 @@ struct TrackState { // Active subscriptions. subscriptions: Vec>, + + // Specific groups requested via `fetch` that aren't cached yet, FIFO for the + // producer to serve (see `TrackProducer::requested_fetch`). + fetches: VecDeque, } impl TrackState { @@ -363,6 +366,32 @@ impl TrackState { fn modify(producer: &kio::Producer) -> Result> { producer.write().map_err(|r| r.abort.clone().unwrap_or(Error::Dropped)) } + + /// Insert a fetched group with an explicit timescale, independent of whether the + /// track has been accepted (info may be absent). Used to serve a one-shot FETCH + /// without a full subscription. The group lands in the cache so a waiting + /// [`TrackFetchPending`] resolves via [`Self::poll_get_group`]. + fn insert_fetch_group(&mut self, sequence: u64, timescale: Option) -> Result { + if let Some(err) = &self.abort { + return Err(err.clone()); + } + if let Some(fin) = self.final_sequence + && sequence >= fin + { + return Err(Error::Closed); + } + if !self.duplicates.insert(sequence) { + return Err(Error::Duplicate); + } + + let group = GroupProducer::new_with_timescale(Group { sequence }, timescale); + let now = tokio::time::Instant::now(); + self.max_sequence = Some(self.max_sequence.unwrap_or(0).max(sequence)); + self.groups.push_back(Some((group.clone(), now))); + let cache = self.info.as_ref().map(|i| i.cache).unwrap_or(DEFAULT_CACHE); + self.evict_expired(now, cache); + Ok(group) + } } /// A producer for a track, used to create new groups. @@ -617,6 +646,34 @@ impl TrackProducer { Poll::Ready(Ok(combined)) } + /// Block until a consumer fetches a group that isn't cached, returning the request. + /// + /// The producer serves it by making the group available (a relay issues a wire + /// FETCH then [`create_group`](Self::create_group); an origin already has it + /// cached, so the fetch resolves without ever reaching here). Errors once the + /// track is aborted. + pub async fn requested_fetch(&mut self) -> Result { + kio::wait(|waiter| self.poll_requested_fetch(waiter)).await + } + + pub fn poll_requested_fetch(&mut self, waiter: &kio::Waiter) -> Poll> { + 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() { + return Poll::Ready(Ok(state.fetches.pop_front().unwrap())); + } + match &state.abort { + Some(err) => Poll::Ready(Err(err.clone())), + None => Poll::Pending, + } + }) { + 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, + } + } + fn modify(&self) -> Result> { TrackState::modify(&self.state) } @@ -718,37 +775,61 @@ 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, - }) + })) } - //pub fn fetch(&self) -> TrackFetchPending { - // TODO - //} + /// Fetch a single past group, without holding a live subscription. + /// + /// Returns a [`TrackFetchPending`] that resolves to the [`GroupConsumer`] once + /// the group is available: immediately if it's already cached, otherwise once + /// the producer serves the request (a wire FETCH for a relay). `options` accepts + /// `None`, a [`Fetch`], or `Fetch::default()`. The pending resolves to + /// [`Error::NotFound`] if the group can never exist (past the final sequence). + pub fn fetch(&self, sequence: u64, options: impl Into>) -> Result { + let options = options.into().unwrap_or_default(); + + match self.state.write() { + Ok(mut state) => { + // 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, + priority: options.priority, + }); + } + } + Err(state) => return Err(state.abort.clone().unwrap_or(Error::Dropped)), + }; + + Ok(kio::Pending::new(TrackFetch { + state: self.state.clone(), + sequence, + })) + } 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())) @@ -766,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; @@ -775,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())) @@ -803,16 +885,69 @@ 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 { + /// 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 specific group requested via [`TrackConsumer::fetch`], handed to the producer +/// (or a relay) to serve via [`TrackProducer::requested_fetch`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FetchRequest { + /// The group sequence the consumer wants. + pub sequence: u64, + /// The requested delivery priority. + pub priority: u8, +} + +/// The pollable state of a [`TrackConsumer::fetch`]. +/// +/// 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, +} + +impl kio::Future for TrackFetch { + type Output = Result; + + 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))??; + + // `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)) + } +} + +/// 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. /// /// Created via [`TrackConsumer::subscribe`](crate::TrackConsumer::subscribe), or @@ -1047,9 +1182,49 @@ impl TrackRequest { } } - //pub async fn requested(&self) -> TrackDynamic { - // TODO Used to return the next group requested for FETCH - //} + /// Block until a consumer fetches a group, returning the request to serve. + /// + /// A relay serves it by opening a wire FETCH and calling [`Self::serve_fetch`]. + /// Unlike [`Self::accept`], this doesn't consume the request, so the same track + /// can serve any number of fetches (and later a subscription). Errors once the + /// track is aborted. + pub async fn requested_fetch(&self) -> Result { + kio::wait(|waiter| self.poll_requested_fetch(waiter)).await + } + + pub fn poll_requested_fetch(&self, waiter: &kio::Waiter) -> Poll> { + match self.state.poll(waiter, |state| { + // Only take the `&mut` (which flags the state modified) once there's a + // request to pop, so idle polls don't wake unrelated waiters. + if !state.fetches.is_empty() { + return Poll::Ready(Ok(state.fetches.pop_front().unwrap())); + } + match &state.abort { + Some(err) => Poll::Ready(Err(err.clone())), + None => Poll::Pending, + } + }) { + Poll::Ready(Ok(res)) => Poll::Ready(res), + // A TrackRequest holds the only producer, so the channel can't close here. + Poll::Ready(Err(_)) => Poll::Pending, + Poll::Pending => Poll::Pending, + } + } + + /// Make a fetched group available in the track's cache, resolving the matching + /// [`TrackConsumer::fetch`]. Returns a [`GroupProducer`] to fill from the wire. + /// + /// `timescale` is the scale from the wire FETCH_OK, used for the group's frames. + /// Returns [`Error::Duplicate`] if the group is already present. + pub fn serve_fetch(&self, sequence: u64, timescale: impl Into>) -> Result { + TrackState::modify(&self.state)?.insert_fetch_group(sequence, timescale.into()) + } + + /// Poll for the request becoming unused (every consumer dropped), so a relay can + /// stop serving and drop the request. + pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<()> { + self.state.poll_unused(waiter).map(|_| ()) + } /// Serve the request with the given track, resolving every waiting subscriber. /// @@ -1873,4 +2048,82 @@ mod test { assert!(matches!(producer.append_group(), Err(Error::BoundsExceeded(_)))); } + + #[tokio::test] + async fn fetch_cache_hit() { + let mut producer = TrackProducer::new("test", None); + + // Produce a cached group. + let mut group = producer.append_group().unwrap(); // seq 0 + group.write_frame(bytes::Bytes::from_static(b"hello")).unwrap(); + group.finish().unwrap(); + + // A cached group resolves immediately and never signals the producer. + let consumer = producer.consume(); + let mut g = consumer.fetch(0, None).unwrap().await.unwrap(); + assert_eq!(g.sequence, 0); + assert_eq!(&g.read_frame().await.unwrap().unwrap()[..], b"hello"); + + // Nothing was queued for the producer to serve. + assert!(producer.poll_requested_fetch(&kio::Waiter::noop()).is_pending()); + } + + #[tokio::test] + async fn fetch_miss_signals_producer() { + let mut producer = TrackProducer::new("test", None); + let consumer = producer.consume(); + + // 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 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() + .now_or_never() + .expect("should not block") + .unwrap(); + assert_eq!( + req, + FetchRequest { + sequence: 5, + priority: 7 + } + ); + + // Serve it by producing the group; the fetch then resolves. + let mut group = producer.create_group(Group { sequence: 5 }).unwrap(); + group.write_frame(bytes::Bytes::from_static(b"hi")).unwrap(); + group.finish().unwrap(); + + let mut g = pending.await.unwrap(); + assert_eq!(g.sequence, 5); + assert_eq!(&g.read_frame().await.unwrap().unwrap()[..], b"hi"); + } + + #[tokio::test] + async fn fetch_past_final_not_found() { + let mut producer = TrackProducer::new("test", None); + producer.append_group().unwrap(); // seq 0 + producer.finish().unwrap(); // final_sequence = 1 + + // A group at or past the final sequence can never exist. + let consumer = producer.consume(); + assert!(matches!(consumer.fetch(5, None).unwrap().await, Err(Error::NotFound))); + + // And it doesn't signal the producer. + assert!(producer.poll_requested_fetch(&kio::Waiter::noop()).is_pending()); + } + + #[tokio::test] + async fn fetch_aborts_with_track() { + let mut producer = TrackProducer::new("test", None); + let consumer = producer.consume(); + + 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()); + } } diff --git a/rs/moq-relay/src/web.rs b/rs/moq-relay/src/web.rs index 52e03d18e8..ae938aa175 100644 --- a/rs/moq-relay/src/web.rs +++ b/rs/moq-relay/src/web.rs @@ -528,27 +528,29 @@ async fn serve_fetch( .announced_broadcast("") .await .ok_or(StatusCode::NOT_FOUND)?; - let mut track = async { broadcast.track(&track)?.subscribe(None)?.await } - .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 async { broadcast.track(&track)?.subscribe(None)?.await }.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) => async { broadcast.track(&track)?.fetch(sequence, None)?.await } + .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 {