fix(net): follow-ups to #5133 (deferred FIN on drop, reprioritize open group streams)

rs/moq-net/src/ietf/publisher.rs

@@ -865,12 +865,10 @@ where
 					.await;
 				return Some((result.and(Err(Error::Unsupported)), filled));
 			}
-			if let Some(priority) = update.priority {
-				let mut subscription = serve.track.subscription();
-				subscription.priority = super::priority::from_wire(priority);
-				if let Err(err) = serve.track.update(subscription) {
-					return Some((Err(err), filled));
-				}
+			if let Some(priority) = update.priority
+				&& let Err(err) = serve.set_priority(super::priority::from_wire(priority))
+			{
+				return Some((Err(err), filled));
 			}
 			let acked = async {
 				stream.writer.encode(&ietf::RequestOk::ID).await?;
@@ -2282,6 +2280,8 @@ struct TrackServe<S: crate::transport::poll::Session> {
 	range: ServeRange,
 	timescale: Option<Timescale>,
 	children: kio::Tasks<GroupServe<S>>,
+	/// The subscriber priority, which every group stream follows while it is open.
+	priority: kio::Producer<u8>,
 	/// The track finished: the in-flight group machines drain, then FIN.
 	draining: bool,
 	/// Group streams opened, shared with the group machines.
@@ -2312,6 +2312,7 @@ impl<S: crate::transport::poll::Session> TrackServe<S> {
 		}
 		track.end_at(range.end.map_or(Bound::Unbounded, |end| Bound::Included(end.group)));
 		let datagrams = session.max_datagram_size() > 0;
+		let priority = kio::Producer::new(track.subscription().priority);
 
 		Self {
 			datagrams,
@@ -2322,12 +2323,25 @@ impl<S: crate::transport::poll::Session> TrackServe<S> {
 			range,
 			timescale,
 			children: kio::Tasks::new(),
+			priority,
 			draining: false,
 			opened: Default::default(),
 			end: None,
 		}
 	}
 
+	/// Move the subscriber priority, for the group streams already open as well as later ones.
+	fn set_priority(&mut self, priority: u8) -> Result<(), Error> {
+		let mut subscription = self.track.subscription();
+		subscription.priority = priority;
+		self.track.update(subscription)?;
+		// Only a closed channel refuses the write, and this is the producer, which never closes it.
+		if let Ok(mut current) = self.priority.write() {
+			*current = priority;
+		}
+		Ok(())
+	}
+
 	/// Group streams opened so far.
 	fn opened(&self) -> u64 {
 		self.opened.load(Ordering::Relaxed)
@@ -2418,7 +2432,7 @@ impl<S: crate::transport::poll::Session> TrackServe<S> {
 						GroupServe::new(
 							self.session.clone(),
 							msg,
-							self.track.subscription().priority,
+							self.priority.consume(),
 							group,
 							self.timescale,
 							self.version,
@@ -2502,7 +2516,10 @@ struct GroupServe<S: crate::transport::poll::Session> {
 	/// The subscription's count of opened streams, bumped once this one opens.
 	opened: Arc<AtomicU64>,
 	msg: ietf::GroupHeader,
-	priority: u8,
+	/// The subscription's priority, followed while the stream is open.
+	priority: kio::Consumer<u8>,
+	/// The priority last handed to the stream.
+	applied: u8,
 	group: group::Consumer,
 	timescale: Option<Timescale>,
 	version: Version,
@@ -2546,7 +2563,7 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 	fn new(
 		session: S,
 		msg: ietf::GroupHeader,
-		priority: u8,
+		priority: kio::Consumer<u8>,
 		mut group: group::Consumer,
 		timescale: Option<Timescale>,
 		version: Version,
@@ -2555,10 +2572,12 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 		group.skip_to(slice.skip);
 		group.end_at(slice.until.map_or(Bound::Unbounded, Bound::Excluded));
 		let object_delta = group.index();
+		let applied = *priority.read();
 		Self {
 			session,
 			opened: Default::default(),
 			msg,
+			applied,
 			priority,
 			group,
 			timescale,
@@ -2576,6 +2595,20 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 
 	fn poll_serve(&mut self, waiter: &kio::Waiter) -> Poll<Result<(), Error>> {
 		let mut cx = std::task::Context::from_waker(waiter.waker());
+		// A REQUEST_UPDATE can move the subscriber priority while this stream is open.
+		// Serve and Closed, the only states holding a writer, apply the change; Open
+		// reads the current value once it gets its stream.
+		let applied = self.applied;
+		let mut moved = match self.priority.poll(waiter, |priority| match **priority == applied {
+			true => Poll::Pending,
+			false => Poll::Ready(**priority),
+		}) {
+			Poll::Ready(Ok(priority)) => {
+				self.applied = priority;
+				Some(priority)
+			}
+			_ => None,
+		};
 		loop {
 			match &mut self.state {
 				GroupState::Open => {
@@ -2592,7 +2625,8 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 					};
 					self.opened.fetch_add(1, Ordering::Relaxed);
 					let mut stream = stream;
-					stream.set_priority(self.priority);
+					self.applied = *self.priority.read();
+					stream.set_priority(self.applied);
 
 					let mut writer = Writer::new(stream, self.version);
 					if let Err(err) = writer.buffer(&self.msg) {
@@ -2614,6 +2648,9 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 					batch,
 					batch_pos,
 				} => {
+					if let Some(priority) = moved.take() {
+						writer.set_priority(priority);
+					}
 					// The peer closing first cancels the group.
 					if writer.poll_closed(&mut cx).is_ready() {
 						self.state = GroupState::Done;
@@ -2731,6 +2768,9 @@ impl<S: crate::transport::poll::Session> GroupServe<S> {
 					}
 				}
 				GroupState::Closed { writer } => {
+					if let Some(priority) = moved.take() {
+						writer.set_priority(priority);
+					}
 					// Wait until everything is acknowledged by the peer so we can still
 					// cancel the stream. poll_close releases the stream on completion so
 					// the Drop fallback cannot reset the acknowledged stream.
@@ -2851,7 +2891,7 @@ mod group_priority_test {
 		let mut serve = GroupServe::new(
 			session,
 			msg,
-			200,
+			kio::Producer::new(200).consume(),
 			consumer,
 			Some(Timescale::default()),
 			Version::Draft14,
@@ -2953,7 +2993,7 @@ mod group_priority_test {
 				publisher_priority: 0,
 				flags: Default::default(),
 			},
-			0,
+			kio::Producer::new(0).consume(),
 			group,
 			Some(Timescale::default()),
 			Version::Draft19,
@@ -2998,7 +3038,7 @@ mod group_priority_test {
 				publisher_priority: 0,
 				flags: Default::default(),
 			},
-			0,
+			kio::Producer::new(0).consume(),
 			group,
 			Some(Timescale::default()),
 			Version::Draft19,
@@ -3021,6 +3061,49 @@ mod group_priority_test {
 		assert_eq!(log.resets(), vec![crate::StreamError::Cancel.to_code()]);
 	}
 
+	/// A stream parked on its FIN acknowledgement still follows each priority update.
+	#[tokio::test]
+	async fn unacknowledged_fin_follows_priority_updates() {
+		let session = SinkSession::new(Default::default()).with_unacked_fin();
+		let log = session.log.clone();
+		let track = track::Producer::new(std::sync::Arc::new(crate::broadcast::Info::default()), "test", None);
+		let mut group = track.create_group(group::Info { sequence: 0 }).unwrap();
+		group.write_frame(crate::Timestamp::ZERO, b"done".as_slice()).unwrap();
+		let consumer = group.consume();
+		group.finish().unwrap();
+
+		let priority = kio::Producer::new(200);
+		let mut serve = GroupServe::new(
+			session,
+			ietf::GroupHeader {
+				track_alias: 0,
+				group_id: 0,
+				sub_group_id: 0,
+				publisher_priority: 0,
+				flags: Default::default(),
+			},
+			priority.consume(),
+			consumer,
+			Some(Timescale::default()),
+			Version::Draft19,
+			GroupSlice::default(),
+		);
+		let mut serving = std::pin::pin!(kio::wait(|waiter| serve.poll_serve(waiter)));
+		assert!(
+			futures::poll!(serving.as_mut()).is_pending(),
+			"the FIN is unacknowledged"
+		);
+		assert_eq!(log.priorities(), vec![200]);
+
+		*priority.write().ok().unwrap() = 100;
+		assert!(futures::poll!(serving.as_mut()).is_pending());
+		assert_eq!(log.priorities(), vec![200, 100]);
+
+		*priority.write().ok().unwrap() = 50;
+		assert!(futures::poll!(serving.as_mut()).is_pending());
+		assert_eq!(log.priorities(), vec![200, 100, 50]);
+	}
+
 	/// The final payload remains guarded after its frame has advanced the group cursor.
 	#[tokio::test]
 	async fn blocked_final_transport_chunk_expires_with_the_group() {
@@ -3048,7 +3131,7 @@ mod group_priority_test {
 				publisher_priority: 0,
 				flags: Default::default(),
 			},
-			0,
+			kio::Producer::new(0).consume(),
 			group,
 			Some(Timescale::default()),
 			Version::Draft19,
@@ -4455,7 +4538,7 @@ mod serve_tests {
 			let mut serve = GroupServe::new(
 				session,
 				header(),
-				0,
+				kio::Producer::new(0).consume(),
 				consumer,
 				Some(Timescale::default()),
 				Version::Draft20,
@@ -4580,6 +4663,31 @@ mod serve_tests {
 		}
 	}
 
+	/// A priority update moves the group streams already in flight too, not only the
+	/// ones opened after it.
+	#[tokio::test(start_paused = true)]
+	async fn subscription_update_reprioritizes_open_group_streams() {
+		let h = serve(Version::Draft22);
+		let requester = ScriptedSession::new(Vec::new());
+		let mut run = std::pin::pin!(start_subscription(&h, requester.clone()).await);
+		assert!(!drive(run.as_mut()).await);
+
+		// An unfinished group keeps its stream open.
+		let mut group = h.track.create_group(group::Info { sequence: 0 }).unwrap();
+		group.write_frame(timestamp(), b"open".as_slice()).unwrap();
+		assert!(!drive(run.as_mut()).await);
+		let before = h.log.priorities();
+		assert_eq!(before.len(), 1, "{before:?}");
+		assert_ne!(before[0], 245);
+
+		// REQUEST_UPDATE with SUBSCRIBER_PRIORITY 10.
+		requester.push(&[0x02, 0, 4, 2, 1, 0x20, 10]);
+		assert!(!drive(run.as_mut()).await);
+		assert_eq!(h.track.subscription().unwrap().priority, 245);
+		let after = h.log.priorities();
+		assert_eq!(after.last(), Some(&245), "{after:?}");
+	}
+
 	/// An update we cannot apply is refused NOT_SUPPORTED and ends the subscription with
 	/// UPDATE_FAILED, rather than silently dropping it.
 	#[tokio::test(start_paused = true)]

rs/moq-tokio/src/transport.rs

@@ -534,15 +534,21 @@ impl<S: web_transport_trait::SendStream + 'static> Drop for SendStream<S> {
 					});
 				}
 			}
+			// Both deferred actions fired the watch's interrupt when recorded, so
+			// the salvage task reclaims the stream without waiting on the peer.
 			Some(SendState::Closing(fut)) => {
 				let reset = self.reset.take();
-				if reset.is_some()
+				let finish = std::mem::take(&mut self.finish);
+				if (reset.is_some() || finish)
 					&& let Ok(handle) = tokio::runtime::Handle::try_current()
 				{
 					handle.spawn(async move {
 						let (mut stream, _) = fut.await;
-						if let Some(code) = reset {
-							stream.reset(code);
+						match reset {
+							Some(code) => stream.reset(code),
+							None => {
+								let _ = stream.finish();
+							}
 						}
 					});
 				}
@@ -1172,6 +1178,25 @@ mod tests {
 		assert_eq!(fake.resets.lock().unwrap().as_slice(), &[9]);
 	}
 
+	// An idle closed watch owns the stream, so a FIN recorded against it and then
+	// dropped must still go out instead of the backend's drop reset.
+	#[tokio::test(start_paused = true)]
+	async fn a_deferred_finish_survives_a_drop_mid_watch() {
+		let fake = FakeSend::default();
+		let mut send = SendStream::new(fake.clone());
+		let mut cx = cx();
+
+		assert!(send.poll_closed(&mut cx).is_pending());
+		send.finish().unwrap();
+		drop(send);
+
+		for _ in 0..10 {
+			tokio::task::yield_now().await;
+		}
+		assert!(fake.finished.load(Ordering::SeqCst), "the deferred FIN was lost");
+		assert!(fake.resets.lock().unwrap().is_empty());
+	}
+
 	// The receive-side twin: a reader dropped mid-read still sends the deferred
 	// STOP_SENDING with the real code.
 	#[tokio::test]