diff --git a/rs/moq-net/src/model/group.rs b/rs/moq-net/src/model/group.rs index 544546434b..bb725dcba1 100644 --- a/rs/moq-net/src/model/group.rs +++ b/rs/moq-net/src/model/group.rs @@ -770,6 +770,20 @@ impl Producer { /// their buffers in memory forever; consumers that haven't drained yet surface the /// abort error instead of the leftover cache. pub fn abort(self, err: Error) -> Result<()> { + self.close_aborted(err) + } + + /// Abort with `err` unless another handle still owns the group, so a retained clone + /// lets go without the unfinished-drop warning and never cuts off a live writer. + /// Sound without a lock: only an existing handle can mint another. + pub(crate) fn abort_if_last(&self, err: Error) -> Result<()> { + if Arc::strong_count(&self.alive) > 1 { + return Ok(()); + } + self.close_aborted(err) + } + + fn close_aborted(&self, err: Error) -> Result<()> { let mut guard = modify(&self.state)?; guard.abort = Some(err); self.alive.aborted.store(true, Ordering::Release); diff --git a/rs/moq-net/src/model/origin.rs b/rs/moq-net/src/model/origin.rs index 4f8c1ae340..94aef76377 100644 --- a/rs/moq-net/src/model/origin.rs +++ b/rs/moq-net/src/model/origin.rs @@ -2233,17 +2233,37 @@ impl Drop for WarmGroup { } } +/// A spliced copy's newest group, held so its head outlives the copy's route (see +/// [`TrackIo::held`]). Borrowed from the route, which may still be writing it for other +/// readers, so it aborts on drop only once it is the last handle: a dead route never +/// finishes it. +struct HeldGroup(group::Producer); + +impl Drop for HeldGroup { + fn drop(&mut self) { + if !self.0.is_finished() { + let _ = self.0.abort_if_last(Error::Cancel); + } + } +} + /// Cache what `source` delivered on a new local track the origin owns: its complete /// groups, and its open live edge rebuilt from the frames already delivered. /// /// `head` is the previous park's edge. A copy spliced after a warm cache continues its /// edge group from the next frame, so its copy of that group lacks the head, which -/// `head` supplies. -fn warm_copy(source: &track::Consumer, head: Option<&WarmGroup>) -> Option { +/// `head` supplies. `held` is the copy's newest group as [`TrackIo::held`] kept it: a +/// copy that died tore its open groups out of its own cache, but not out of the held one. +fn warm_copy(source: &track::Consumer, head: Option<&WarmGroup>, held: Option<&group::Producer>) -> Option { let info = source.cached_info()?; let mut track = track::Producer::new(Arc::new(source.broadcast().clone()), source.name(), info); let head = head.map(|head| &head.0); - let groups = source.cached_groups(); + let mut groups = source.cached_groups(); + if let Some(held) = held + && !groups.iter().any(|(group, _)| group.sequence == held.sequence) + { + groups.push((held.clone(), true)); + } // Not `source.latest()`: datagrams share the sequence counter and can run past it. let latest = groups .iter() @@ -2403,6 +2423,14 @@ struct TrackIo { /// The last warm copy's newest group, outliving it so the copy spliced after it /// can continue the group (see [`WarmGroup`]). Released at the next park. head: Option, + /// The spliced copy's newest group. A route that ends abruptly releases its open + /// groups, and the takeover after it asks the next copy only for the frames past + /// the break, so without this handle no route holds the group's head for a reader + /// that arrives later. See [`TrackIo::retire`]. + held: Option, + /// The spliced copy, once it died: its segment keeps ending readers with its error + /// unless a takeover follows, which retires it. + dead: Option, /// Whether the track had a reader as of the last demand edge. used: bool, } @@ -2416,6 +2444,39 @@ impl TrackIo { self.copy = None; self.warm = None; self.head = None; + self.held = None; + self.dead = None; + } + + /// Keep the spliced copy as dead. The front only refreshes `held` between wakes, so + /// first catch a newest group the copy surfaced since, while its cache still has it. + fn bury(&mut self) { + let Some((_, copy)) = self.copy.take() else { return }; + let after = self.held.as_ref().map(|held| held.0.sequence); + if let Poll::Ready(group) = copy.poll_latest_group(after, &kio::Waiter::noop()) { + self.held = Some(HeldGroup(group)); + } + self.dead = Some(copy); + } + + /// Retire `copy` as a takeover replaces it. Mid-group, park its delivery warm, open + /// group included, so a reader arriving later still gets the group's head and the + /// takeover continues it. Between groups its own cache keeps everything finished, so + /// there is nothing to do. + fn retire(&mut self, copy: &track::Consumer) -> Result<(), Error> { + let Some(held) = self.held.take() else { + return Ok(()); + }; + if held.0.is_finished() { + return Ok(()); + } + let Some(warm) = warm_copy(copy, self.head.as_ref(), Some(&held.0)) else { + return Ok(()); + }; + self.head = None; + self.resume.park(&warm.track)?; + self.warm = Some(warm); + Ok(()) } } @@ -2441,6 +2502,7 @@ async fn run_front(task: FrontTask) { SourceClosed(u64), Info(Arc, u64, Result), Ended(Arc, u64, Result<(), Error>), + Held(Arc, group::Producer), Demand(Arc), Deadline, Table, @@ -2592,7 +2654,7 @@ async fn run_front(task: FrontTask) { // spliced until a replacement resumes past them. for io in tracks.values_mut() { if io.copy.as_ref().is_some_and(|(s, _)| *s == source) { - io.copy = None; + io.bury(); } if io.query.as_ref().is_some_and(|(s, ..)| *s == source) { io.query = None; @@ -2634,7 +2696,12 @@ async fn run_front(task: FrontTask) { if staged != source { continue; } - if let Err(err) = io.resume.takeover(©) { + let outgoing = io.copy.take().map(|(_, copy)| copy).or_else(|| io.dead.take()); + let retired = match outgoing { + Some(outgoing) => io.retire(&outgoing), + None => Ok(()), + }; + if let Err(err) = retired.and_then(|()| io.resume.takeover(©)) { // Closed means the logical track already ended. Anything // else is a boundary bug; abort rather than strand // subscribers on a track nobody serves. @@ -2648,6 +2715,7 @@ async fn run_front(task: FrontTask) { // The new segment has produced nothing yet: this is the // edge the copy is asked to advance. io.edge = io.resume.resume_position(); + io.held = None; io.copy = Some((source, copy)); } Action::Park { track: name } => { @@ -2656,9 +2724,11 @@ async fn run_front(task: FrontTask) { // Drop the source copy so its producer goes idle at once; keep // the groups it delivered on a local track so resume stays // spliced until the linger expires. - let warm = warm_copy(©, io.head.as_ref()); + let warm = warm_copy(©, io.head.as_ref(), io.held.as_ref().map(|held| &held.0)); drop(copy); io.head = None; + io.held = None; + io.dead = None; let parked = match &warm { Some(warm) => io.resume.park(&warm.track), None => io.resume.release(), @@ -2773,6 +2843,12 @@ async fn run_front(task: FrontTask) { { return Poll::Ready(Step::Ended(name.clone(), *source, result)); } + if let Some((_, copy)) = &io.copy + && let Poll::Ready(group) = + copy.poll_latest_group(io.held.as_ref().map(|held| held.0.sequence), waiter) + { + return Poll::Ready(Step::Held(name.clone(), group)); + } // Watch the demand edge in whichever direction is unmet. let edge = match io.used { true => io.resume.poll_unused(waiter), @@ -2801,6 +2877,8 @@ async fn run_front(task: FrontTask) { edge: None, warm: None, head: None, + held: None, + dead: None, used: false, }, ); @@ -2860,7 +2938,7 @@ async fn run_front(task: FrontTask) { Step::Ended(name, source, result) => { let closing = sources.get(&source).is_some_and(|s| s.is_closing()); let Some(io) = tracks.get_mut(&name) else { continue }; - io.copy = None; + io.bury(); let delivered = io.resume.resume_position() != io.edge; Event::TrackEnded { track: name, @@ -2870,6 +2948,12 @@ async fn run_front(task: FrontTask) { delivered, } } + Step::Held(name, group) => { + if let Some(io) = tracks.get_mut(&name) { + io.held = Some(HeldGroup(group)); + } + continue; + } Step::Demand(name) => { let Some(io) = tracks.get_mut(&name) else { continue }; io.used = io.resume.is_used(); @@ -5992,6 +6076,119 @@ mod tests { } } + /// A route that ends abruptly releases its open groups, and the takeover asks the + /// next route only for the frames past the break. The front keeps the head it + /// spliced, so a reader arriving afterwards still gets the whole group: whether the + /// old route dies before the new one takes over (a reconnect) or after (a cheaper + /// route preempting it). + #[tokio::test] + async fn a_takeover_keeps_the_open_group_head_for_later_readers() { + tokio::time::pause(); + for dies_first in [true, false] { + takeover_keeps_the_open_group_head(dies_first).await; + } + } + + async fn takeover_keeps_the_open_group_head(dies_first: bool) { + let producer = origin(1).produce(); + let first_server = producer + .dynamic("room/alice", Route::default().with_hops(hops(&[10])).with_cost(5)) + .unwrap(); + let pending = producer.consume().request_broadcast("room/alice"); + let upstream = broadcast::Info::new().produce(); + let mut dynamic = upstream.dynamic(); + queued(&first_server).await.accept(&upstream); + let resolved = pending.await.unwrap(); + + async fn read(group: &mut group::Consumer) -> Vec { + let frame = tokio::time::timeout(Duration::from_secs(1), group.read_frame()) + .await + .expect("frame") + .unwrap() + .expect("group ended"); + frame.payload.to_vec() + } + + // A reader is mid-way through the open group when the route changes. + let track = resolved.track("log").unwrap(); + let subscribing = tokio::spawn(async move { track.subscribe(None).await }); + let source = dynamic.requested_track().await.unwrap().accept(None); + let mut group = source.create_group(0u64.into()).unwrap(); + group.write_frame(crate::Timestamp::ZERO, b"a".as_ref()).unwrap(); + group.write_frame(crate::Timestamp::ZERO, b"b".as_ref()).unwrap(); + let mut subscription = subscribing.await.unwrap().unwrap(); + let mut reading = subscription.recv_group().await.unwrap().unwrap(); + assert_eq!(read(&mut reading).await, b"a"); + assert_eq!(read(&mut reading).await, b"b"); + + // Either the first route's session ends and reconnects, announcing the same route + // afresh while the old one is still held, or a cheaper route preempts the live one. + let mut first = Some((group, source, upstream, dynamic)); + if dies_first { + drop(first.take()); + } + let cost = if dies_first { 5 } else { 0 }; + let replacement_server = producer + .dynamic("room/alice", Route::default().with_hops(hops(&[10])).with_cost(cost)) + .unwrap(); + let replacement = broadcast::Info::new().produce(); + let mut replacement_dynamic = replacement.dynamic(); + queued(&replacement_server).await.accept(&replacement); + let mut next = tokio::time::timeout(Duration::from_secs(1), replacement_dynamic.requested_track()) + .await + .expect("the front asked the replacement") + .unwrap() + .resolving_start() + .accept(None); + next.start_at(0).unwrap(); + let mut next_group = next.create_group(0u64.into()).unwrap(); + next_group.start_at(2).unwrap(); + next_group.write_frame(crate::Timestamp::ZERO, b"c".as_ref()).unwrap(); + assert_eq!(read(&mut reading).await, b"c", "dies_first={dies_first}"); + + // A preempted route still owns its open group: its writer and other readers carry on. + if let Some((group, ..)) = &mut first { + let mut independent = group.consume(); + group.write_frame(crate::Timestamp::ZERO, b"x".as_ref()).unwrap(); + for expect in [b"a", b"b", b"x"] { + assert_eq!(read(&mut independent).await, expect); + } + group.finish().unwrap(); + } + drop(first); + + let track = resolved.track("log").unwrap(); + let mut later = tokio::time::timeout(Duration::from_secs(1), track.subscribe(None)) + .await + .expect("subscribe") + .unwrap(); + let mut late = tokio::time::timeout(Duration::from_secs(1), later.recv_group()) + .await + .unwrap_or_else(|_| panic!("dies_first={dies_first}: the later reader never got the group")) + .unwrap() + .expect("track ended"); + assert_eq!(late.sequence, 0); + for expect in [b"a", b"b", b"c"] { + assert_eq!(read(&mut late).await, expect, "dies_first={dies_first}"); + } + + // The continuation still flows to both. + next_group.write_frame(crate::Timestamp::ZERO, b"d".as_ref()).unwrap(); + assert_eq!(read(&mut reading).await, b"d"); + assert_eq!(read(&mut late).await, b"d"); + next_group.finish().unwrap(); + drop(( + reading, + late, + subscription, + later, + next, + replacement, + replacement_server, + first_server, + )); + } + /// A returning reader is handed the parked open group spliced onto the source's /// continuation, which starts past the warm head. The group's end stays unknown /// until the source finishes it, even before a frame of it is read: readers drop a diff --git a/rs/moq-net/src/model/track.rs b/rs/moq-net/src/model/track.rs index f63d786dcc..097b162a4b 100644 --- a/rs/moq-net/src/model/track.rs +++ b/rs/moq-net/src/model/track.rs @@ -2647,6 +2647,31 @@ impl Consumer { })) } + /// Poll for the newest cached group once its sequence passes `after`, as a producer + /// handle: holding one keeps its frames when the track's own producer tears down + /// abruptly, which otherwise releases every unfinished group. Pending once closed. + pub(crate) fn poll_latest_group(&self, after: Option, waiter: &kio::Waiter) -> Poll { + let ConsumerKind::Plain(state) = &self.inner else { + return Poll::Pending; + }; + let res = state.poll(waiter, |state| { + match state + .latest_group + .filter(|latest| after.is_none_or(|after| *latest > after)) + { + Some(latest) => match state.lookup.get(&latest) { + Some(slot) => Poll::Ready(slot.group.clone()), + None => Poll::Pending, + }, + None => Poll::Pending, + } + }); + match res { + Poll::Ready(Ok(group)) => Poll::Ready(group), + _ => Poll::Pending, + } + } + /// Poll for a cached group by sequence, parking `waiter` until it lands. /// /// `Ready(None)` once it can never arrive: the track ended below the sequence, or