diff --git a/rs/moq-net/src/ietf/publisher.rs b/rs/moq-net/src/ietf/publisher.rs index 5d19955eb4..b780d2bbe9 100644 --- a/rs/moq-net/src/ietf/publisher.rs +++ b/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 { range: ServeRange, timescale: Option, children: kio::Tasks>, + /// The subscriber priority, which every group stream follows while it is open. + priority: kio::Producer, /// 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 TrackServe { } 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 TrackServe { 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 TrackServe { GroupServe::new( self.session.clone(), msg, - self.track.subscription().priority, + self.priority.consume(), group, self.timescale, self.version, @@ -2502,7 +2516,10 @@ struct GroupServe { /// The subscription's count of opened streams, bumped once this one opens. opened: Arc, msg: ietf::GroupHeader, - priority: u8, + /// The subscription's priority, followed while the stream is open. + priority: kio::Consumer, + /// The priority last handed to the stream. + applied: u8, group: group::Consumer, timescale: Option, version: Version, @@ -2546,7 +2563,7 @@ impl GroupServe { fn new( session: S, msg: ietf::GroupHeader, - priority: u8, + priority: kio::Consumer, mut group: group::Consumer, timescale: Option, version: Version, @@ -2555,10 +2572,12 @@ impl GroupServe { 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 GroupServe { fn poll_serve(&mut self, waiter: &kio::Waiter) -> Poll> { 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 GroupServe { }; 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 GroupServe { 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 GroupServe { } } 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)] diff --git a/rs/moq-tokio/src/transport.rs b/rs/moq-tokio/src/transport.rs index 31d7b34d6b..01814ee444 100644 --- a/rs/moq-tokio/src/transport.rs +++ b/rs/moq-tokio/src/transport.rs @@ -534,15 +534,21 @@ impl Drop for SendStream { }); } } + // 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]