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 <noreply@anthropic.com>

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<typeof import("node:timers/promises")>();
+  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 });
+    }
+  });
+});

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<number, number>();
+    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("<InitiateMultipartUploadResult><UploadId>upload</UploadId></InitiateMultipartUploadResult>");
+        } 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("<Error><Code>SlowDown</Code><Message>Please reduce your request rate.</Message><RequestId>throttled</RequestId></Error>");
+          } 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("<CompleteMultipartUploadResult><ETag>complete</ETag></CompleteMultipartUploadResult>");
+        } else {
+          response.writeHead(500).end();
+        }
+      });
+    });
+    await new Promise<void>((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(/<PartNumber>(\d+)<\/PartNumber>/g)].map(match => Number(match[1]))).toEqual([1, 2, 3]);
+    } finally {
+      random.mockRestore(); warn.mockRestore(); vi.unstubAllEnvs();
+      server.closeAllConnections();
+      await new Promise<void>(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<number, number>();
+    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 }); }
+  });
+});

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<PutObjectReceipt> {
@@ -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();