fix(net): drain IETF serves before session close

doc/lib/rs/moq-net.md

@@ -65,11 +65,11 @@ second, for finished tracks to deliver their last groups and FIN, returning
 `Error::Timeout` if it gave up. Finish or abort live tracks before calling it.
 Both protocols withdraw the session's announcements and wait for delivery
 under that same deadline. IETF drafts 14 through 16 send withdrawals without
-waiting, and IETF media streams are not drained yet. Dropping the driver
-cancels the session. `moq-tokio` and `moq-wasm` drive sessions for their
-callers. A deliberate local session close ends received tracks cleanly after
-their delivered groups. A peer close ends received tracks and open group
-readers with the session error, preserving its close code.
+waiting. Dropping the driver cancels the session. `moq-tokio` and `moq-wasm`
+drive sessions for their callers. A deliberate local session close ends
+received tracks cleanly after their delivered groups. A peer close ends
+received tracks and open group readers with the session error, preserving its
+close code.
 
 `origin::Producer::new` returns a driver with the same `time::Driver`
 interface. It calls `cache::Pool::gc(now)` after each poll and folds the next

quest/m1/README.md

@@ -36,7 +36,6 @@ QUIC studies there on that rule.
 - [Binding audio delay](/quest/m1/binding-surface.md) - moq-ffi and every wrapper configure and observe audio playout delay
 - [FFI shape](/quest/m1/ffi-shape/README.md) - the bindings mirror Rust's layers: net at the root, then media, json, flate, audio, and video namespaces built from the handle below
 - [moq-net owns its transport seam](/quest/m1/transport-seam.md) - moq-net names only its own transport traits, so a web-transport-trait or qmux major bump is a patch for it
-- [IETF drain before close](/quest/m1/ietf-drain-before-close.md) - moq-transport sessions deliver finished tracks before a graceful close, as moq-lite does
 - [Close waits for the tail](/quest/m1/close-tail.md) - on lite-07, `close()` returns `Ok` only after each subscriber FINs its Subscribe Stream, having read the track to its end
 - [moq play drain tail](/quest/m1/play-drain-tail.md) - retired renditions and finite tracks play their last 10 ms of audio
 - [A moq-relay release refuses --cluster-mesh](/quest/m1/relay-mesh-refusal-release.md) - the release that lets the cluster flag shims go

quest/m1/close-tail.md

@@ -38,7 +38,7 @@ Decided (maintainer, 2026-09-30):
   so close keeps today's ack-based drain rather than time out every close.
   #4508's repro is on lite-05, and that path stays as it is: this quest closes
   the issue by fixing the close on the version that can carry the FIN rule.
-  IETF sessions stay with [IETF drain before close](/quest/m1/ietf-drain-before-close.md).
+  IETF sessions keep their ack-based drain.
 - Both subscribers change: moq-net and `@moq/net`.
 
 Update `drafts/draft-lcurley-moq-lite.md` (the Subscribe section gains the
@@ -65,5 +65,4 @@ after reading the track's end, and a publisher closing gracefully waits for it.
 
 ## Related
 
-- [IETF drain before close](/quest/m1/ietf-drain-before-close.md) - the same drain for moq-transport sessions
 - [SUBSCRIBE_DROP](/quest/m1/subscribe-drop.md) - replaces the lite-07 tail accounting the FIN rule waits on

rs/moq-net/src/driver.rs

@@ -94,11 +94,11 @@ impl<S: crate::transport::poll::Session> Protocol<S> {
 	}
 
 	/// Whether the protocol owes the peer no queued data, so a draining close can
-	/// proceed. IETF currently tracks namespace withdrawals only.
+	/// proceed.
 	fn drained(&self) -> bool {
 		match self {
 			Self::Lite(driver) => driver.drained(),
-			Self::Ietf(driver) => driver.withdrawal.drained(),
+			Self::Ietf(driver) => driver.drained(),
 		}
 	}
 }

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

@@ -5,7 +5,7 @@ use std::{
 	ops::Bound,
 	sync::{
 		Arc,
-		atomic::{AtomicU64, Ordering},
+		atomic::{AtomicU64, AtomicUsize, Ordering},
 	},
 	task::{Poll, ready},
 	time::Duration,
@@ -325,6 +325,16 @@ pub(super) struct Publisher<S: crate::transport::poll::Session> {
 	// Shared across request handlers; None marks a dispatched subscription still resolving.
 	joins: kio::Shared<HashMap<RequestId, Option<Joined>>>,
 	version: Version,
+	// Dispatched finite serves, including those not yet polled.
+	pub(super) owed: Arc<AtomicUsize>,
+}
+
+struct Serve(Arc<AtomicUsize>);
+
+impl Drop for Serve {
+	fn drop(&mut self) {
+		self.0.fetch_sub(1, Ordering::Relaxed);
+	}
 }
 
 /// The snapshot a joining FETCH inherits from its subscription.
@@ -381,6 +391,7 @@ where
 			peer_setup,
 			joins: Default::default(),
 			version,
+			owed: Default::default(),
 		}
 	}
 
@@ -465,6 +476,11 @@ where
 		stream: Stream<S, Version>,
 	) -> Result<MaybeSendBox<'static, ()>, Error> {
 		let this = self.clone();
+		// Count at dispatch so close cannot pass a request waiting for its first poll.
+		let serve = matches!(id, ietf::Subscribe::ID | ietf::Fetch::ID).then(|| {
+			self.owed.fetch_add(1, Ordering::Relaxed);
+			Serve(self.owed.clone())
+		});
 		let task = match id {
 			ietf::Subscribe::ID => {
 				let msg = ietf::Subscribe::decode_msg(&mut data, this.version)?;
@@ -474,6 +490,7 @@ where
 				tracing::debug!(message = ?msg, "received subscribe");
 				let task = this.run_subscribe_stream(stream, msg);
 				async move {
+					let _serve = serve;
 					if let Err(err) = task.await {
 						tracing::debug!(%err, "subscribe stream error");
 					}
@@ -487,6 +504,7 @@ where
 				}
 				tracing::debug!(message = ?msg, "received fetch");
 				async move {
+					let _serve = serve;
 					if let Err(err) = this.run_fetch_stream(stream, msg).await {
 						tracing::debug!(%err, "fetch stream error");
 					}
@@ -3187,6 +3205,64 @@ mod serve_tests {
 		}
 	}
 
+	/// A close must count requests before their first poll and release them even
+	/// when a dispatched task is cancelled or its message is refused.
+	#[tokio::test(start_paused = true)]
+	async fn drain_counts_dispatched_subscribe_and_fetch() {
+		for version in [Version::Draft14, Version::Draft19, Version::Draft20] {
+			let h = serve(version);
+			let mut subscribe_body = bytes::BytesMut::new();
+			subscribe(Filter::NextObject, None)
+				.encode_msg(&mut subscribe_body, version)
+				.unwrap();
+			let mut fetch_body = bytes::BytesMut::new();
+			ietf::Fetch {
+				request_id: FETCH_ID,
+				subscriber_priority: 128,
+				group_order: GroupOrder::Ascending,
+				// Draft-20 dropped the Fetch Type tag, leaving only the filtered form.
+				fetch_type: match version {
+					Version::Draft20 => FetchType::Filtered {
+						namespace: crate::Path::new("missing"),
+						track: "video".into(),
+						filter: Filter::Unfiltered,
+					},
+					_ => FetchType::Standalone {
+						namespace: crate::Path::new("missing"),
+						track: "video".into(),
+						start: Location { group: 0, object: 0 },
+						end: Location { group: 0, object: 1 },
+					},
+				},
+				range_filters: false,
+				fill_timeout: false,
+			}
+			.encode_msg(&mut fetch_body, version)
+			.unwrap();
+
+			for (id, body) in [
+				(ietf::Subscribe::ID, subscribe_body.freeze()),
+				(ietf::Fetch::ID, fetch_body.freeze()),
+			] {
+				let stream = Stream::open(&mut h.session.clone(), version).await.unwrap();
+				let task = h.publisher.handle_stream(id, body.clone(), stream).unwrap();
+				assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 1, "count before polling");
+				drop(task);
+				assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 0, "release a cancelled task");
+
+				let stream = Stream::open(&mut h.session.clone(), version).await.unwrap();
+				assert!(h.publisher.handle_stream(id, bytes::Bytes::new(), stream).is_err());
+				assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 0, "release malformed input");
+
+				if id == ietf::Fetch::ID {
+					let stream = Stream::open(&mut h.session.clone(), version).await.unwrap();
+					h.publisher.handle_stream(id, body, stream).unwrap().await;
+					assert_eq!(h.publisher.owed.load(Ordering::Relaxed), 0, "release a completed fetch");
+				}
+			}
+		}
+	}
+
 	/// A distinctive request id, so `[FetchHeader::TYPE, REQUEST_ID]` is a usable needle.
 	const REQUEST_ID: u64 = 0x2B;
 

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

@@ -76,9 +76,18 @@ pub struct Config<S: crate::transport::poll::Session> {
 pub(crate) struct Driver {
 	pub(crate) withdrawal: crate::session::Withdrawal,
 	pub(crate) local_close: std::sync::Arc<std::sync::atomic::AtomicBool>,
+	// Dispatched SUBSCRIBE and FETCH serves still owing the peer data.
+	owed: std::sync::Arc<std::sync::atomic::AtomicUsize>,
 	future: MaybeSendBox<'static, Result<(), Error>>,
 }
 
+impl Driver {
+	/// Whether withdrawals and dispatched serves have reached the peer.
+	pub(crate) fn drained(&self) -> bool {
+		self.withdrawal.drained() && self.owed.load(std::sync::atomic::Ordering::Relaxed) == 0
+	}
+}
+
 impl std::future::Future for Driver {
 	type Output = Result<(), Error>;
 	fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
@@ -122,6 +131,8 @@ where
 	let withdrawing = withdrawal.clone();
 	let local_close = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
 	let closing = local_close.clone();
+	let owed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+	let serving = owed.clone();
 	let driver = async move {
 		// Our own Hop ID, taken from whichever origin the caller actually supplied so
 		// every session out of this process stamps the same one and cross-session loop
@@ -171,6 +182,7 @@ where
 				);
 				let (tasks, mut task_set) = TaskSet::new();
 				publisher.withdrawal = withdrawing.clone();
+				publisher.owed = serving.clone();
 
 				let subscriber = Subscriber::new(
 					runtime.clone(),
@@ -338,6 +350,7 @@ where
 				);
 				let (tasks, mut task_set) = TaskSet::new();
 				publisher.withdrawal = withdrawing.clone();
+				publisher.owed = serving.clone();
 
 				let subscriber = Subscriber::new(
 					runtime.clone(),
@@ -483,6 +496,7 @@ where
 		Driver {
 			withdrawal,
 			local_close,
+			owed,
 			future: driver,
 		},
 		goaway_handle,

rs/moq-net/src/session.rs

@@ -170,8 +170,7 @@ impl Session {
 	/// is still live never finishes, so finish or abort tracks before closing.
 	///
 	/// Both protocols withdraw this session's announcements and wait for their
-	/// delivery. IETF drafts 14 through 16 send withdrawals without waiting, and
-	/// IETF media streams are not drained yet.
+	/// delivery. IETF drafts 14 through 16 send withdrawals without waiting.
 	pub async fn close(self) -> Result<(), Error> {
 		if let Ok(mut close) = self.close.write()
 			&& close.is_none()

rs/moq-net/tests/session_close.rs

@@ -29,14 +29,14 @@ struct Setup {
 	_broadcast: moq_net::broadcast::Producer,
 }
 
-async fn setup() -> Setup {
+async fn setup(version: &str) -> Setup {
 	let publisher = produce_origin(1);
 	let broadcast = publisher.create_broadcast("bcast").unwrap();
 	let track = broadcast.create_track("video", None).unwrap();
 	broadcast.announce(Default::default()).unwrap();
 
 	let subscriber = produce_origin(2);
-	let mut options = MockConnectOptions::new("moq-lite-05".parse::<Version>().unwrap());
+	let mut options = MockConnectOptions::new(version.parse::<Version>().unwrap());
 	options.client_publish = Some(publisher.consume());
 	options.server_subscribe = Some(subscriber.clone());
 	let pair = connect_mock(options).await;
@@ -88,9 +88,18 @@ async fn setup() -> Setup {
 /// close is requested before the session wrote them.
 #[tokio::test(start_paused = true)]
 async fn close_delivers_a_finished_track() {
+	close_delivers_a_finished_track_for("moq-lite-05").await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn ietf_close_delivers_a_finished_track() {
+	close_delivers_a_finished_track_for("moq-transport-20").await;
+}
+
+async fn close_delivers_a_finished_track_for(version: &str) {
 	let Setup {
 		pair, track, reader, ..
-	} = setup().await;
+	} = setup(version).await;
 
 	let mut group = track.append_group().unwrap();
 	group.write_frame(Timestamp::ZERO, b"last".as_slice()).unwrap();
@@ -119,9 +128,18 @@ async fn close_delivers_a_finished_track() {
 /// session closes anyway.
 #[tokio::test(start_paused = true)]
 async fn close_gives_up_on_a_live_track() {
+	close_gives_up_on_a_live_track_for("moq-lite-05").await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn ietf_close_gives_up_on_a_live_track() {
+	close_gives_up_on_a_live_track_for("moq-transport-20").await;
+}
+
+async fn close_gives_up_on_a_live_track_for(version: &str) {
 	let Setup {
 		pair, track: _track, ..
-	} = setup().await;
+	} = setup(version).await;
 
 	let started = tokio::time::Instant::now();
 	let res = tokio::time::timeout(TIMEOUT, pair.client.close())
@@ -134,9 +152,18 @@ async fn close_gives_up_on_a_live_track() {
 /// An abort from another handle cuts a drain short.
 #[tokio::test(start_paused = true)]
 async fn abort_cuts_a_drain_short() {
+	abort_cuts_a_drain_short_for("moq-lite-05").await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn ietf_abort_cuts_a_drain_short() {
+	abort_cuts_a_drain_short_for("moq-transport-20").await;
+}
+
+async fn abort_cuts_a_drain_short_for(version: &str) {
 	let Setup {
 		pair, track: _track, ..
-	} = setup().await;
+	} = setup(version).await;
 
 	let other = pair.client.clone();
 	let started = tokio::time::Instant::now();