fix(net): a takeover keeps the open group's head for later readers

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

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<WarmCopy> {
+/// `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<WarmCopy> {
 	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<WarmGroup>,
+	/// 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<HeldGroup>,
+	/// The spliced copy, once it died: its segment keeps ending readers with its error
+	/// unless a takeover follows, which retires it.
+	dead: Option<track::Consumer>,
 	/// 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<str>, u64, Result<track::Info, Error>),
 		Ended(Arc<str>, u64, Result<(), Error>),
+		Held(Arc<str>, group::Producer),
 		Demand(Arc<str>),
 		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(&copy) {
+						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(&copy)) {
 							// 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(&copy, io.head.as_ref());
+						let warm = warm_copy(&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<u8> {
+			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

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<u64>, waiter: &kio::Waiter) -> Poll<group::Producer> {
+		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