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
This commit is contained in:
Miguel Ángel
2026-10-22 20:21:55 +00:00
committed by GitHub
parent a8ac55dc30
commit 77394fe9e5
2 changed files with 120 additions and 56 deletions
+68 -23
View File
@@ -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<string, unknown> }[] = [];
readonly partSizes: number[] = [];
head: Record<string, unknown> | 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<string, { bytes: Buffer; sha256: string; lastModified?: Date }>();
readonly copies: Record<string, unknown>[] = [];
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,
+52 -33
View File
@@ -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<void> {
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<CompletedPart> {
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<T>(
path: string,
range: ByteRange | undefined,
send: (body: ReadStream) => Promise<T>,
): Promise<T> {
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<string, unknown> {
return value !== null && typeof value === "object" && !Array.isArray(value);
}