Skip to content

Commit bc53298

Browse files
committed
fix: resume large Drive outputs in bounded upload chunks
1 parent f139bfe commit bc53298

5 files changed

Lines changed: 353 additions & 8 deletions

File tree

.github/workflows/docker-build-media-server.yml

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,10 @@ jobs:
8484
src/__tests__/lib/media-size.test.ts \
8585
src/__tests__/lib/media-probe.integration.test.ts
8686
docker run --rm --network none --entrypoint bun "$MEDIA_IMAGE" test \
87-
src/__tests__/lib/media-transfer.test.ts
87+
src/__tests__/lib/media-transfer.test.ts \
88+
src/__tests__/lib/drive-resumable-upload.test.ts \
89+
src/__tests__/lib/storage-upload.test.ts \
90+
src/__tests__/lib/container-memory.test.ts
8891
docker run --rm --network none --entrypoint bun "$MEDIA_IMAGE" test \
8992
src/__tests__/routes/recording-verification.test.ts
9093
docker run --rm --network none --entrypoint bun "$MEDIA_IMAGE" test \
@@ -106,7 +109,8 @@ jobs:
106109
docker run --rm --network none --cpus 2 --memory 2g \
107110
-e MEDIA_SERVER_TRANSFER_PERFORMANCE_TESTS=1 \
108111
--entrypoint bun "$MEDIA_IMAGE" test \
109-
src/__tests__/lib/media-transfer.integration.test.ts
112+
src/__tests__/lib/media-transfer.integration.test.ts \
113+
src/__tests__/lib/drive-resumable-upload.test.ts
110114
111115
- name: Export Digest
112116
if: github.event_name != 'pull_request'

apps/media-server/Dockerfile

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@ RUN bun install --frozen-lockfile --production
1111

1212
COPY apps/media-server/src ./src
1313

14-
RUN bun test src/__tests__/lib/media-size.test.ts src/__tests__/lib/media-probe.integration.test.ts \
14+
RUN bun test src/__tests__/lib/drive-resumable-upload.test.ts src/__tests__/lib/storage-upload.test.ts src/__tests__/lib/container-memory.test.ts \
15+
&& bun test src/__tests__/lib/media-size.test.ts src/__tests__/lib/media-probe.integration.test.ts \
1516
&& bun test src/__tests__/lib/recording-verification.integration.test.ts src/__tests__/lib/job-manager.test.ts \
1617
&& bun test src/__tests__/lib/media-transfer.test.ts \
1718
&& bun test src/__tests__/routes/recording-verification.test.ts \
@@ -24,7 +25,8 @@ RUN MEDIA_SERVER_RECORDING_PERFORMANCE_TESTS=1 bun test \
2425
--test-name-pattern 'streams a complete long recording'
2526

2627
RUN MEDIA_SERVER_TRANSFER_PERFORMANCE_TESTS=1 bun test \
27-
src/__tests__/lib/media-transfer.integration.test.ts
28+
src/__tests__/lib/media-transfer.integration.test.ts \
29+
src/__tests__/lib/drive-resumable-upload.test.ts
2830

2931
ENV PORT=3456
3032
EXPOSE 3456
Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
import { afterEach, expect, test } from "bun:test";
2+
import { mkdtemp, open, rm } from "node:fs/promises";
3+
import { tmpdir } from "node:os";
4+
import { join } from "node:path";
5+
import { uploadDriveResumable } from "../../lib/drive-resumable-upload";
6+
7+
const originalFetch = globalThis.fetch;
8+
const chunk = 32 * 1024 ** 2;
9+
const dirs: string[] = [];
10+
afterEach(async () => {
11+
globalThis.fetch = originalFetch;
12+
await Promise.all(
13+
dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })),
14+
);
15+
});
16+
17+
async function sparseFile(size: number) {
18+
const dir = await mkdtemp(join(tmpdir(), "cap-drive-upload-"));
19+
dirs.push(dir);
20+
const path = join(dir, "output.mp4");
21+
const handle = await open(path, "w");
22+
try {
23+
await handle.truncate(size);
24+
await handle.write(new Uint8Array([99]), 0, 1, size - 1);
25+
} finally {
26+
await handle.close();
27+
}
28+
return Bun.file(path);
29+
}
30+
31+
function incomplete(end?: number) {
32+
return new Response(null, {
33+
status: 308,
34+
headers: end === undefined ? {} : { Range: `bytes=0-${end - 1}` },
35+
});
36+
}
37+
38+
test("slices real files beyond 8 GiB without truncating offsets or allocating the file", async () => {
39+
const size = 9 * 1024 ** 3 + 123;
40+
const body = await sparseFile(size);
41+
let received = 0;
42+
globalThis.fetch = (async (_url, init) => {
43+
const piece = init?.body as Blob;
44+
const end = Math.min(received + chunk, size);
45+
expect(piece.size).toBe(end - received);
46+
expect(new Headers(init?.headers).get("Content-Range")).toBe(
47+
`bytes ${received}-${end - 1}/${size}`,
48+
);
49+
if (end === size)
50+
expect(new Uint8Array(await piece.slice(-1).arrayBuffer())[0]).toBe(99);
51+
received = end;
52+
return received === size
53+
? Response.json({ size: String(size) })
54+
: incomplete(received);
55+
}) as typeof fetch;
56+
const result = await uploadDriveResumable(
57+
"https://drive.test/upload",
58+
body,
59+
"video/mp4",
60+
);
61+
expect(received).toBe(size);
62+
expect(await result.json()).toEqual({ size: String(size) });
63+
});
64+
65+
test.each(["disconnect", "503"])(
66+
"resumes confirmed partial bytes after %s without replaying the prefix",
67+
async (mode) => {
68+
const body = await sparseFile(chunk + 123);
69+
const requests: string[] = [];
70+
globalThis.fetch = (async (_url, init) => {
71+
requests.push(new Headers(init?.headers).get("Content-Range") ?? "");
72+
if (requests.length === 1) {
73+
if (mode === "disconnect") throw new Error("Connection lost");
74+
return new Response(null, { status: 503 });
75+
}
76+
if (requests.length === 2) return incomplete(chunk / 2);
77+
expect((init?.body as Blob).size).toBe(chunk / 2 + 123);
78+
return Response.json({ done: true });
79+
}) as typeof fetch;
80+
await uploadDriveResumable("https://drive.test/upload", body, "video/mp4");
81+
expect(requests).toEqual([
82+
`bytes 0-${chunk - 1}/${body.size}`,
83+
`bytes */${body.size}`,
84+
`bytes ${chunk / 2}-${body.size - 1}/${body.size}`,
85+
]);
86+
},
87+
);
88+
89+
test("reconciles a lost final response without sending the output again", async () => {
90+
const requests: string[] = [];
91+
globalThis.fetch = (async (_url, init) => {
92+
requests.push(new Headers(init?.headers).get("Content-Range") ?? "");
93+
if (requests.length === 1) throw new Error("Final response lost");
94+
return Response.json({ done: true });
95+
}) as typeof fetch;
96+
await uploadDriveResumable(
97+
"https://drive.test/upload",
98+
new Blob(["abc"]),
99+
"video/mp4",
100+
);
101+
expect(requests).toEqual(["bytes 0-2/3", "bytes */3"]);
102+
});
103+
104+
test.each(["bytes=0-999", "bytes=1-2", "nonsense", "bytes=0-9007199254740993"])(
105+
"rejects untrustworthy offsets %s",
106+
async (range) => {
107+
globalThis.fetch = (async () =>
108+
new Response(null, {
109+
status: 308,
110+
headers: { Range: range },
111+
})) as typeof fetch;
112+
await expect(
113+
uploadDriveResumable(
114+
"https://drive.test/upload",
115+
new Blob(["abc"]),
116+
"video/mp4",
117+
),
118+
).rejects.toThrow("invalid upload offset");
119+
},
120+
);
121+
122+
test("bounds an upload that never acknowledges progress", async () => {
123+
let requests = 0;
124+
globalThis.fetch = (async () => {
125+
requests++;
126+
return incomplete();
127+
}) as typeof fetch;
128+
await expect(
129+
uploadDriveResumable(
130+
"https://drive.test/upload",
131+
new Blob(["abc"]),
132+
"video/mp4",
133+
),
134+
).rejects.toThrow("no progress");
135+
expect(requests).toBe(9);
136+
}, 10000);
137+
138+
test("cancels backoff without resending media", async () => {
139+
const controller = new AbortController();
140+
let requests = 0;
141+
globalThis.fetch = (async () => {
142+
requests++;
143+
controller.abort();
144+
throw new Error("lost");
145+
}) as typeof fetch;
146+
await expect(
147+
uploadDriveResumable(
148+
"https://drive.test/upload",
149+
new Blob(["abc"]),
150+
"video/mp4",
151+
undefined,
152+
controller.signal,
153+
),
154+
).rejects.toThrow();
155+
expect(requests).toBe(1);
156+
});
157+
158+
test.skipIf(process.env.MEDIA_SERVER_TRANSFER_PERFORMANCE_TESTS !== "1")(
159+
"streams a 9 GiB output over HTTP within bounded memory",
160+
async () => {
161+
const size = 9 * 1024 ** 3 + 123;
162+
const body = await sparseFile(size);
163+
let received = 0;
164+
let lastByte = 0;
165+
let requests = 0;
166+
const baseline = process.memoryUsage().rss;
167+
let peak = baseline;
168+
const started = performance.now();
169+
const server = Bun.serve({
170+
port: 0,
171+
hostname: "127.0.0.1",
172+
async fetch(request) {
173+
requests++;
174+
const start = received;
175+
const reader = request.body?.getReader();
176+
if (!reader) {
177+
console.log(
178+
JSON.stringify({
179+
event: "drive_status_probe",
180+
received,
181+
range: request.headers.get("Content-Range"),
182+
}),
183+
);
184+
return incomplete(received || undefined);
185+
}
186+
for (;;) {
187+
const next = await reader.read();
188+
if (next.done) break;
189+
received += next.value.length;
190+
lastByte = next.value[next.value.length - 1];
191+
}
192+
peak = Math.max(peak, process.memoryUsage().rss);
193+
expect(request.headers.get("Content-Range")).toBe(
194+
`bytes ${start}-${received - 1}/${size}`,
195+
);
196+
return received === size
197+
? Response.json({ done: true })
198+
: incomplete(received);
199+
},
200+
});
201+
try {
202+
await uploadDriveResumable(
203+
`http://127.0.0.1:${server.port}/upload`,
204+
body,
205+
"video/mp4",
206+
);
207+
expect(received).toBe(size);
208+
expect(lastByte).toBe(99);
209+
expect(requests).toBe(Math.ceil(size / chunk));
210+
expect(peak - baseline).toBeLessThan(384 * 1024 ** 2);
211+
console.log(
212+
JSON.stringify({
213+
event: "drive_large_upload_verified",
214+
bytes: received,
215+
requests,
216+
elapsedMs: Math.round(performance.now() - started),
217+
peakRssGrowthMiB: Math.round((peak - baseline) / 1024 ** 2),
218+
}),
219+
);
220+
} finally {
221+
await server.stop(true);
222+
}
223+
},
224+
120000,
225+
);
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import { setTimeout as sleep } from "node:timers/promises";
2+
import { UPLOAD_TIMEOUT_MS } from "./media-common";
3+
4+
const CHUNK_BYTES = 32 * 1024 * 1024;
5+
const MAX_RETRIES = 4;
6+
const MAX_RETRANSMITTED_BYTES = 16 * CHUNK_BYTES;
7+
8+
export async function uploadDriveResumable(
9+
url: string,
10+
body: Blob,
11+
contentType: string,
12+
ifNoneMatch?: "*",
13+
abortSignal?: AbortSignal,
14+
): Promise<Response> {
15+
const size = body.size;
16+
if (!Number.isSafeInteger(size) || size <= 0)
17+
throw new Error("Invalid Drive upload size");
18+
let offset = 0;
19+
let sentThrough = 0;
20+
let transmitted = 0;
21+
let failures = 0;
22+
let queryStatus = false;
23+
for (;;) {
24+
abortSignal?.throwIfAborted();
25+
const end = Math.min(offset + CHUNK_BYTES, size);
26+
const querying = queryStatus;
27+
const headers: Record<string, string> = {
28+
"Content-Type": contentType,
29+
"Content-Length": querying ? "0" : String(end - offset),
30+
"Content-Range": querying
31+
? `bytes */${size}`
32+
: `bytes ${offset}-${end - 1}/${size}`,
33+
};
34+
if (ifNoneMatch) headers["If-None-Match"] = ifNoneMatch;
35+
if (!querying) {
36+
transmitted += end - offset;
37+
sentThrough = Math.max(sentThrough, end);
38+
if (transmitted > size + MAX_RETRANSMITTED_BYTES)
39+
throw new Error("Drive upload exceeded its retransmission limit");
40+
}
41+
let response: Response | undefined;
42+
let failure: unknown;
43+
try {
44+
response = await fetch(url, {
45+
method: "PUT",
46+
headers,
47+
body: querying ? undefined : body.slice(offset, end),
48+
redirect: "manual",
49+
signal: abortSignal
50+
? AbortSignal.any([
51+
abortSignal,
52+
AbortSignal.timeout(UPLOAD_TIMEOUT_MS),
53+
])
54+
: AbortSignal.timeout(UPLOAD_TIMEOUT_MS),
55+
});
56+
} catch (error) {
57+
abortSignal?.throwIfAborted();
58+
failure = error;
59+
}
60+
if (response?.status === 200 || response?.status === 201) {
61+
if (sentThrough === size) return response;
62+
await response.body?.cancel().catch(() => {});
63+
throw new Error("Drive completed an upload before all bytes were sent");
64+
}
65+
if (response?.status === 308) {
66+
const range = response.headers.get("range");
67+
await response.body?.cancel().catch(() => {});
68+
const match = range?.match(/^bytes=0-(\d+)$/i);
69+
const nextOffset =
70+
range === null ? 0 : match ? Number(match[1]) + 1 : NaN;
71+
if (
72+
!Number.isSafeInteger(nextOffset) ||
73+
nextOffset < offset ||
74+
nextOffset > sentThrough
75+
)
76+
throw new Error("Drive returned an invalid upload offset");
77+
if (nextOffset > offset) {
78+
offset = nextOffset;
79+
failures = 0;
80+
queryStatus = offset === size;
81+
continue;
82+
}
83+
if (querying && offset < size) {
84+
queryStatus = false;
85+
continue;
86+
}
87+
failure = new Error("Drive upload made no progress");
88+
} else if (response) {
89+
await response.body?.cancel().catch(() => {});
90+
failure = new Error(`Drive upload failed: HTTP ${response.status}`);
91+
if (![408, 425, 429, 500, 502, 503, 504].includes(response.status))
92+
throw failure;
93+
}
94+
if (++failures > MAX_RETRIES)
95+
throw failure instanceof Error
96+
? failure
97+
: new Error("Drive upload failed after retries");
98+
queryStatus = true;
99+
await sleep(250 * 2 ** (failures - 1), undefined, { signal: abortSignal });
100+
}
101+
}

0 commit comments

Comments
 (0)