fix(net): finish legacy tracks whose FIN races a graceful close

js/net/src/lite/connection.test.ts

@@ -1,5 +1,6 @@
 import { expect, spyOn, test } from "bun:test";
 import { Manual, Milli, Nano } from "@moq/clock";
+import Qmux from "@moq/qmux";
 import { accept } from "../connection/accept.ts";
 import { connect } from "../connection/connect.ts";
 import { SessionCode, SessionError } from "../error.ts";
@@ -23,7 +24,17 @@ import {
 	SubscribeStart,
 } from "./subscribe.ts";
 import { TrackInfo, Track as TrackMessage } from "./track.ts";
-import { ALPN_04, ALPN_05, ALPN_06, ALPN_07_WIP, Version } from "./version.ts";
+import {
+	ALPN,
+	ALPN_03,
+	ALPN_04,
+	ALPN_05,
+	ALPN_06,
+	ALPN_07_WIP,
+	hasTrackStream,
+	Version,
+	versionName,
+} from "./version.ts";
 
 /** A transport whose `getStats` behaves as described, or is absent entirely. */
 function transport(getStats?: () => Promise<unknown>): WebTransport {
@@ -203,6 +214,70 @@ function holdFins(transport: MockTransport, streams: "uni" | "incoming-bidi") {
 	};
 }
 
+// Hand `transport` each incoming unidirectional stream only once its FIN arrived, buffered outside
+// the transport so a session close cannot discard it. QMux keeps a finished stream's unread bytes
+// across a close this way, so the reads still owed land after the close itself.
+function deliverWhole(transport: MockTransport) {
+	Object.defineProperty(transport, "incomingUnidirectionalStreams", {
+		value: transport.incomingUnidirectionalStreams.pipeThrough(
+			new TransformStream<ReadableStream<Uint8Array>, ReadableStream<Uint8Array>>({
+				transform(stream, controller) {
+					controller.enqueue(
+						new ReadableStream<Uint8Array>({
+							async start(whole) {
+								const chunks = await Array.fromAsync(stream);
+								for (const chunk of chunks) whole.enqueue(chunk);
+								whole.close();
+							},
+						}),
+					);
+				},
+			}),
+		),
+	});
+}
+
+// Deliver the bytes of every bidirectional stream `transport` opens at once, but each FIN only a task
+// after the session closed, so every reaction to the close runs before a SUBSCRIBE FIN sent first.
+function finAfterClose(transport: MockTransport) {
+	const create = transport.createBidirectionalStream.bind(transport);
+	transport.createBidirectionalStream = async (options) => {
+		const stream = await create(options);
+		const readable = stream.readable.pipeThrough(
+			new TransformStream<Uint8Array, Uint8Array>({
+				async flush() {
+					await transport.closed;
+					await new Promise((resolve) => setTimeout(resolve, 0));
+				},
+			}),
+		);
+		return { readable, writable: stream.writable };
+	};
+}
+
+// A qmux session pair over an in-memory socket. Unlike the mock, its receive streams are
+// pull-driven, so a stream's FIN surfaces only to a read.
+function qmuxPair(): { client: WebTransport; server: WebTransport } {
+	const socket = (readable: ReadableStream<Uint8Array>, writable: WritableStream<Uint8Array>) => {
+		const closed = Promise.withResolvers<{ closeCode: number; reason: string }>();
+		return {
+			url: "",
+			opened: Promise.resolve({ readable, writable, extensions: "", protocol: "qmux-01" }),
+			closed: closed.promise,
+			close: () => closed.resolve({ closeCode: 1000, reason: "" }),
+		};
+	};
+	// Copied, since a sender may reuse its buffer once a write resolves.
+	const pipe = () =>
+		new TransformStream<Uint8Array, Uint8Array>({ transform: (chunk, c) => c.enqueue(chunk.slice()) });
+	const up = pipe();
+	const down = pipe();
+	const config = { idleTimeout: null };
+	const client = new Qmux(socket(down.readable, up.writable), { config });
+	const server = Qmux.accept(socket(up.readable, down.writable), { config });
+	return { client, server };
+}
+
 // Holds the one-second close deadline on a manual clock, so a test fires it by hand instead of waiting.
 function mockDeadline() {
 	const clock = new Manual();
@@ -217,9 +292,17 @@ async function settle() {
 	for (let i = 0; i < 5; i++) await new Promise((resolve) => setTimeout(resolve, 0));
 }
 
-for (const alpn of [ALPN_05, ALPN_06, ALPN_07_WIP]) {
+for (const [alpn, version] of [
+	[ALPN, Version.DRAFT_01],
+	[ALPN, Version.DRAFT_02],
+	[ALPN_03, Version.DRAFT_03],
+	[ALPN_04, Version.DRAFT_04],
+	[ALPN_05, Version.DRAFT_05],
+	[ALPN_06, Version.DRAFT_06],
+	[ALPN_07_WIP, Version.DRAFT_07],
+] as const) {
 	// A publisher that finishes a track and closes must not cut the final group short.
-	test(`close delivers a finished track's final group over ${alpn}`, async () => {
+	test(`close delivers a finished track's final group over ${versionName(version)}`, async () => {
 		const pair = createMockTransportPair(alpn);
 		const fin = holdFins(pair.server, "uni");
 		const origin = new Producer();
@@ -228,18 +311,18 @@ for (const alpn of [ALPN_05, ALPN_06, ALPN_07_WIP]) {
 		broadcastPublication.announce();
 		const producer = broadcast.createTrack("video", { timescale: Timescale.MILLI });
 		const url = new URL("https://localhost/test");
-		const [client, server] = await Promise.all([
-			connect({ url, transport: pair.client }),
-			accept({ url, transport: pair.server, publish: origin.consume() }),
-		]);
+		const client = new Connection({ url, quic: pair.client, version });
+		const server = new Connection({ url, quic: pair.server, version, publish: origin.consume(), client: false });
 		const remote = wireOf(client).consume(Path.from("room"));
 		const track = remote.track("video").subscribe().ordered();
 		let closed = false;
 		void pair.server.closed.then(() => {
 			closed = true;
 		});
 		try {
-			producer.appendGroup().writeFrame(textFrame("first"));
+			const first = producer.appendGroup();
+			first.writeFrame(textFrame("first"));
+			first.finish();
 			expect(await (await track.nextGroup())?.readString()).toBe("first");
 
 			fin.enable();
@@ -258,6 +341,7 @@ for (const alpn of [ALPN_05, ALPN_06, ALPN_07_WIP]) {
 			fin.release();
 			await closing;
 			expect(closed).toBe(true);
+			expect(await track.nextGroup()).toBeUndefined();
 		} finally {
 			fin.release();
 			track.close();
@@ -268,6 +352,100 @@ for (const alpn of [ALPN_05, ALPN_06, ALPN_07_WIP]) {
 			origin.close();
 		}
 	});
+
+	// The peer's close can overtake the final group or, where it is the track's only end, the
+	// SUBSCRIBE FIN it sent first; the closing session reads them.
+	for (const overtaken of hasTrackStream(version) ? ["group"] : ["group", "fin"]) {
+		test(`a peer close drains the overtaken ${overtaken} over ${versionName(version)}`, async () => {
+			const pair = createMockTransportPair(alpn);
+			deliverWhole(pair.client);
+			if (overtaken === "fin") finAfterClose(pair.client);
+			const origin = new Producer();
+			const broadcastPublication = origin.createBroadcast(Path.from("room"));
+			const broadcast = broadcastPublication.broadcast;
+			broadcastPublication.announce();
+			const producer = broadcast.createTrack("video", { timescale: Timescale.MILLI });
+			const url = new URL("https://localhost/test");
+			const client = new Connection({ url, quic: pair.client, version });
+			const server = new Connection({
+				url,
+				quic: pair.server,
+				version,
+				publish: origin.consume(),
+				client: false,
+			});
+			const remote = wireOf(client).consume(Path.from("room"));
+			const track = remote.track("video").subscribe().ordered();
+			try {
+				const first = producer.appendGroup();
+				first.writeFrame(textFrame("first"));
+				first.finish();
+				expect(await (await track.nextGroup())?.readString()).toBe("first");
+
+				const frames = Array.from({ length: 32 }, (_, i) => `frame ${i}`);
+				const last = producer.appendGroup();
+				for (const frame of frames) last.writeFrame(textFrame(frame));
+				last.finish();
+				producer.finish();
+				producer.close();
+				await server.close();
+
+				const group = await track.nextGroup();
+				for (const frame of frames) expect(await group?.readString()).toBe(frame);
+				expect(await group?.readString()).toBeUndefined();
+				expect(await track.nextGroup()).toBeUndefined();
+			} finally {
+				track.close();
+				remote.close();
+				client.abort();
+				server.abort();
+				broadcast.close();
+				origin.close();
+			}
+		});
+	}
+
+	// The subscriber reads a legacy subscribe stream to its FIN, since QMux closes a stream only
+	// on a read; awaiting the FIN instead would hold the closing session's drain forever.
+	if (!hasTrackStream(version)) {
+		test(`a peer close finishes a FIN'd track over qmux ${versionName(version)}`, async () => {
+			const pair = qmuxPair();
+			const origin = new Producer();
+			const broadcastPublication = origin.createBroadcast(Path.from("room"));
+			const broadcast = broadcastPublication.broadcast;
+			broadcastPublication.announce();
+			const producer = broadcast.createTrack("video", { timescale: Timescale.MILLI });
+			const url = new URL("https://localhost/test");
+			const client = new Connection({ url, quic: pair.client, version });
+			const server = new Connection({
+				url,
+				quic: pair.server,
+				version,
+				publish: origin.consume(),
+				client: false,
+			});
+			const remote = wireOf(client).consume(Path.from("room"));
+			const track = remote.track("video").subscribe().ordered();
+			try {
+				const group = producer.appendGroup();
+				group.writeFrame(textFrame("only"));
+				group.finish();
+				expect(await (await track.nextGroup())?.readString()).toBe("only");
+
+				producer.finish();
+				producer.close();
+				await server.close();
+				expect(await track.nextGroup()).toBeUndefined();
+			} finally {
+				track.close();
+				remote.close();
+				client.abort();
+				server.abort();
+				broadcast.close();
+				origin.close();
+			}
+		});
+	}
 }
 
 /** A lite-07 publisher serving a finished track to a raw subscriber that has read up to the FIN. */

js/net/src/lite/connection.ts

@@ -104,6 +104,9 @@ export class Connection implements Established {
 	// Module for distributing tracks.
 	#subscriber: Subscriber;
 
+	// The received unidirectional streams still being read, which a closing session drains.
+	#unis = new Set<Promise<void>>();
+
 	/** The peer's PROBE estimates; see {@link Established.probe}. */
 	readonly probe: Getter<Probe>;
 
@@ -265,6 +268,10 @@ export class Connection implements Established {
 		this.#publisher.close();
 		this.#subscriber.close();
 
+		this.#closeTransport();
+	}
+
+	#closeTransport(): void {
 		try {
 			// TODO: For whatever reason, this try/catch doesn't seem to work..?
 			this.#quic.close();
@@ -316,9 +323,16 @@ export class Connection implements Established {
 			probes.close();
 			// A session that ends before the peer's SETUP never learns its implementation.
 			this.#peerImplementation.resolve(undefined);
+			const err = fatal ?? (await closeError(this.#quic));
+			// A failed task can end the session over a live transport. Closing it first ends
+			// every read the drain waits on.
+			this.#closeTransport();
+			// The peer's close can overtake the group streams and SUBSCRIBE FIN it sent first.
+			// Read what already arrived, so a track its publisher ended finishes rather than aborts.
+			await Promise.all([...this.#unis, this.#subscriber.drain()]);
 			// The session died under every track it was receiving, so they end with its
 			// error. A deliberate close() already ended them cleanly, which makes this a no-op.
-			this.#fail(fatal ?? (await closeError(this.#quic)));
+			this.#fail(err);
 		}
 	}
 
@@ -448,7 +462,7 @@ export class Connection implements Established {
 			const stream = await readers.next();
 			if (!stream) break;
 
-			this.#runUni(stream)
+			const task = this.#runUni(stream)
 				.then(() => {
 					stream.stop(StreamCode.Cancel);
 				})
@@ -458,7 +472,9 @@ export class Connection implements Established {
 						closeViolation(this.#quic, err);
 						this.#fail(err);
 					}
-				});
+				})
+				.finally(() => this.#unis.delete(task));
+			this.#unis.add(task);
 		}
 	}
 

js/net/src/lite/subscriber.test.ts

@@ -30,12 +30,13 @@ import {
 	encodeSubscribeResponse,
 	SubscribeDrop,
 	SubscribeEnd,
+	SubscribeOk,
 	type SubscribeResponse,
 	SubscribeStart,
 } from "./subscribe.ts";
 import { SUBSCRIBE_SETUP_TIMEOUT_MS, Subscriber } from "./subscriber.ts";
 import { Track, TrackInfo } from "./track.ts";
-import { Version } from "./version.ts";
+import { Version, versionName } from "./version.ts";
 
 test("closing the subscriber suppresses probe stream warnings", async () => {
 	let readable!: ReadableStreamDefaultController<Uint8Array>;
@@ -2255,6 +2256,42 @@ test("a lite-07 publisher FIN ends a covered bounded subscription", async () =>
 	}
 });
 
+// Before lite-05 the publisher FINs the subscribe stream only once the track is over, so that FIN
+// finishes the track once the tail settles, rather than releasing it for another route.
+test.each(
+	[Version.DRAFT_01, Version.DRAFT_02, Version.DRAFT_03, Version.DRAFT_04].map((v) => [versionName(v), v] as const),
+)("a %s publisher FIN finishes the track", async (_, version) => {
+	const clock = new Manual();
+	const installed = clock.install();
+	const { quic, streams, opened, closes } = fakeSession();
+	const subscriber = new Subscriber(quic, version, HopSchema.parse(1n));
+	const broadcast = subscriber.consume(Path.from("room"));
+	const track = broadcast.track("video").subscribe();
+	try {
+		const stream = await opened(0);
+		await stream.reading;
+		const chunks: Uint8Array[] = [];
+		const writer = new Writer(
+			new WritableStream<Uint8Array>({ write: (chunk) => void chunks.push(new Uint8Array(chunk)) }),
+			version,
+		);
+		await encodeSubscribeResponse(writer, { ok: new SubscribeOk({}) }, version);
+		for (const chunk of chunks) streams[0].inbound.enqueue(chunk);
+		streams[0].inbound.close();
+		await drainReactions();
+
+		await clock.advance(Nano.fromMilli(GRACE));
+		expect(await track.finished()).toEqual({ group: 0, frame: 0 });
+		expect(await track.closed).toBeNull();
+		expect(closes).toEqual([]);
+	} finally {
+		track.close();
+		broadcast.close();
+		subscriber.close();
+		installed.close();
+	}
+});
+
 // The FIN may answer the bounded range before the publisher read an update that widened it,
 // so the track fails rather than ending clean short of the demand.
 test.each([