fix(net): drain subscriptions awaiting track info

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

@@ -4881,6 +4881,54 @@ mod tests {
 		}
 	}
 
+	/// An in-flight TRACK survives withdrawal, like its pipelined SUBSCRIBE, but
+	/// never accepts a replacement instance's info.
+	#[kio::test]
+	async fn an_in_flight_track_info_keeps_its_instance() {
+		for version in [Version::Lite05, Version::Lite06, Version::Lite07] {
+			for replace in [false, true] {
+				let session = crate::lite::test_transport::ScriptedSession::new(Vec::new());
+				let subscriber = pipelining(crate::driver::Transport::new(session.clone()), version);
+				let (_waiting, mut run) = subscribing(&subscriber, Subscription::default());
+				let route = kio::Producer::new(());
+				run.serve.route = Some(route.consume());
+				let mut selection = crate::model::TestSelection::new();
+				run.serve.validity = selection.validity();
+				let mut fetch = TrackInfoFetch::new(&run.serve);
+				let waiter = kio::Waiter::noop();
+				assert!(fetch.poll_fetch(&run.serve, true, &waiter).is_pending());
+				drop(route);
+				selection.withdraw();
+				if replace {
+					selection.replace();
+				} else {
+					assert!(
+						fetch.poll_fetch(&run.serve, true, &waiter).is_pending(),
+						"{version}: the request already went out"
+					);
+				}
+
+				let mut info = Vec::new();
+				lite::TrackInfo {
+					refresh: None,
+					priority: 0,
+					max_age: None,
+					timescale: Some(Timescale::MILLI),
+					compress: false,
+				}
+				.encode(&mut crate::coding::Encoder::new(&mut info, version.into()), version)
+				.unwrap();
+				session.push(&info);
+				let result = fetch.poll_fetch(&run.serve, true, &waiter);
+				if replace {
+					assert!(matches!(result, Poll::Ready(Err(Error::Unroutable))), "{version}");
+				} else {
+					assert!(matches!(result, Poll::Ready(Ok(_))), "{version}");
+				}
+			}
+		}
+	}
+
 	/// Accepting the TRACK_INFO sends no SUBSCRIBE_UPDATE, even when the publisher keeps
 	/// groups for less than the subscriber waits: the publisher applies its own max age.
 	#[kio::test]
@@ -8595,7 +8643,10 @@ impl kio::Task for TrackServeRun {
 					}
 					// TRACK opens first: its stream fails fast without credit, where a fetch's
 					// waits for it, so fetches opened ahead could starve their own prerequisite.
-					let fetched = info.poll_fetch(&self.serve, waiter);
+					let subscribed = pipelined
+						.as_ref()
+						.is_some_and(|subscribing| matches!(subscribing.state, PipelinedState::Open(_)));
+					let fetched = info.poll_fetch(&self.serve, subscribed, waiter);
 					let failed = match pipelined {
 						Some(subscribing) => subscribing.poll(waiter),
 						None => Poll::Pending,
@@ -8762,15 +8813,19 @@ impl TrackInfoFetch {
 	fn poll_fetch(
 		&mut self,
 		serve: &TrackServe,
+		subscribed: bool,
 		waiter: &kio::Waiter,
 	) -> Poll<Result<(track::Info, HeldTrack), Error>> {
 		#[every_poll]
-		self.validity.check(waiter)?;
-		// A withdrawn route fails an info still owed. One already delivered stays usable, like
-		// one read before the withdrawal: a SUBSCRIBE sent alongside may already be served.
-		#[every_poll]
-		if !matches!(self.state, TrackInfoState::Read(_)) && withdrawn(serve, waiter) {
-			return Poll::Ready(Err(Error::Unroutable));
+		if subscribed && matches!(self.state, TrackInfoState::Read(..) | TrackInfoState::Properties(..)) {
+			// The request is on the wire: withdrawal leaves it to its answer, just like
+			// the SUBSCRIBE sent alongside it. A replacement still invalidates that answer.
+			self.validity.check_unanswered(waiter)?;
+		} else {
+			self.validity.check(waiter)?;
+			if withdrawn(serve, waiter) {
+				return Poll::Ready(Err(Error::Unroutable));
+			}
 		}
 
 		let mut stream: Stream<Transport, Version> =
@@ -8788,14 +8843,10 @@ impl TrackInfoFetch {
 			if let Poll::Ready(err) = self.closed.poll_closed(waiter) {
 				return Poll::Ready(Err(crate::lite::error::from_transport(err)));
 			}
-			let read = match serve.subscriber.peer_setup.poll_ietf(waiter) {
+			match serve.subscriber.peer_setup.poll_ietf(waiter) {
 				Poll::Ready(ietf) => stream.reader.poll_decode::<lite::TrackInfo>(waiter).map_ok(|info| (ietf, info)),
 				Poll::Pending => Poll::Pending,
-			};
-			if read.is_pending() && withdrawn(serve, waiter) {
-				return Poll::Ready(Err(Error::Unroutable));
 			}
-			read
 		})?;
 
 		// Publisher Max Age rides on the wire, so the local retention

rs/moq-net/src/model/front.rs

@@ -123,8 +123,8 @@ pub(super) enum Action {
 	/// Request the path through `route`; feed the outcome back as
 	/// [`Event::Resolved`].
 	Request { route: u64 },
-	/// Drop the driver's handle on a source.
-	Detach { source: u64 },
+	/// Drop the driver's source handle, keeping its pending track queries when `drain` is true.
+	Detach { source: u64, drain: bool },
 	/// Resolve the requesters parked on the front: it has a source.
 	Resolve,
 	/// Ask `source` for its copy of `track`; feed the answer back as
@@ -395,7 +395,7 @@ impl Front {
 		let Some(upstream) = self.upstream.filter(|upstream| upstream.route == route) else {
 			// A source for a request we no longer want: let it go.
 			if let Ok(source) = result {
-				actions.push(Action::Detach { source });
+				actions.push(Action::Detach { source, drain: false });
 			}
 			return;
 		};
@@ -423,11 +423,14 @@ impl Front {
 
 	fn attach(&mut self, source: u64, candidate: Candidate, actions: &mut Vec<Action>) {
 		if let Some((old, _)) = self.serving.take() {
-			actions.push(Action::Detach { source: old });
+			actions.push(Action::Detach {
+				source: old,
+				drain: false,
+			});
 		}
 		// Every copy is from an older source, including one a track still drains from a
 		// source the previous one replaced.
-		self.drain_copies();
+		self.drain_copies(false);
 		self.serving = Some((source, candidate));
 		self.serving_closing = false;
 		if !self.resolved {
@@ -459,22 +462,23 @@ impl Front {
 		self.excluded.insert(candidate.route);
 		self.serving = None;
 		self.serving_closing = false;
-		actions.push(Action::Detach { source });
-		self.drain_copies();
+		actions.push(Action::Detach { source, drain: true });
+		self.drain_copies(true);
 		actions.push(Action::Reselect);
 	}
 
 	/// The serving source is going: a spliced copy keeps feeding its track (readers read it
 	/// until a replacement's is spliced, or it runs out), and one still being asked for is
-	/// dropped.
-	fn drain_copies(&mut self) {
+	/// kept while that source can still answer, until a replacement attaches.
+	fn drain_copies(&mut self, keep_queries: bool) {
 		for track in self.tracks.values_mut() {
 			match track.state {
 				TrackState::Spliced { source } => {
 					track.draining = Some(source);
 					track.state = TrackState::Idle;
 				}
-				TrackState::Querying { .. } => track.state = TrackState::Idle,
+				TrackState::Querying { .. } if !keep_queries => track.state = TrackState::Idle,
+				TrackState::Querying { .. } => {}
 				TrackState::Idle | TrackState::Parked { .. } => {}
 			}
 		}
@@ -759,8 +763,8 @@ mod tests {
 		assert_eq!(front.next_deadline(), scanned.first().map(|(at, _)| *at));
 	}
 
-	/// A front serving `source` through `candidate`, with `video` spliced and read.
-	fn serving(candidate: Candidate, source: u64) -> Front {
+	/// A front querying `source` through `candidate` for the read track `video`.
+	fn querying(candidate: Candidate, source: u64) -> Front {
 		let mut front = Front::new(LINGER);
 		assert_actions(
 			front.step(Event::Selected {
@@ -794,6 +798,12 @@ mod tests {
 				Action::Arm { at: None },
 			],
 		);
+		front
+	}
+
+	/// A front serving `source` through `candidate`, with `video` spliced and read.
+	fn serving(candidate: Candidate, source: u64) -> Front {
+		let mut front = querying(candidate, source);
 		assert_actions(
 			front.step(Event::TrackInfo {
 				track: name("video"),
@@ -857,7 +867,10 @@ mod tests {
 				result: Ok(200),
 			}),
 			&[
-				Action::Detach { source: 100 },
+				Action::Detach {
+					source: 100,
+					drain: false,
+				},
 				Action::Query {
 					track: name("video"),
 					source: 200,
@@ -867,6 +880,28 @@ mod tests {
 		assert_eq!(front.serving, Some((200, remote(2))));
 	}
 
+	#[test]
+	fn a_closing_source_answers_while_another_route_is_resolving() {
+		let mut front = querying(remote(1), 100);
+		front.step(Event::SourceClosed { source: 100 });
+		front.step(Event::Selected {
+			best: Some(remote(2)),
+			serving_closing: false,
+		});
+		assert_actions(
+			front.step(Event::TrackInfo {
+				track: name("video"),
+				source: 100,
+				closing: true,
+				result: Ok(info()),
+			}),
+			&[Action::Splice {
+				track: name("video"),
+				source: 100,
+			}],
+		);
+	}
+
 	#[test]
 	fn a_stale_resolution_is_let_go() {
 		let mut front = serving(remote(1), 100);
@@ -887,7 +922,10 @@ mod tests {
 				route: 2,
 				result: Ok(200),
 			}),
-			&[Action::Detach { source: 200 }],
+			&[Action::Detach {
+				source: 200,
+				drain: false,
+			}],
 		);
 	}
 
@@ -896,7 +934,13 @@ mod tests {
 		let mut front = serving(remote(1), 100);
 		assert_actions(
 			front.step(Event::SourceClosed { source: 100 }),
-			&[Action::Detach { source: 100 }, Action::Reselect],
+			&[
+				Action::Detach {
+					source: 100,
+					drain: true,
+				},
+				Action::Reselect,
+			],
 		);
 		assert_eq!(front.excluded_routes(), &HashSet::from([1]));
 		assert_actions(
@@ -1350,7 +1394,13 @@ mod tests {
 		front.step(Event::Forgotten { track: name("video") });
 		assert_actions(
 			front.step(Event::SourceClosed { source: 100 }),
-			&[Action::Detach { source: 100 }, Action::Reselect],
+			&[
+				Action::Detach {
+					source: 100,
+					drain: true,
+				},
+				Action::Reselect,
+			],
 		);
 		assert_actions(front.step(Event::Unheld), &[]);
 		assert_actions(
@@ -1425,7 +1475,10 @@ mod tests {
 				route: 2,
 				result: Ok(200),
 			}),
-			&[Action::Detach { source: 100 }],
+			&[Action::Detach {
+				source: 100,
+				drain: false,
+			}],
 		);
 		let t1 = t0 + LINGER;
 		front.step(Event::Unwanted {
@@ -1574,7 +1627,13 @@ mod tests {
 		let mut front = serving(local(1), 100);
 		assert_actions(
 			front.step(Event::SourceClosed { source: 100 }),
-			&[Action::Detach { source: 100 }, Action::Reselect],
+			&[
+				Action::Detach {
+					source: 100,
+					drain: true,
+				},
+				Action::Reselect,
+			],
 		);
 		assert_actions(
 			front.step(Event::Selected {
@@ -1760,7 +1819,7 @@ mod tests {
 				assert!(splices <= 1);
 				// A Detach names a source the front no longer serves from.
 				for action in &actions {
-					if let Action::Detach { source } = action {
+					if let Action::Detach { source, .. } = action {
 						assert_ne!(next.serving.map(|(s, _)| s), Some(*source));
 					}
 				}
@@ -1808,7 +1867,10 @@ mod tests {
 				result: Ok(300),
 			}),
 			&[
-				Action::Detach { source: 200 },
+				Action::Detach {
+					source: 200,
+					drain: false,
+				},
 				Action::Query {
 					track: name("video"),
 					source: 300,