From 32b3648758e9ac3832efa0acb194a6e4fdbbe3ba Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sun, 8 Nov 2026 20:42:12 -0800 Subject: [PATCH] fix(storage): back off and reduce parallel parts when S3 asks to slow down (#20165) A workspace snapshot upload sends up to eight file-backed parts at once. The SDK never retries a part because its body is a stream, so a part that received 503 SlowDown had only the uploader's three fixed-delay attempts (0.5 s and 1 s). Under a request-rate limit the save failed and the run failed in setup. A rate-limited part now retries with exponential full-jitter backoff (up to eight attempts, at most 20 s per wait), and the upload halves the parts it admits, once per round of throttled requests. Other part failures keep the existing three attempts. Tests cover a real-client SlowDown recovery, the single halving, the attempt bound, cancellation during a wait, and the backoff schedule through its ceiling. Co-Authored-By: Claude Opus 5.5 --- .../storage-s3-multipart-backoff.test.ts | 54 ++++++ .../storage-s3-multipart-throttle.test.ts | 164 ++++++++++++++++++ server/src/storage/s3-multipart.ts | 44 ++++- 3 files changed, 255 insertions(+), 7 deletions(-) create mode 100644 server/src/__tests__/storage-s3-multipart-backoff.test.ts create mode 100644 server/src/__tests__/storage-s3-multipart-throttle.test.ts diff --git a/server/src/__tests__/storage-s3-multipart-backoff.test.ts b/server/src/__tests__/storage-s3-multipart-backoff.test.ts new file mode 100644 index 00000000000..7fccdc437d1 --- /dev/null +++ b/server/src/__tests__/storage-s3-multipart-backoff.test.ts @@ -0,0 +1,54 @@ +import { CreateMultipartUploadCommand, UploadPartCommand, S3Client } from "@aws-sdk/client-s3"; +import type { Readable } from "node:stream"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it, vi } from "vitest"; +import { THROTTLED_PART_ATTEMPTS, putS3MultipartFile } from "../storage/s3-multipart.js"; + +// Record each requested wait and resolve it at once, so the schedule is +// observable without waiting for it. +const waits = vi.hoisted(() => [] as number[]); +vi.mock("node:timers/promises", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + setTimeout: (ms: number, value?: unknown, options?: { signal?: AbortSignal }) => { + waits.push(ms); + return actual.setTimeout(0, value, options); + }, + }; +}); + +const slowDown = () => Object.assign(new Error("Please reduce your request rate."), { name: "SlowDown", $metadata: { httpStatusCode: 503 } }); + +describe("file-backed S3 multipart backoff schedule", () => { + it("doubles the jittered wait for each throttled attempt up to the 20 s ceiling", async () => { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "pc-s3-backoff-")); + const file = path.join(dir, "source.bin"), size = 3 * 8 * 1024 * 1024; + await fs.writeFile(file, Buffer.alloc(size, 7), { mode: 0o600 }); + const client = new S3Client({ region: "us-east-1" }); + let partOneAttempts = 0; + const random = vi.spyOn(Math, "random").mockReturnValue(0.5); + const send = vi.spyOn(client, "send").mockImplementation(async (command: any) => { + if (command instanceof CreateMultipartUploadCommand) return { UploadId: "upload" }; + if (command instanceof UploadPartCommand) { + (command.input.Body as Readable).resume(); + if (command.input.PartNumber === 1 && ++partOneAttempts < THROTTLED_PART_ATTEMPTS) throw slowDown(); + return { ETag: `part-${command.input.PartNumber}` }; + } + return {}; + }); + waits.length = 0; + try { + await expect(putS3MultipartFile(client, "bucket", "key", { objectKey: "key", contentType: "application/octet-stream", + contentLength: size, body: Buffer.alloc(0), sourceFile: file })).resolves.toEqual({ integrity: "server_verified_md5" }); + expect(partOneAttempts).toBe(THROTTLED_PART_ATTEMPTS); + // Half of each ceiling: 1, 2, 4, 8, 16, then 20, 20 seconds. + expect(waits).toEqual([500, 1_000, 2_000, 4_000, 8_000, 10_000, 10_000]); + } finally { + random.mockRestore(); send.mockRestore(); client.destroy(); + await fs.rm(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/server/src/__tests__/storage-s3-multipart-throttle.test.ts b/server/src/__tests__/storage-s3-multipart-throttle.test.ts new file mode 100644 index 00000000000..0572359cd79 --- /dev/null +++ b/server/src/__tests__/storage-s3-multipart-throttle.test.ts @@ -0,0 +1,164 @@ +import { AbortMultipartUploadCommand, CreateMultipartUploadCommand, UploadPartCommand, S3Client } from "@aws-sdk/client-s3"; +import { Readable } from "node:stream"; +import { createHash } from "node:crypto"; +import { createServer } from "node:http"; +import type { AddressInfo } from "node:net"; +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { describe, expect, it, vi } from "vitest"; +import { FILE_PART_CONCURRENCY, THROTTLED_PART_ATTEMPTS, putS3MultipartFile } from "../storage/s3-multipart.js"; +import { createS3StorageProvider } from "../storage/s3-provider.js"; +import { createStorageService } from "../storage/service.js"; + +const partBytes = 8 * 1024 * 1024; + +async function sourceFile(parts: number) { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "pc-s3-throttle-")); + const file = path.join(dir, "source.bin"); + const handle = await fs.open(file, "wx", 0o600); + try { + for (let index = 0; index < parts; index++) await handle.write(Buffer.alloc(partBytes, index + 1)); + } finally { await handle.close(); } + return { dir, file, size: parts * partBytes }; +} + +const md5 = (bytes: Buffer) => createHash("md5").update(bytes).digest("base64"); + +async function drain(body: unknown) { + const chunks: Buffer[] = []; + for await (const chunk of body as Readable) chunks.push(Buffer.from(chunk)); + return Buffer.concat(chunks); +} + +const slowDown = () => Object.assign(new Error("Please reduce your request rate."), { name: "SlowDown", $metadata: { httpStatusCode: 503 } }); + +describe("file-backed S3 multipart uploads under a request-rate limit", () => { + it("backs off and completes when the real S3 client receives SlowDown for one part", async () => { + const source = await sourceFile(3); + const attempts = new Map(); + let completed = ""; + const server = createServer((request, response) => { + const url = new URL(request.url!, "http://localhost"); + const chunks: Buffer[] = []; + request.on("data", chunk => chunks.push(chunk)); + request.on("end", () => { + if (url.searchParams.has("uploads")) { + response.end("upload"); + } else if (request.method === "PUT") { + const number = Number(url.searchParams.get("partNumber")); + attempts.set(number, (attempts.get(number) ?? 0) + 1); + if (number === 2 && attempts.get(number)! <= 3) { + response.writeHead(503, { "Content-Type": "application/xml" }) + .end("SlowDownPlease reduce your request rate.throttled"); + } else { + expect(md5(Buffer.concat(chunks))).toBe(request.headers["content-md5"]); + response.writeHead(200, { ETag: `"part-${number}"` }).end(); + } + } else if (request.method === "POST") { + completed = Buffer.concat(chunks).toString(); + response.end("complete"); + } else { + response.writeHead(500).end(); + } + }); + }); + await new Promise((resolve, reject) => { server.once("error", reject); server.listen(0, "127.0.0.1", resolve); }); + vi.stubEnv("AWS_ACCESS_KEY_ID", "fixture"); vi.stubEnv("AWS_SECRET_ACCESS_KEY", "fixture"); + const random = vi.spyOn(Math, "random").mockReturnValue(0); + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const service = createStorageService(createS3StorageProvider({ bucket: "bucket", region: "us-east-1", forcePathStyle: true, + endpoint: `http://127.0.0.1:${(server.address() as AddressInfo).port}` })); + try { + const result = await service.putFile({ companyId: "company", objectKey: "company/workspace-snapshots/throttled.pcws", namespace: "workspace-snapshots", + originalFilename: null, contentType: "application/octet-stream", body: Readable.from([Buffer.alloc(0)]), byteSize: source.size, + sha256: createHash("sha256").update(await fs.readFile(source.file)).digest("hex"), sourceFile: source.file }); + expect(result.integrity).toBe("server_verified_md5"); + expect(attempts.get(2)).toBe(4); + expect([...completed.matchAll(/(\d+)<\/PartNumber>/g)].map(match => Number(match[1]))).toEqual([1, 2, 3]); + } finally { + random.mockRestore(); warn.mockRestore(); vi.unstubAllEnvs(); + server.closeAllConnections(); + await new Promise(resolve => server.close(() => resolve())); + await fs.rm(source.dir, { recursive: true, force: true }); + } + }); + + it("halves new parts in flight once when a parallel burst of parts is throttled", async () => { + const source = await sourceFile(16); + const client = new S3Client({ region: "us-east-1" }); + const attempts = new Map(); + let laterActive = 0, laterPeak = 0; + const random = vi.spyOn(Math, "random").mockReturnValue(0); + const send = vi.spyOn(client, "send").mockImplementation(async (command: any) => { + if (command instanceof CreateMultipartUploadCommand) return { UploadId: "upload" }; + if (command instanceof UploadPartCommand) { + const number = command.input.PartNumber!; + attempts.set(number, (attempts.get(number) ?? 0) + 1); + await drain(command.input.Body); + if (number <= FILE_PART_CONCURRENCY && attempts.get(number) === 1) throw slowDown(); + if (number > FILE_PART_CONCURRENCY) { + laterActive++; laterPeak = Math.max(laterPeak, laterActive); + await new Promise(resolve => setTimeout(resolve, 20)); + laterActive--; + } + return { ETag: `part-${number}` }; + } + if (command instanceof AbortMultipartUploadCommand) throw new Error("must not abort a recovered upload"); + return {}; + }); + try { + await expect(putS3MultipartFile(client, "bucket", "key", { objectKey: "key", contentType: "application/octet-stream", + contentLength: source.size, body: Buffer.alloc(0), sourceFile: source.file })).resolves.toEqual({ integrity: "server_verified_md5" }); + expect([...attempts.keys()].sort((a, b) => a - b)).toEqual([...Array(16).keys()].map(n => n + 1)); + // Eight simultaneous throttles are one round: four slots remain, not one. + expect(laterPeak).toBe(FILE_PART_CONCURRENCY / 2); + } finally { random.mockRestore(); send.mockRestore(); client.destroy(); await fs.rm(source.dir, { recursive: true, force: true }); } + }); + + it("stops after a bounded number of throttled attempts and aborts once", async () => { + const source = await sourceFile(3); + const client = new S3Client({ region: "us-east-1" }); + const random = vi.spyOn(Math, "random").mockReturnValue(0); + const send = vi.spyOn(client, "send").mockImplementation(async (command: any) => { + if (command instanceof CreateMultipartUploadCommand) return { UploadId: "upload" }; + if (command instanceof UploadPartCommand) { + await drain(command.input.Body); + if (command.input.PartNumber === 1) throw slowDown(); + return { ETag: `part-${command.input.PartNumber}` }; + } + if (command instanceof AbortMultipartUploadCommand) return {}; + throw new Error("must not complete a failed upload"); + }); + try { + await expect(putS3MultipartFile(client, "bucket", "key", { objectKey: "key", contentType: "application/octet-stream", + contentLength: source.size, body: Buffer.alloc(0), sourceFile: source.file })).rejects.toMatchObject({ name: "SlowDown" }); + const partOne = send.mock.calls.filter(([command]) => command instanceof UploadPartCommand && command.input.PartNumber === 1); + expect(partOne).toHaveLength(THROTTLED_PART_ATTEMPTS); + expect(send.mock.calls.filter(([command]) => command instanceof AbortMultipartUploadCommand)).toHaveLength(1); + } finally { random.mockRestore(); send.mockRestore(); client.destroy(); await fs.rm(source.dir, { recursive: true, force: true }); } + }); + + it("cancels promptly while a throttled part is backing off", async () => { + const source = await sourceFile(3); + const client = new S3Client({ region: "us-east-1" }); + const cancellation = new AbortController(); + const random = vi.spyOn(Math, "random").mockReturnValue(0.99); + const send = vi.spyOn(client, "send").mockImplementation(async (command: any) => { + if (command instanceof CreateMultipartUploadCommand) return { UploadId: "upload" }; + if (command instanceof UploadPartCommand) { + await drain(command.input.Body); + if (command.input.PartNumber === 1) { setTimeout(() => cancellation.abort(new Error("run cancelled")), 10); throw slowDown(); } + return { ETag: `part-${command.input.PartNumber}` }; + } + return {}; + }); + const started = Date.now(); + try { + await expect(putS3MultipartFile(client, "bucket", "key", { objectKey: "key", contentType: "application/octet-stream", + contentLength: source.size, body: Buffer.alloc(0), sourceFile: source.file, signal: cancellation.signal })).rejects.toThrow("run cancelled"); + expect(Date.now() - started).toBeLessThan(900); + expect(send.mock.calls.filter(([command]) => command instanceof UploadPartCommand && command.input.PartNumber === 1)).toHaveLength(1); + } finally { random.mockRestore(); send.mockRestore(); client.destroy(); await fs.rm(source.dir, { recursive: true, force: true }); } + }); +}); diff --git a/server/src/storage/s3-multipart.ts b/server/src/storage/s3-multipart.ts index 0c187233eae..408443c9c1f 100644 --- a/server/src/storage/s3-multipart.ts +++ b/server/src/storage/s3-multipart.ts @@ -129,8 +129,23 @@ export async function putS3Multipart(client: S3Client, bucket: string, key: stri export const FILE_PART_CONCURRENCY = 8; const FILE_PART_ATTEMPTS = 3; +/** A provider that asks for a lower request rate gets more, longer-spaced attempts. */ +export const THROTTLED_PART_ATTEMPTS = 8; +const THROTTLE_BACKOFF_BASE_MS = 1_000; +const THROTTLE_BACKOFF_MAX_MS = 20_000; const FILE_READ_BYTES = 1024 * 1024; +/** + * S3 answers a request-rate limit with 503 SlowDown; compatible gateways also + * use 429. The SDK never retries a part because its body is a stream, so the + * upload must slow down itself. + */ +function isRequestRateLimited(error: unknown) { + if (!error || typeof error !== "object") return false; + const { name, $metadata } = error as { name?: unknown; $metadata?: { httpStatusCode?: unknown } }; + return name === "SlowDown" || $metadata?.httpStatusCode === 503 || $metadata?.httpStatusCode === 429; +} + /** Base64 MD5 of an inclusive byte range, as the Content-MD5 header expects. */ export async function fileRangeMd5(file: string, start: number, end: number, signal?: AbortSignal) { const digest = createHash("md5"); @@ -146,7 +161,9 @@ export async function assertSourceFileSize(file: string, contentLength: number) /** A file source is read per part instead of through one stream. Parts upload * in parallel without the shared memory window, each carries Content-MD5 so * the backend rejects bytes it did not receive intact, and a failed part is - * resent from disk instead of failing the whole upload. */ + * resent from disk instead of failing the whole upload. When the backend + * limits the request rate, the part backs off with jitter and the upload + * halves its parts in flight, once per round of throttled requests. */ export async function putS3MultipartFile( client: S3Client, bucket: string, key: string, input: PutObjectInput & { sourceFile: string }, ): Promise { @@ -167,12 +184,17 @@ export async function putS3MultipartFile( if (!uploadId) throw new Error("S3 did not return a multipart upload id"); const id = uploadId; const parts: { ETag: string; PartNumber: number }[] = new Array(count); + // Workers at or above this slot stop taking new parts. A reduction counts + // only for requests sent after the previous one, so a burst of parallel + // throttled parts halves the rate once rather than once per part. + let concurrency = FILE_PART_CONCURRENCY, reductions = 0; const sendPart = async (index: number) => { const start = index * partSize, length = Math.min(partSize, input.contentLength - start), end = start + length - 1; const contentMd5 = await fileRangeMd5(file, start, end, controller.signal); - for (let attempt = 1; ; attempt++) { + for (let failures = 0, throttles = 0; ;) { controller.signal.throwIfAborted(); const body = createReadStream(file, { start, end, highWaterMark: 64 * 1024 }); + const sentAfter = reductions; try { const result = await withS3FailureDiagnostics("upload_part", () => client.send(new UploadPartCommand({ Bucket: bucket, Key: key, UploadId: id, PartNumber: index + 1, Body: body, ContentLength: length, ContentMD5: contentMd5, @@ -182,21 +204,29 @@ export async function putS3MultipartFile( return; } catch (error) { // A part read from disk can be resent; the stream path cannot. - if (controller.signal.aborted || attempt >= FILE_PART_ATTEMPTS) throw error; - await delay(attempt * 500, undefined, { signal: controller.signal }); + if (controller.signal.aborted) throw error; + if (isRequestRateLimited(error)) { + if (++throttles >= THROTTLED_PART_ATTEMPTS) throw error; + if (sentAfter === reductions) { concurrency = Math.max(1, Math.floor(concurrency / 2)); reductions++; } + const ceiling = Math.min(THROTTLE_BACKOFF_MAX_MS, THROTTLE_BACKOFF_BASE_MS * 2 ** (throttles - 1)); + await delay(Math.floor(Math.random() * ceiling), undefined, { signal: controller.signal }); + continue; + } + if (++failures >= FILE_PART_ATTEMPTS) throw error; + await delay(failures * 500, undefined, { signal: controller.signal }); } finally { body.destroy(); } } }; let next = 0; - const worker = async () => { - while (next < count) { + const worker = async (slot: number) => { + while (next < count && slot < concurrency) { const index = next++; try { await sendPart(index); } catch (error) { if (!controller.signal.aborted) controller.abort(error); throw error; } } }; // Wait for every sibling to settle before the multipart upload is aborted. - const settled = await Promise.allSettled(Array.from({ length: Math.min(FILE_PART_CONCURRENCY, count) }, worker)); + const settled = await Promise.allSettled(Array.from({ length: Math.min(FILE_PART_CONCURRENCY, count) }, (_, slot) => worker(slot))); const rejected = settled.find((entry): entry is PromiseRejectedResult => entry.status === "rejected"); if (rejected) throw controller.signal.aborted && controller.signal.reason !== undefined ? controller.signal.reason : rejected.reason; controller.signal.throwIfAborted();