From 77394fe9e52e972123d0868e90f2d56530ef8443 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20=C3=81ngel?= Date: Thu, 22 Oct 2026 20:21:55 +0000 Subject: [PATCH] fix(aws-lambda): retry a file upload after a dropped connection (#6082) * fix(aws-lambda): retry a file upload after a dropped connection * fix(aws-lambda): keep retrying an S3 request timeout --- packages/aws-lambda/src/s3Transport.test.ts | 91 +++++++++++++++------ packages/aws-lambda/src/s3Transport.ts | 85 +++++++++++-------- 2 files changed, 120 insertions(+), 56 deletions(-) diff --git a/packages/aws-lambda/src/s3Transport.test.ts b/packages/aws-lambda/src/s3Transport.test.ts index af6281eef7..ebc5e6f283 100644 --- a/packages/aws-lambda/src/s3Transport.test.ts +++ b/packages/aws-lambda/src/s3Transport.test.ts @@ -24,10 +24,19 @@ import { untarDirectory, putContentStoreFile, uploadContentAddressedFileToS3, + uploadFileToS3, } from "./s3Transport.js"; let scratchRoot: string; +/** An S3 error as the SDK throws it; a dropped connection carries no HTTP status. */ +function s3Error(name: string, httpStatusCode?: number): Error { + const error = Object.assign(new Error(name), { name }); + return httpStatusCode === undefined + ? error + : Object.assign(error, { $metadata: { httpStatusCode } }); +} + beforeAll(() => { scratchRoot = mkdtempSync(join(tmpdir(), "hf-s3transport-test-")); }); @@ -134,6 +143,7 @@ describe("content-addressed v2 artifacts", () => { uploadContentAddressedFileToS3(s3.asClient(), source, uri, digest), ).rejects.toMatchObject({ name: "BadDigest" }); expect(s3.objects.has(uri)).toBe(false); + expect(s3.putAttempts).toBe(1); }); it("refuses to overwrite an immutable key with conflicting metadata", async () => { @@ -278,6 +288,34 @@ describe("presignVerifiedS3GetUrl", () => { }); }); +const TRANSIENT_FAILURES: [string, Error][] = [ + ["a dropped connection", s3Error("TimeoutError")], + ["an S3 request timeout (HTTP 400)", s3Error("RequestTimeout", 400)], +]; + +describe("single-PUT uploads", () => { + it.each(TRANSIENT_FAILURES)("resends the file from a fresh read after %s", async (_, failure) => { + const source = join(scratchRoot, "single-retry.bin"); + writeFileSync(source, "chunk bytes"); + const s3 = new ContentAddressedFakeS3(); + s3.failNextPut = failure; + await uploadFileToS3(s3.asClient(), source, "s3://bucket/chunks/0.mp4"); + expect(s3.objects.get("s3://bucket/chunks/0.mp4")?.bytes.toString()).toBe("chunk bytes"); + expect(s3.putAttempts).toBe(2); + }); + + it("does not resend after any other S3 answer", async () => { + const source = join(scratchRoot, "single-final.bin"); + writeFileSync(source, "chunk bytes"); + const s3 = new ContentAddressedFakeS3(); + s3.failNextPut = s3Error("InvalidArgument", 400); + await expect( + uploadFileToS3(s3.asClient(), source, "s3://bucket/chunks/0.mp4"), + ).rejects.toMatchObject({ name: "InvalidArgument" }); + expect(s3.putAttempts).toBe(1); + }); +}); + describe("objects over the single-PUT limit", () => { const bigSize = MAX_SINGLE_PUT_BYTES + 1; // sha256 of bigSize zero bytes, computed outside this code. @@ -315,20 +353,24 @@ describe("objects over the single-PUT limit", () => { expect(s3.maxPartsInFlight).toBeLessThanOrEqual(8); }, 120_000); - it("resends a part from a fresh read after a dropped connection", async () => { - const source = sparseFile("big-retry.bin"); - const s3 = new MultipartFakeS3(); - s3.failNextPart = true; - try { - expect(await putContentStoreFile(s3.asClient(), source, "bucket", bigDigest)).toBe( - "uploaded", - ); - } finally { - rmSync(source); - } - expect(s3.partSizes.reduce((a, b) => a + b, 0)).toBe(bigSize); - expect(s3.calls.at(-1)).toMatchObject({ name: "CompleteMultipartUploadCommand" }); - }, 120_000); + it.each(TRANSIENT_FAILURES)( + "resends a part from a fresh read after %s", + async (_, failure) => { + const source = sparseFile("big-retry.bin"); + const s3 = new MultipartFakeS3(); + s3.failNextPart = failure; + try { + expect(await putContentStoreFile(s3.asClient(), source, "bucket", bigDigest)).toBe( + "uploaded", + ); + } finally { + rmSync(source); + } + expect(s3.partSizes.reduce((a, b) => a + b, 0)).toBe(bigSize); + expect(s3.calls.at(-1)).toMatchObject({ name: "CompleteMultipartUploadCommand" }); + }, + 120_000, + ); it("aborts instead of completing when the file no longer has the digest it was stored under", async () => { const source = sparseFile("big-changed.bin"); @@ -388,7 +430,7 @@ class MultipartFakeS3 { readonly calls: { name: string; input: Record }[] = []; readonly partSizes: number[] = []; head: Record | undefined; - failNextPart = false; + failNextPart: Error | undefined; maxPartsInFlight = 0; #inFlight = 0; @@ -433,12 +475,11 @@ class MultipartFakeS3 { hash.update(chunk); length += chunk.length; } - if (this.failNextPart) { - this.failNextPart = false; - throw Object.assign(new Error("connection reset"), { name: "ECONNRESET" }); - } + const failure = this.failNextPart; + this.failNextPart = undefined; + if (failure) throw failure; if (length !== input.ContentLength || hash.digest("base64") !== input.ChecksumSHA256) { - throw Object.assign(new Error("part checksum mismatch"), { name: "BadDigest" }); + throw s3Error("BadDigest", 400); } this.partSizes.push(length); return { ETag: `"part-${String(input.PartNumber)}"` }; @@ -460,6 +501,8 @@ class ContentAddressedFakeS3 { readonly objects = new Map(); readonly copies: Record[] = []; putCount = 0; + putAttempts = 0; + failNextPut: Error | undefined; raceOnNextPut: { bytes: Buffer; sha256: string } | undefined; asClient(): import("@aws-sdk/client-s3").S3Client { @@ -533,11 +576,13 @@ class ContentAddressedFakeS3 { const chunks: Buffer[] = []; for await (const chunk of value.input.Body ?? []) chunks.push(Buffer.from(chunk)); const bytes = Buffer.concat(chunks); + this.putAttempts += 1; + const failure = this.failNextPut; + this.failNextPut = undefined; + if (failure) throw failure; const checksum = createHash("sha256").update(bytes).digest("base64"); if (value.input.ChecksumSHA256 !== undefined && value.input.ChecksumSHA256 !== checksum) { - const error = new Error("checksum mismatch"); - error.name = "BadDigest"; - throw error; + throw s3Error("BadDigest", 400); } this.objects.set(uri, { bytes, diff --git a/packages/aws-lambda/src/s3Transport.ts b/packages/aws-lambda/src/s3Transport.ts index 8601b5bb18..8bc824e208 100644 --- a/packages/aws-lambda/src/s3Transport.ts +++ b/packages/aws-lambda/src/s3Transport.ts @@ -22,6 +22,7 @@ import { createWriteStream, existsSync, mkdirSync, + type ReadStream, rmSync, statSync, } from "node:fs"; @@ -52,7 +53,7 @@ export const MAX_SINGLE_PUT_BYTES = 5 * 1024 ** 3; const MAX_PARTS = 10_000; const MIN_FILE_PART_BYTES = 64 * 1024 * 1024; const FILE_PARTS_IN_FLIGHT = 8; -const PART_ATTEMPTS = 3; +const UPLOAD_ATTEMPTS = 3; /** Parsed `s3://bucket/key` URI. */ export interface S3Location { @@ -180,14 +181,16 @@ export async function uploadFileToS3( }); return; } - await client.send( - new PutObjectCommand({ - Bucket: bucket, - Key: key, - Body: createReadStream(localPath), - ContentType: contentType, - ContentLength: size, - }), + await sendFileBody(localPath, undefined, (body) => + client.send( + new PutObjectCommand({ + Bucket: bucket, + Key: key, + Body: body, + ContentType: contentType, + ContentLength: size, + }), + ), ); } @@ -254,9 +257,8 @@ async function putImmutableBody( expectedSha256: string, contentType?: string, ): Promise { - const body = createReadStream(localPath); - try { - await client.send( + await sendFileBody(localPath, undefined, (body) => + client.send( new PutObjectCommand({ Bucket: bucket, Key: key, @@ -267,13 +269,8 @@ async function putImmutableBody( ChecksumSHA256: Buffer.from(expectedSha256, "hex").toString("base64"), IfNoneMatch: "*", }), - ); - } finally { - // A failed conditional request may reject before consuming the stream. - // Explicit teardown avoids retaining the source descriptor on a warm - // Lambda planner. - body.destroy(); - } + ), + ); } export type MultipartTarget = { Bucket: string; Key: string; UploadId: string }; @@ -386,7 +383,6 @@ async function hashRange(path: string, { start, end }: ByteRange, whole: Hash): return part.digest("base64"); } -// A streamed body cannot be replayed by the SDK's own retry, so each attempt opens a fresh stream. async function sendFilePart( client: S3Client, target: MultipartTarget, @@ -395,22 +391,34 @@ async function sendFilePart( { start, end }: ByteRange, checksum: string, ): Promise { + const result = await sendFileBody(path, { start, end }, (body) => + client.send( + new UploadPartCommand({ + ...target, + PartNumber: partNumber, + Body: body, + ContentLength: end - start, + ChecksumSHA256: checksum, + }), + ), + ); + if (!result.ETag) throw new Error(`[s3Transport] part ${partNumber} returned no ETag`); + return { ETag: result.ETag, PartNumber: partNumber, ChecksumSHA256: checksum }; +} + +// The SDK never retries a streamed body, so each attempt opens a fresh stream. Only a transient +// failure is retried: any other S3 answer (412, BadDigest) would repeat. +async function sendFileBody( + path: string, + range: ByteRange | undefined, + send: (body: ReadStream) => Promise, +): Promise { for (let attempt = 1; ; attempt++) { - const body = createReadStream(path, { start, end: end - 1 }); + const body = createReadStream(path, range && { start: range.start, end: range.end - 1 }); try { - const result = await client.send( - new UploadPartCommand({ - ...target, - PartNumber: partNumber, - Body: body, - ContentLength: end - start, - ChecksumSHA256: checksum, - }), - ); - if (!result.ETag) throw new Error(`[s3Transport] part ${partNumber} returned no ETag`); - return { ETag: result.ETag, PartNumber: partNumber, ChecksumSHA256: checksum }; + return await send(body); } catch (error) { - if (attempt >= PART_ATTEMPTS) throw error; + if (attempt >= UPLOAD_ATTEMPTS || !isTransientS3Error(error)) throw error; } finally { body.destroy(); } @@ -618,6 +626,17 @@ function isS3PreconditionFailed(error: unknown): boolean { return error.name === "PreconditionFailed" || metadata?.httpStatusCode === 412; } +// The AWS SDK's own transient codes; S3 sends RequestTimeout as HTTP 400. +const TRANSIENT_S3_ERRORS = new Set(["TimeoutError", "RequestTimeout", "RequestTimeoutException"]); + +/** A dropped connection (no HTTP status), a timeout, throttling (429) or a 5xx. */ +function isTransientS3Error(error: unknown): boolean { + if (!isRecord(error)) return false; + if (typeof error.name === "string" && TRANSIENT_S3_ERRORS.has(error.name)) return true; + const status = isRecord(error.$metadata) ? error.$metadata.httpStatusCode : undefined; + return typeof status !== "number" || status === 429 || status >= 500; +} + function isRecord(value: unknown): value is Record { return value !== null && typeof value === "object" && !Array.isArray(value); }