Add concurrency and timeout controls to audio extraction

Author: richiemcilroy

apps/media-server/src/__tests__/index.test.ts

@@ -10,7 +10,7 @@ describe("GET /", () => {
 		expect(data).toEqual({
 			name: "@cap/media-server",
 			version: "1.0.0",
-			endpoints: ["/health", "/audio/check", "/audio/extract"],
+			endpoints: ["/health", "/audio/status", "/audio/check", "/audio/extract"],
 		});
 	});
 });

apps/media-server/src/app.ts

@@ -14,7 +14,7 @@ app.get("/", (c) => {
 	return c.json({
 		name: "@cap/media-server",
 		version: "1.0.0",
-		endpoints: ["/health", "/audio/check", "/audio/extract"],
+		endpoints: ["/health", "/audio/status", "/audio/check", "/audio/extract"],
 	});
 });
 

apps/media-server/src/lib/ffmpeg.ts

@@ -1,4 +1,4 @@
-import { spawn } from "bun";
+import { type Subprocess, spawn } from "bun";
 
 export interface AudioExtractionOptions {
 	format?: "mp3";
@@ -12,23 +12,91 @@ const DEFAULT_OPTIONS: Required<AudioExtractionOptions> = {
 	bitrate: "128k",
 };
 
+const CHECK_TIMEOUT_MS = 30_000;
+const EXTRACT_TIMEOUT_MS = 120_000;
+const MAX_AUDIO_SIZE_BYTES = 100 * 1024 * 1024;
+
+let activeProcesses = 0;
+const MAX_CONCURRENT_PROCESSES = 6;
+
+export function getActiveProcessCount(): number {
+	return activeProcesses;
+}
+
+export function canAcceptNewProcess(): boolean {
+	return activeProcesses < MAX_CONCURRENT_PROCESSES;
+}
+
+function killProcess(proc: Subprocess): void {
+	try {
+		proc.kill();
+	} catch {}
+}
+
+async function withTimeout<T>(
+	promise: Promise<T>,
+	timeoutMs: number,
+	cleanup?: () => void,
+): Promise<T> {
+	let timeoutId: ReturnType<typeof setTimeout> | undefined;
+	const timeoutPromise = new Promise<never>((_, reject) => {
+		timeoutId = setTimeout(() => {
+			cleanup?.();
+			reject(new Error(`Operation timed out after ${timeoutMs}ms`));
+		}, timeoutMs);
+	});
+
+	try {
+		const result = await Promise.race([promise, timeoutPromise]);
+		if (timeoutId) clearTimeout(timeoutId);
+		return result;
+	} catch (err) {
+		if (timeoutId) clearTimeout(timeoutId);
+		throw err;
+	}
+}
+
 export async function checkHasAudioTrack(videoUrl: string): Promise<boolean> {
+	if (!canAcceptNewProcess()) {
+		throw new Error("Server is busy, please try again later");
+	}
+
+	activeProcesses++;
+
 	const proc = spawn({
 		cmd: ["ffmpeg", "-i", videoUrl, "-hide_banner"],
 		stdout: "pipe",
 		stderr: "pipe",
 	});
 
-	const stderrText = await new Response(proc.stderr).text();
-	await proc.exited;
+	try {
+		const result = await withTimeout(
+			(async () => {
+				const stderrText = await new Response(proc.stderr).text();
+				await proc.exited;
+				return /Stream #\d+:\d+.*Audio:/.test(stderrText);
+			})(),
+			CHECK_TIMEOUT_MS,
+			() => killProcess(proc),
+		);
 
-	return /Stream #\d+:\d+.*Audio:/.test(stderrText);
+		return result;
+	} finally {
+		activeProcesses--;
+		killProcess(proc);
+	}
 }
 
 export async function extractAudio(
 	videoUrl: string,
 	options: AudioExtractionOptions = {},
 ): Promise<Uint8Array> {
+	if (!canAcceptNewProcess()) {
+		throw new Error("Server is busy, please try again later");
+	}
+
+	activeProcesses++;
+
 	const opts = { ...DEFAULT_OPTIONS, ...options };
 
 	const ffmpegArgs = [
@@ -51,23 +119,53 @@ export async function extractAudio(
 		stderr: "pipe",
 	});
 
-	const [stdout, stderrText, exitCode] = await Promise.all([
-		new Response(proc.stdout).arrayBuffer(),
-		new Response(proc.stderr).text(),
-		proc.exited,
-	]);
+	try {
+		const result = await withTimeout(
+			(async () => {
+				const [stdout, stderrText, exitCode] = await Promise.all([
+					new Response(proc.stdout).arrayBuffer(),
+					new Response(proc.stderr).text(),
+					proc.exited,
+				]);
+
+				if (exitCode !== 0) {
+					throw new Error(`FFmpeg exited with code ${exitCode}: ${stderrText}`);
+				}
+
+				if (stdout.byteLength > MAX_AUDIO_SIZE_BYTES) {
+					throw new Error(
+						`Audio too large: ${stdout.byteLength} bytes exceeds ${MAX_AUDIO_SIZE_BYTES} byte limit`,
+					);
+				}
+
+				return new Uint8Array(stdout);
+			})(),
+			EXTRACT_TIMEOUT_MS,
+			() => killProcess(proc),
+		);
 
-	if (exitCode !== 0) {
-		throw new Error(`FFmpeg exited with code ${exitCode}: ${stderrText}`);
+		return result;
+	} finally {
+		activeProcesses--;
+		killProcess(proc);
 	}
+}
 
-	return new Uint8Array(stdout);
+export interface StreamingExtractResult {
+	stream: ReadableStream<Uint8Array>;
+	cleanup: () => void;
 }
 
-export async function extractAudioStream(
+export function extractAudioStream(
 	videoUrl: string,
 	options: AudioExtractionOptions = {},
-): Promise<ReadableStream<Uint8Array>> {
+): StreamingExtractResult {
+	if (!canAcceptNewProcess()) {
+		throw new Error("Server is busy, please try again later");
+	}
+
+	activeProcesses++;
+
 	const opts = { ...DEFAULT_OPTIONS, ...options };
 
 	const ffmpegArgs = [
@@ -90,5 +188,52 @@ export async function extractAudioStream(
 		stderr: "pipe",
 	});
 
-	return proc.stdout as ReadableStream<Uint8Array>;
+	let timeoutId: ReturnType<typeof setTimeout> | undefined;
+	let cleaned = false;
+
+	const cleanup = () => {
+		if (cleaned) return;
+		cleaned = true;
+		if (timeoutId) clearTimeout(timeoutId);
+		activeProcesses--;
+		killProcess(proc);
+	};
+
+	timeoutId = setTimeout(() => {
+		console.error("[ffmpeg] Stream extraction timed out");
+		cleanup();
+	}, EXTRACT_TIMEOUT_MS);
+
+	proc.exited.then((code) => {
+		if (code !== 0) {
+			console.error(`[ffmpeg] Stream extraction exited with code ${code}`);
+		}
+		cleanup();
+	});
+
+	const originalStream = proc.stdout as ReadableStream<Uint8Array>;
+
+	const wrappedStream = new ReadableStream<Uint8Array>({
+		async start(controller) {
+			const reader = originalStream.getReader();
+			try {
+				while (true) {
+					const { done, value } = await reader.read();
+					if (done) break;
+					controller.enqueue(value);
+				}
+				controller.close();
+			} catch (err) {
+				controller.error(err);
+			} finally {
+				reader.releaseLock();
+				cleanup();
+			}
+		},
+		cancel() {
+			cleanup();
+		},
+	});
+
+	return { stream: wrappedStream, cleanup };
 }

apps/media-server/src/routes/audio.ts

@@ -1,13 +1,39 @@
 import { Hono } from "hono";
 import { z } from "zod";
-import { checkHasAudioTrack, extractAudio } from "../lib/ffmpeg";
+import {
+	canAcceptNewProcess,
+	checkHasAudioTrack,
+	extractAudio,
+	extractAudioStream,
+	getActiveProcessCount,
+} from "../lib/ffmpeg";
 
 const audio = new Hono();
 
 const videoUrlSchema = z.object({
 	videoUrl: z.string().url(),
 });
 
+const extractSchema = z.object({
+	videoUrl: z.string().url(),
+	stream: z.boolean().optional().default(false),
+});
+
+function isBusyError(err: unknown): boolean {
+	return err instanceof Error && err.message.includes("Server is busy");
+}
+
+function isTimeoutError(err: unknown): boolean {
+	return err instanceof Error && err.message.includes("timed out");
+}
+
+audio.get("/status", (c) => {
+	return c.json({
+		activeProcesses: getActiveProcessCount(),
+		canAcceptNewProcess: canAcceptNewProcess(),
+	});
+});
+
 audio.post("/check", async (c) => {
 	const body = await c.req.json();
 	const result = videoUrlSchema.safeParse(body);
@@ -28,6 +54,29 @@ audio.post("/check", async (c) => {
 		return c.json({ hasAudio });
 	} catch (err) {
 		console.error("[audio/check] Error:", err);
+
+		if (isBusyError(err)) {
+			return c.json(
+				{
+					error: "Server is busy",
+					code: "SERVER_BUSY",
+					details: "Too many concurrent requests, please retry later",
+				},
+				503,
+			);
+		}
+
+		if (isTimeoutError(err)) {
+			return c.json(
+				{
+					error: "Request timed out",
+					code: "TIMEOUT",
+					details: err instanceof Error ? err.message : String(err),
+				},
+				504,
+			);
+		}
+
 		return c.json(
 			{
 				error: "Failed to check audio track",
@@ -41,7 +90,7 @@ audio.post("/check", async (c) => {
 
 audio.post("/extract", async (c) => {
 	const body = await c.req.json();
-	const result = videoUrlSchema.safeParse(body);
+	const result = extractSchema.safeParse(body);
 
 	if (!result.success) {
 		return c.json(
@@ -54,16 +103,28 @@ audio.post("/extract", async (c) => {
 		);
 	}
 
+	const { videoUrl, stream: useStreaming } = result.data;
+
 	try {
-		const hasAudio = await checkHasAudioTrack(result.data.videoUrl);
+		const hasAudio = await checkHasAudioTrack(videoUrl);
 		if (!hasAudio) {
 			return c.json(
 				{ error: "Video has no audio track", code: "NO_AUDIO_TRACK" },
 				422,
 			);
 		}
 
-		const audioData = await extractAudio(result.data.videoUrl);
+		if (useStreaming) {
+			const { stream } = extractAudioStream(videoUrl);
+			return new Response(stream, {
+				headers: {
+					"Content-Type": "audio/mpeg",
+					"Transfer-Encoding": "chunked",
+				},
+			});
+		}
+
+		const audioData = await extractAudio(videoUrl);
 
 		return new Response(Buffer.from(audioData), {
 			headers: {
@@ -73,6 +134,29 @@ audio.post("/extract", async (c) => {
 		});
 	} catch (err) {
 		console.error("[audio/extract] Error:", err);
+
+		if (isBusyError(err)) {
+			return c.json(
+				{
+					error: "Server is busy",
+					code: "SERVER_BUSY",
+					details: "Too many concurrent requests, please retry later",
+				},
+				503,
+			);
+		}
+
+		if (isTimeoutError(err)) {
+			return c.json(
+				{
+					error: "Request timed out",
+					code: "TIMEOUT",
+					details: err instanceof Error ? err.message : String(err),
+				},
+				504,
+			);
+		}
+
 		return c.json(
 			{
 				error: "Failed to extract audio",