Skip to content

Commit ed80066

Browse files
richiemcilroyrichiemcilroy
committed
fix: bound processing transfers and download Drive revisions directly
1 parent 24e3c74 commit ed80066

21 files changed

Lines changed: 5657 additions & 51 deletions

‎apps/media-server/src/__tests__/lib/media-transfer.test.ts‎

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
import { afterEach, describe, expect, spyOn, test } from "bun:test";
2+
import { createHash, randomUUID } from "node:crypto";
3+
import { link, mkdtemp, readFile, rm, stat } from "node:fs/promises";
4+
import { tmpdir } from "node:os";
5+
import { join } from "node:path";
6+
import {
7+
cleanupMediaTransferCache,
8+
downloadDriveRevision,
9+
getMediaDownloadTarget,
10+
type MediaDownloadTarget,
11+
MediaTransferBudgetError,
12+
materializeMedia,
13+
releaseMaterializedMedia,
14+
withMediaTransfers,
15+
} from "../../lib/media-transfer";
16+
17+
const roots: string[] = [];
18+
const originalFetch = globalThis.fetch;
19+
const originalSecret = process.env.MEDIA_SERVER_WEBHOOK_SECRET;
20+
afterEach(async () => {
21+
globalThis.fetch = originalFetch;
22+
if (originalSecret === undefined)
23+
delete process.env.MEDIA_SERVER_WEBHOOK_SECRET;
24+
else process.env.MEDIA_SERVER_WEBHOOK_SECRET = originalSecret;
25+
await cleanupMediaTransferCache();
26+
for (const root of roots.splice(0))
27+
await rm(root, { recursive: true, force: true });
28+
});
29+
async function destination() {
30+
const root = await mkdtemp(join(tmpdir(), "media-transfer-test-"));
31+
roots.push(root);
32+
return join(root, "recording.mp4");
33+
}
34+
const content = Buffer.from("a real transfer with an interrupted stream");
35+
function target(): MediaDownloadTarget {
36+
return {
37+
version: 1,
38+
url: "https://www.googleapis.com/drive/v3/files/file/revisions/revision?alt=media",
39+
authorization: "Bearer private",
40+
objectIdentity: '"identity"',
41+
size: content.length,
42+
sha256: createHash("sha256").update(content).digest("hex"),
43+
};
44+
}
45+
function fetcher(
46+
run: (
47+
input: string | URL | Request,
48+
init?: RequestInit,
49+
) => Response | Promise<Response>,
50+
): typeof fetch {
51+
return Object.assign(run, { preconnect: fetch.preconnect }) as typeof fetch;
52+
}
53+
describe("bounded revision downloads", () => {
54+
test("resumes an interrupted pinned revision without rereading completed bytes", async () => {
55+
const path = await destination();
56+
let calls = 0;
57+
let received = 0;
58+
const network = fetcher((_input, init) => {
59+
calls++;
60+
if (calls === 1) {
61+
let pull = 0;
62+
return new Response(
63+
new ReadableStream({
64+
pull(controller) {
65+
if (pull++ === 0) controller.enqueue(content.subarray(0, 12));
66+
else controller.error(new Error("network disconnected"));
67+
},
68+
}),
69+
{ headers: { "Content-Length": String(content.length) } },
70+
);
71+
}
72+
expect(new Headers(init?.headers).get("range")).toBe("bytes=12-");
73+
return new Response(content.subarray(12), {
74+
status: 206,
75+
headers: {
76+
"Content-Length": String(content.length - 12),
77+
"Content-Range": `bytes 12-${content.length - 1}/${content.length}`,
78+
},
79+
});
80+
});
81+
await downloadDriveRevision(target(), path, {
82+
fetcher: network,
83+
onBytes: (bytes) => {
84+
received += bytes;
85+
},
86+
});
87+
expect(await readFile(path)).toEqual(content);
88+
expect(received).toBe(content.length);
89+
expect(calls).toBe(2);
90+
});
91+
test("rejects corruption and removes the incomplete file", async () => {
92+
const path = await destination();
93+
await expect(
94+
downloadDriveRevision(target(), path, {
95+
fetcher: fetcher(
96+
() =>
97+
new Response(Buffer.alloc(content.length), {
98+
headers: { "Content-Length": String(content.length) },
99+
}),
100+
),
101+
}),
102+
).rejects.toThrow("checksum");
103+
expect(await stat(path).catch(() => null)).toBeNull();
104+
});
105+
test("does not retry after the byte budget is exhausted", async () => {
106+
let calls = 0;
107+
await expect(
108+
downloadDriveRevision(target(), await destination(), {
109+
fetcher: fetcher(() => {
110+
calls++;
111+
return new Response(content, {
112+
headers: { "Content-Length": String(content.length) },
113+
});
114+
}),
115+
onBytes: () => {
116+
throw new MediaTransferBudgetError();
117+
},
118+
}),
119+
).rejects.toBeInstanceOf(MediaTransferBudgetError);
120+
expect(calls).toBe(1);
121+
});
122+
test("does not forward worker credentials to an arbitrary host", async () => {
123+
await expect(
124+
getMediaDownloadTarget("https://attacker.example/api/storage/object"),
125+
).rejects.toThrow("Untrusted");
126+
});
127+
test("bounds a repeatedly truncated source to three requests", async () => {
128+
let calls = 0;
129+
await expect(
130+
downloadDriveRevision(target(), await destination(), {
131+
fetcher: fetcher(() => {
132+
calls++;
133+
return new Response(
134+
new ReadableStream({
135+
start(controller) {
136+
controller.close();
137+
},
138+
}),
139+
{ headers: { "Content-Length": String(content.length) } },
140+
);
141+
}),
142+
}),
143+
).rejects.toThrow("checksum or size");
144+
expect(calls).toBe(3);
145+
});
146+
});
147+
148+
describe("parallel transfer admission and cache reuse", () => {
149+
function network() {
150+
process.env.MEDIA_SERVER_WEBHOOK_SECRET = "worker-secret";
151+
let downloads = 0;
152+
globalThis.fetch = fetcher((input) => {
153+
if (String(input).startsWith("https://cap.so/"))
154+
return Response.json(target());
155+
downloads++;
156+
return new Response(content, {
157+
headers: { "Content-Length": String(content.length) },
158+
});
159+
});
160+
return () => downloads;
161+
}
162+
const url = () => `https://cap.so/api/storage/object?key=${randomUUID()}`;
163+
164+
test("reuses a verified file across retries without another media download", async () => {
165+
const downloads = network();
166+
const input = url();
167+
await withMediaTransfers(content.length, async () => {
168+
const first = await materializeMedia(input);
169+
const second = await materializeMedia(input);
170+
expect(first?.path).toBe(second?.path);
171+
});
172+
await withMediaTransfers(1, async () => {
173+
expect(await materializeMedia(input)).toBeDefined();
174+
});
175+
expect(downloads()).toBe(1);
176+
});
177+
178+
test("releases one cache retention after concurrent reads in the same job", async () => {
179+
const downloads = network();
180+
const input = url();
181+
await withMediaTransfers(content.length, () => materializeMedia(input));
182+
await withMediaTransfers(content.length, async () => {
183+
const reads = await Promise.all(
184+
Array.from({ length: 10 }, () => materializeMedia(input)),
185+
);
186+
expect(new Set(reads.map((read) => read?.path)).size).toBe(1);
187+
});
188+
const now = Date.now();
189+
const clock = spyOn(Date, "now").mockReturnValue(now + 31 * 60_000);
190+
try {
191+
await withMediaTransfers(content.length, () => materializeMedia(input));
192+
} finally {
193+
clock.mockRestore();
194+
}
195+
expect(downloads()).toBe(2);
196+
});
197+
198+
test("retains a worker input after its cache pathname is removed", async () => {
199+
network();
200+
await withMediaTransfers(content.length, async () => {
201+
const local = await materializeMedia(url());
202+
expect(local).toBeDefined();
203+
if (!local) throw new Error("Input was not materialized");
204+
const path = await destination();
205+
await link(local.path, path);
206+
await releaseMaterializedMedia(local.path);
207+
await cleanupMediaTransferCache();
208+
expect(await readFile(path)).toEqual(content);
209+
});
210+
});
211+
212+
test("does not let simultaneous files overbook the remaining byte budget", async () => {
213+
const downloads = network();
214+
await withMediaTransfers(content.length, async () => {
215+
const results = await Promise.allSettled([
216+
materializeMedia(url()),
217+
materializeMedia(url()),
218+
]);
219+
expect(
220+
results.filter((result) => result.status === "fulfilled"),
221+
).toHaveLength(1);
222+
const failed = results.find((result) => result.status === "rejected");
223+
expect(
224+
failed?.status === "rejected" &&
225+
failed.reason instanceof MediaTransferBudgetError,
226+
).toBe(true);
227+
});
228+
expect(downloads()).toBe(1);
229+
});
230+
231+
test("does not follow a descriptor redirect with worker credentials", async () => {
232+
process.env.MEDIA_SERVER_WEBHOOK_SECRET = "worker-secret";
233+
let calls = 0;
234+
globalThis.fetch = fetcher((_input, init) => {
235+
calls++;
236+
expect(init?.redirect).toBe("manual");
237+
return Response.redirect("https://attacker.example/", 307);
238+
});
239+
await expect(getMediaDownloadTarget(url())).rejects.toThrow(
240+
"authorization failed",
241+
);
242+
expect(calls).toBe(1);
243+
});
244+
});

‎apps/media-server/src/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import app from "./app";
22
import { abortAllJobs } from "./lib/job-manager";
33
import { cancelAllMediaOperations } from "./lib/media-operations";
4+
import { cleanupMediaTransferCache } from "./lib/media-transfer";
45

56
const port = Number(process.env.PORT) || 3456;
67

@@ -17,6 +18,7 @@ const shutdown = async () => {
1718
console.log(`[media-server] Aborted ${abortedJobs} active jobs`);
1819
}
1920
await cancelAllMediaOperations();
21+
await cleanupMediaTransferCache();
2022
process.exit(0);
2123
};
2224

‎apps/media-server/src/lib/job-manager.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@ export type RecordingErrorCode =
2424
| "source-missing"
2525
| "source-changed"
2626
| "output-invalid"
27-
| "processing-unavailable";
27+
| "processing-unavailable"
28+
| "processing-budget-exhausted";
2829

2930
export interface RecordingVerificationRequest {
3031
version: 1;
@@ -324,10 +325,17 @@ export function hasCriticalMemoryPressure(): boolean {
324325
return getSystemResources().memoryPressure >= MEMORY_REJECT_THRESHOLD;
325326
}
326327

327-
export function canAcceptNewVideoProcess(): boolean {
328+
export function canAcceptNewVideoProcess(
329+
priority: "interactive" | "recovery" = "interactive",
330+
): boolean {
328331
const active = getActiveVideoProcessCount();
329332
const resources = getSystemResources();
330-
return active < resources.effectiveMax;
333+
return (
334+
active <
335+
(priority === "recovery" && resources.effectiveMax > 1
336+
? Math.max(1, Math.floor(resources.effectiveMax / 2))
337+
: resources.effectiveMax)
338+
);
331339
}
332340

333341
export function generateJobId(): string {

‎apps/media-server/src/lib/media-probe.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
getActiveProbeOperationCount,
1212
withMediaOperation,
1313
} from "./media-operations";
14+
import { materializeMedia } from "./media-transfer";
1415

1516
const PROBE_TIMEOUT_MS = 30_000;
1617
const probeFetch: typeof fetch = globalThis.fetch.bind(globalThis);
@@ -119,6 +120,8 @@ async function probeMedia(path: string): Promise<VideoMetadata> {
119120
}
120121

121122
export async function probeVideo(videoUrl: string): Promise<VideoMetadata> {
123+
const local = await materializeMedia(videoUrl);
124+
if (local) return probeVideoFile(local.path);
122125
if (!canAcceptNewProbeOperation()) {
123126
throw new Error("Server is busy, please try again later");
124127
}

0 commit comments

Comments
 (0)