perf(media-server): refill the initial download slot

apps/media-server/scripts/benchmark-parallel-download.ts

@@ -0,0 +1,104 @@
+import { createHash } from "node:crypto";
+import { mkdtemp, readFile, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+const concurrency = Number(process.argv[2] ?? 2);
+const mode = process.argv[3] ?? "large";
+if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 64)
+	throw new Error("Concurrency must be between 1 and 64");
+if (!["large", "small", "plain"].includes(mode))
+	throw new Error("Mode must be large, small or plain");
+
+import { downloadUrlToPathParallel } from "../src/lib/parallel-download";
+
+const responseDelayMs = 30;
+const partBytes = 8 * 1024 * 1024;
+const bytes = mode === "small" ? 1024 * 1024 : 128 * 1024 * 1024 + 123;
+const payload = new Uint8Array(bytes);
+for (let i = 0; i < bytes; i++) payload[i] = (i * 31 + (i >>> 20)) & 255;
+const fingerprint = createHash("sha256").update(payload).digest("hex");
+const directory = await mkdtemp(join(tmpdir(), "cap-download-slot-bench-"));
+let requests = 0;
+let active = 0;
+let peak = 0;
+const sockets = new Set<number>();
+const server = Bun.serve({
+	port: 0,
+	fetch(request, server) {
+		const port = server.requestIP(request)?.port;
+		if (port !== undefined) sockets.add(port);
+		requests++;
+		const match = request.headers.get("range")?.match(/^bytes=(\d+)-(\d+)$/);
+		const start = mode === "plain" ? 0 : Number(match?.[1] ?? 0);
+		const end =
+			mode === "plain"
+				? bytes - 1
+				: Math.min(Number(match?.[2] ?? bytes - 1), bytes - 1);
+		active++;
+		peak = Math.max(peak, active);
+		let settled = false;
+		const finish = () => {
+			if (!settled) {
+				settled = true;
+				active--;
+			}
+		};
+		return new Response(
+			new ReadableStream({
+				async start(controller) {
+					await Bun.sleep(responseDelayMs);
+					if (settled) return;
+					controller.enqueue(payload.subarray(start, end + 1));
+					controller.close();
+					finish();
+				},
+				cancel() {
+					finish();
+				},
+			}),
+			{
+				status: mode === "plain" ? 200 : 206,
+				headers: {
+					ETag: '"fixture"',
+					...(mode === "plain"
+						? {}
+						: { "Content-Range": `bytes ${start}-${end}/${bytes}` }),
+				},
+			},
+		);
+	},
+});
+try {
+	const path = join(directory, "download.bin");
+	const begin = performance.now();
+	const result = await downloadUrlToPathParallel(
+		`http://127.0.0.1:${server.port}/recording`,
+		path,
+		{ concurrency, partBytes },
+	);
+	const elapsedMs = performance.now() - begin;
+	const actual = createHash("sha256")
+		.update(await readFile(path))
+		.digest("hex");
+	if (actual !== fingerprint || result.bytes !== bytes)
+		throw new Error("byte mismatch");
+	console.log(
+		JSON.stringify({
+			source: import.meta.resolve("../src/lib/parallel-download.ts"),
+			responseDelayMs,
+			concurrency,
+			mode,
+			bytes,
+			elapsedMs,
+			result,
+			requests,
+			peakDelayedResponses: peak,
+			tcpConnections: sockets.size,
+			fingerprint,
+		}),
+	);
+} finally {
+	await server.stop(true);
+	await rm(directory, { recursive: true, force: true });
+}

apps/media-server/src/__tests__/lib/parallel-download-concurrency.test.ts

@@ -0,0 +1,121 @@
+import { afterAll, beforeAll, expect, test } from "bun:test";
+import { mkdtemp, readFile, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { setImmediate as nextTurn } from "node:timers/promises";
+import { downloadUrlToPathParallel } from "../../lib/parallel-download";
+
+const MB = 1024 * 1024;
+const payload = new Uint8Array(4 * MB + 123).fill(47);
+let directory: string;
+beforeAll(async () => {
+	directory = await mkdtemp(join(tmpdir(), "cap-download-slots-"));
+});
+afterAll(async () => {
+	await rm(directory, { recursive: true, force: true });
+});
+
+test.each([1, 2, 4])(
+	"keeps all body streams within %i configured slots",
+	async (concurrency) => {
+		let active = 0;
+		let maximum = 0;
+		const fetchImpl = Object.assign(
+			async (_input: RequestInfo | URL, init?: RequestInit) => {
+				const range = new Headers(init?.headers)
+					.get("range")
+					?.match(/^bytes=(\d+)-(\d+)$/);
+				if (!range) throw new Error("Expected a range");
+				const start = Number(range[1]);
+				const end = Math.min(Number(range[2]), payload.length - 1);
+				return new Response(
+					new ReadableStream<Uint8Array>({
+						async start(controller) {
+							active++;
+							maximum = Math.max(maximum, active);
+							await nextTurn();
+							controller.enqueue(payload.subarray(start, end + 1));
+							controller.close();
+							active--;
+						},
+					}),
+					{
+						status: 206,
+						headers: {
+							"Content-Range": `bytes ${start}-${end}/${payload.length}`,
+						},
+					},
+				);
+			},
+			{ preconnect: fetch.preconnect },
+		);
+		const path = join(directory, `limit-${concurrency}.bin`);
+		const result = await downloadUrlToPathParallel(
+			"https://media.test/video",
+			path,
+			{ concurrency, partBytes: MB, fetchImpl },
+		);
+		expect(maximum).toBeLessThanOrEqual(concurrency);
+		expect(active).toBe(0);
+		expect(result).toEqual({ bytes: payload.length, parts: 5, ranged: true });
+		expect(new Uint8Array(await readFile(path))).toEqual(payload);
+	},
+);
+
+test("reuses the first body's slot while another body stays blocked", async () => {
+	const secondStarted = Promise.withResolvers<void>();
+	const thirdStarted = Promise.withResolvers<void>();
+	const pending: Array<() => void> = [];
+	let requests = 0;
+	let autoRelease = false;
+	const fetchImpl = Object.assign(
+		async (_input: RequestInfo | URL, init?: RequestInit) => {
+			const range = new Headers(init?.headers)
+				.get("range")
+				?.match(/^bytes=(\d+)-(\d+)$/);
+			if (!range) throw new Error("Expected a range");
+			const start = Number(range[1]);
+			const end = Math.min(Number(range[2]), payload.length - 1);
+			const body = new ReadableStream<Uint8Array>({
+				start(controller) {
+					const release = () => {
+						controller.enqueue(payload.subarray(start, end + 1));
+						controller.close();
+					};
+					if (autoRelease) release();
+					else pending.push(release);
+				},
+			});
+			requests++;
+			if (requests === 2) secondStarted.resolve();
+			if (requests === 3) thirdStarted.resolve();
+			return new Response(body, {
+				status: 206,
+				headers: { "Content-Range": `bytes ${start}-${end}/${payload.length}` },
+			});
+		},
+		{ preconnect: fetch.preconnect },
+	);
+	const path = join(directory, "refill.bin");
+	const download = downloadUrlToPathParallel("https://media.test/video", path, {
+		concurrency: 2,
+		partBytes: MB,
+		fetchImpl,
+	});
+	const timeout = setTimeout(
+		() => thirdStarted.reject(new Error("First slot was never refilled")),
+		2000,
+	);
+	try {
+		await secondStarted.promise;
+		pending.shift()?.();
+		await thirdStarted.promise;
+		expect(requests).toBe(3);
+	} finally {
+		clearTimeout(timeout);
+		autoRelease = true;
+		for (const release of pending) release();
+		await download;
+	}
+	expect(new Uint8Array(await readFile(path))).toEqual(payload);
+});

apps/media-server/src/lib/parallel-download.ts

@@ -177,10 +177,10 @@ export async function downloadUrlToPathParallel(
 			}
 		};
 		const workers = Array.from(
-			{ length: Math.max(1, Math.min(options.concurrency - 1, ranges.count)) },
+			{ length: Math.max(0, Math.min(options.concurrency - 1, ranges.count)) },
 			() => worker(),
 		);
-		const firstRead = readInto(first, write, 0, firstBytes);
+		const firstRead = readInto(first, write, 0, firstBytes).then(worker);
 		try {
 			await Promise.all([firstRead, ...workers]);
 		} catch (error) {