From 51f9fa7c541fa2a8f57ae700d4189396058bbd62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20=C3=81ngel?= Date: Wed, 28 Oct 2026 09:03:42 +0000 Subject: [PATCH] perf(aws-lambda): small-chunk cloud renders stitch their output in parallel (#6635) * perf(aws-lambda): small-chunk cloud renders stitch their output in parallel When every chunk's payload is under S3's 5 MiB part minimum (typical for short 4K chunks), the assemble step could not copy any part server-side and pulled the whole output through one Lambda's network link. Those range-only parts now go to parallel invocations of the warm function that just rendered the chunks, about 32 MiB each and never more at once than the stack's reserved concurrency allows. The output bytes are unchanged. * fix(aws-lambda): a throttled stitch helper builds its parts in the assemble instead Another render can hold the stack's reserved concurrency while this one assembles, so a helper invocation may be refused. The assemble now builds that batch itself, logs it, and counts it in assemble_stages; the helper's error type is kept as the error name. Comments trimmed back to each file's share at the base. * chore(aws-lambda): raise the package size budget to the measured size plus headroom The parallel stitch helpers add about 1.7 KB packed and 9.9 KB unpacked; the new budget is the measured size times the checker's own 1.05 headroom. --- examples/aws-lambda/template.yaml | 12 ++ examples/aws-lambda/video-function.yaml | 5 + .../HyperframesRenderStack.snapshot.test.ts | 34 +++++ .../src/cdk/HyperframesRenderStack.ts | 15 +- packages/aws-lambda/src/events.ts | 15 ++ packages/aws-lambda/src/handler.test.ts | 102 +++++++++++++- packages/aws-lambda/src/handler.ts | 78 ++++++++++- packages/aws-lambda/src/s3Stitch.test.ts | 76 ++++++++++- packages/aws-lambda/src/s3Stitch.ts | 129 +++++++++++++++--- scripts/package-size-budgets.json | 4 +- 10 files changed, 438 insertions(+), 32 deletions(-) diff --git a/examples/aws-lambda/template.yaml b/examples/aws-lambda/template.yaml index 604138b020..20743446e1 100644 --- a/examples/aws-lambda/template.yaml +++ b/examples/aws-lambda/template.yaml @@ -170,9 +170,21 @@ Resources: TMPDIR: /tmp HYPERFRAMES_RENDER_BUCKET: !Ref RenderBucket HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: !Ref ReservedConcurrency + # Assemble spreads the parts S3 cannot copy across the warm function that ran the chunks. + HYPERFRAMES_STITCH_FUNCTION: !If + - HasVideoFunction + - !Sub "${ProjectName}-video" + - !Sub "${ProjectName}-render" Policies: - S3CrudPolicy: BucketName: !Ref RenderBucket + - Statement: + - Effect: Allow + Action: lambda:InvokeFunction + Resource: !If + - HasVideoFunction + - !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-video" + - !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-render" - Statement: - Effect: Allow Action: s3:AbortMultipartUpload diff --git a/examples/aws-lambda/video-function.yaml b/examples/aws-lambda/video-function.yaml index de8046432b..9b2b0bb7ee 100644 --- a/examples/aws-lambda/video-function.yaml +++ b/examples/aws-lambda/video-function.yaml @@ -62,6 +62,10 @@ Resources: - Effect: Allow Action: lambda:GetAccountSettings Resource: "*" + # Stitch helpers, should this function ever run an assemble. + - Effect: Allow + Action: lambda:InvokeFunction + Resource: !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-video" Tags: - Key: Project Value: !Ref ProjectName @@ -95,6 +99,7 @@ Resources: TMPDIR: /tmp HYPERFRAMES_RENDER_BUCKET: !Ref RenderBucket HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: !Ref ReservedConcurrency + HYPERFRAMES_STITCH_FUNCTION: !Sub "${ProjectName}-video" Tags: - Key: Project Value: !Ref ProjectName diff --git a/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts b/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts index 71198f2330..bdc51a8c01 100644 --- a/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts +++ b/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts @@ -493,6 +493,40 @@ describe("HyperframesRenderStack — snapshot", () => { expect(samActions).toContain("lambda:GetAccountSettings"); }, 30000); + it("lets assemble invoke exactly the stitch helper it names, in CDK and SAM", () => { + const cdkFunction = Object.values( + SYNTHED.template.findResources("AWS::Lambda::Function"), + )[0] as { + Properties: { Environment: { Variables: Record } }; + }; + expect(cdkFunction.Properties.Environment.Variables.HYPERFRAMES_STITCH_FUNCTION).toBe( + "hyperframes-render", + ); + const cdkInvoke = Object.values(SYNTHED.template.findResources("AWS::IAM::Policy")) + .flatMap(policyStatements) + .filter((statement) => statement.Action === "lambda:InvokeFunction"); + expect(JSON.stringify(cdkInvoke.map((s) => s.Resource))).toContain( + ":function:hyperframes-render", + ); + + const samFunction = requireRecordProperty( + requireRecordProperty(readSamResources(), "RenderFunction", "SAM RenderFunction"), + "Properties", + "SAM function properties", + ) as { Environment: { Variables: Record }; Policies: unknown[] }; + // The YAML parser drops the !If and !Sub tags and keeps [condition, then, else]. + const [, withVideo, without] = samFunction.Environment.Variables + .HYPERFRAMES_STITCH_FUNCTION as string[]; + expect([withVideo, without]).toEqual(["${ProjectName}-video", "${ProjectName}-render"]); + const samInvoke = samFunction.Policies.flatMap( + (policy) => (policy as { Statement?: Record[] }).Statement ?? [], + ).filter((statement) => statement.Action === "lambda:InvokeFunction"); + expect(samInvoke).toHaveLength(1); + const [, grantWithVideo, grantWithout] = samInvoke[0]!.Resource as string[]; + expect(grantWithVideo!.endsWith(`:function:${withVideo}`)).toBe(true); + expect(grantWithout!.endsWith(`:function:${without}`)).toBe(true); + }, 30000); + it("fails at once when a chunk cannot fit the worker timeout, in CDK and SAM", () => { for (const definition of [SYNTHED.definition, readSamDefinition()]) { for (const name of ["Plan", "PlanV2"]) { diff --git a/packages/aws-lambda/src/cdk/HyperframesRenderStack.ts b/packages/aws-lambda/src/cdk/HyperframesRenderStack.ts index 864ab61b43..8dc9d61830 100644 --- a/packages/aws-lambda/src/cdk/HyperframesRenderStack.ts +++ b/packages/aws-lambda/src/cdk/HyperframesRenderStack.ts @@ -21,7 +21,7 @@ import { dirname, resolve } from "node:path"; import { fileURLToPath } from "node:url"; -import { Duration, RemovalPolicy, Size } from "aws-cdk-lib"; +import { ArnFormat, Duration, RemovalPolicy, Size, Stack } from "aws-cdk-lib"; import * as cloudwatch from "aws-cdk-lib/aws-cloudwatch"; import * as iam from "aws-cdk-lib/aws-iam"; import * as lambda from "aws-cdk-lib/aws-lambda"; @@ -120,6 +120,7 @@ export class HyperframesRenderStack extends Construct { HYPERFRAMES_LAMBDA_CHROME_SOURCE: chromeSource, HYPERFRAMES_RENDER_BUCKET: this.bucket.bucketName, HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: String(props.reservedConcurrency ?? -1), + HYPERFRAMES_STITCH_FUNCTION: `${projectName}-render`, }, }); @@ -131,6 +132,18 @@ export class HyperframesRenderStack extends Construct { this.renderFunction.addToRolePolicy( new iam.PolicyStatement({ actions: ["lambda:GetAccountSettings"], resources: ["*"] }), ); + const renderFunctionArnWithoutRoleCycle = Stack.of(this).formatArn({ + service: "lambda", + resource: "function", + resourceName: `${projectName}-render`, + arnFormat: ArnFormat.COLON_RESOURCE_NAME, + }); + this.renderFunction.addToRolePolicy( + new iam.PolicyStatement({ + actions: ["lambda:InvokeFunction"], + resources: [renderFunctionArnWithoutRoleCycle], + }), + ); const stateMachineLogGroup = new logs.LogGroup(this, "RenderStateMachineLogGroup", { logGroupName: `/aws/states/${projectName}-render`, diff --git a/packages/aws-lambda/src/events.ts b/packages/aws-lambda/src/events.ts index 7fe1b2e0f1..6727152fe2 100644 --- a/packages/aws-lambda/src/events.ts +++ b/packages/aws-lambda/src/events.ts @@ -28,6 +28,14 @@ export type LambdaAction = "plan" | "audio" | "renderChunk" | "assemble"; export interface WarmEvent { Action: "warm"; } +/** Assemble hands a helper part of the final mp4's multipart upload: each part's S3 ranges, in order. */ +export interface StitchPartsEvent { + Action: "stitchParts"; + Bucket: string; + Key: string; + UploadId: string; + Parts: { partNumber: number; ranges: { uri: string; start: number; end: number }[] }[]; +} /** Transport protocol selected for one complete distributed render. */ export type LambdaPlanProtocol = "v1" | "v2"; @@ -45,6 +53,7 @@ export type LambdaEvent = | RenderChunkEvent | AssembleEvent | WarmEvent + | StitchPartsEvent | { Payload: LambdaEvent } | { Input: LambdaEvent }; @@ -239,8 +248,14 @@ export interface WarmLambdaResult { Action: "warm"; } +export interface StitchPartsLambdaResult { + Action: "stitchParts"; + ETags: string[]; +} + export type LambdaResult = | WarmLambdaResult + | StitchPartsLambdaResult | PlanLambdaResult | AudioLambdaResult | RenderChunkLambdaResult diff --git a/packages/aws-lambda/src/handler.test.ts b/packages/aws-lambda/src/handler.test.ts index 013c531ffd..873e1328fe 100644 --- a/packages/aws-lambda/src/handler.test.ts +++ b/packages/aws-lambda/src/handler.test.ts @@ -38,8 +38,14 @@ import { } from "@hyperframes/producer/distributed"; import { resolveMaxParallelChunks } from "../../producer/src/services/distributed/plan.js"; import { recomputePlanHashFromPlanDir } from "../../producer/src/services/render/stages/freezePlan.js"; -import type { AssembleEvent, LambdaEvent, PlanEvent, RenderChunkEvent } from "./events.js"; -import { handler as handleEvent, unwrapEvent } from "./handler.js"; +import type { + AssembleEvent, + LambdaEvent, + PlanEvent, + RenderChunkEvent, + StitchPartsEvent, +} from "./events.js"; +import { handler as handleEvent, stitchHelper, unwrapEvent } from "./handler.js"; import { S3PlanV2ArtifactPublisher } from "./s3PlanV2Publisher.js"; interface FakeS3Op { @@ -1626,6 +1632,98 @@ describe("handler — S3 URI allowlist (security: F-004)", () => { await expect(handler(event, deps)).rejects.toMatchObject({ name: "S3_URI_NOT_ALLOWED" }); expect(s3.ops).toHaveLength(0); }); + + it("rejects a stitchParts range outside the allowed bucket", async () => { + process.env.HYPERFRAMES_RENDER_BUCKET = "good-bucket"; + const s3 = new FakeS3Client(); + const event: StitchPartsEvent = { + Action: "stitchParts", + Bucket: "good-bucket", + Key: "renders/abc/output.mp4", + UploadId: "upload-1", + Parts: [{ partNumber: 2, ranges: [{ uri: "s3://evil-bucket/secret", start: 0, end: 4 }] }], + }; + await expect( + handler(event, { s3: s3 as unknown as S3Client, skipChromeResolution: true }), + ).rejects.toMatchObject({ name: "S3_URI_NOT_ALLOWED" }); + expect(s3.ops).toHaveLength(0); + }); +}); + +describe("handler — stitch helpers", () => { + const event = (uri: string): StitchPartsEvent => ({ + Action: "stitchParts", + Bucket: "bucket", + Key: "renders/abc/output.mp4", + UploadId: "upload-1", + Parts: [ + { partNumber: 3, ranges: [{ uri, start: 1, end: 4 }] }, + { + partNumber: 4, + ranges: [ + { uri, start: 0, end: 2 }, + { uri, start: 4, end: 6 }, + ], + }, + ], + }); + + it("builds the parts a stitchParts call names and returns their ETags in order", async () => { + const s3 = new FakeS3Client(); + s3.objects.set("s3://bucket/chunks/0.mp4", Buffer.from("abcdef")); + const result = await handler(event("s3://bucket/chunks/0.mp4"), { + s3: s3 as unknown as S3Client, + skipChromeResolution: true, + }); + expect(result).toEqual({ Action: "stitchParts", ETags: ['"part-3"', '"part-4"'] }); + expect([s3.parts.get(3)?.toString(), s3.parts.get(4)?.toString()]).toEqual(["bcd", "abef"]); + }); + + it("the stitch helper invokes the function the stack names and fails with its error", async () => { + const prev = process.env.HYPERFRAMES_STITCH_FUNCTION; + const prevReserved = process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY; + delete process.env.HYPERFRAMES_STITCH_FUNCTION; + expect(stitchHelper()).toBeUndefined(); + process.env.HYPERFRAMES_STITCH_FUNCTION = "proj-video"; + try { + // A reservation of 8 leaves 7 slots once this assemble holds one. + process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = "8"; + expect(stitchHelper()!.maxInvocations).toBe(7); + process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = "-1"; + expect(stitchHelper()!.maxInvocations).toBe(Infinity); + const s3 = new FakeS3Client(); + s3.objects.set("s3://bucket/chunks/0.mp4", Buffer.from("abcdef")); + const invoked: string[] = []; + const lambda = { + send: async (command: { input: { FunctionName: string; Payload: string } }) => { + invoked.push(command.input.FunctionName); + const sent = JSON.parse(command.input.Payload) as StitchPartsEvent; + if (sent.Parts[0]!.ranges[0]!.uri.endsWith("missing")) { + const payload = JSON.stringify({ errorMessage: "NoSuchKey", errorType: "NoSuchKey" }); + return { FunctionError: "Unhandled", Payload: new TextEncoder().encode(payload) }; + } + const out = await handler(sent, { s3: s3 as unknown as S3Client }); + return { Payload: new TextEncoder().encode(JSON.stringify(out)) }; + }, + } as unknown as LambdaClient; + const { send } = stitchHelper(lambda)!; + const target = { Bucket: "bucket", Key: "renders/abc/output.mp4", UploadId: "upload-1" }; + expect(await send(target, event("s3://bucket/chunks/0.mp4").Parts)).toEqual([ + '"part-3"', + '"part-4"', + ]); + await expect(send(target, event("s3://bucket/missing").Parts)).rejects.toMatchObject({ + name: "NoSuchKey", + message: expect.stringContaining("stitchParts on proj-video failed: NoSuchKey"), + }); + expect(invoked).toEqual(["proj-video", "proj-video"]); + } finally { + if (prev === undefined) delete process.env.HYPERFRAMES_STITCH_FUNCTION; + else process.env.HYPERFRAMES_STITCH_FUNCTION = prev; + if (prevReserved === undefined) delete process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY; + else process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = prevReserved; + } + }); }); // ── helpers ───────────────────────────────────────────────────────────────── diff --git a/packages/aws-lambda/src/handler.ts b/packages/aws-lambda/src/handler.ts index 7e2d72b559..684f4407bb 100644 --- a/packages/aws-lambda/src/handler.ts +++ b/packages/aws-lambda/src/handler.ts @@ -14,7 +14,7 @@ import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, statSync } from "node:fs"; import { availableParallelism, tmpdir } from "node:os"; import { basename, join } from "node:path"; -import { GetAccountSettingsCommand, LambdaClient } from "@aws-sdk/client-lambda"; +import { GetAccountSettingsCommand, InvokeCommand, LambdaClient } from "@aws-sdk/client-lambda"; import { S3Client } from "@aws-sdk/client-s3"; import { assemble, @@ -44,7 +44,7 @@ import { renderChunk, } from "@hyperframes/producer/distributed"; import { resolveChromeExecutablePath } from "./chromium.js"; -import { stitchMp4InS3 } from "./s3Stitch.js"; +import { type StitchHelper, sendRangeParts, stitchMp4InS3 } from "./s3Stitch.js"; import { type DistributedFormat, formatExtension } from "./formatExtension.js"; import type { AssembleEvent, @@ -61,6 +61,8 @@ import type { RenderChunkV2Event, RenderChunkLambdaResult, AssembleV2Event, + StitchPartsEvent, + StitchPartsLambdaResult, } from "./events.js"; import { downloadS3ObjectToFile, @@ -126,6 +128,9 @@ export async function handler(event: LambdaEvent, deps?: HandlerDeps): Promise 0 ? reserved : undefined; +} + /** The function's reserved concurrency, else its share of the account's unreserved pool. */ async function functionConcurrency(lambda = new LambdaClient({})): Promise { - const reserved = Number(process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY); - if (reserved > 0) return reserved; + const reserved = reservedConcurrency(); + if (reserved) return reserved; const { AccountLimit } = await lambda.send(new GetAccountSettingsCommand({})); const unreserved = AccountLimit?.UnreservedConcurrentExecutions; if (unreserved === undefined) throw new Error("no UnreservedConcurrentExecutions in the reply"); @@ -859,7 +870,7 @@ async function handleAssembleV2( throw error; } if (chunksAreStitchedMp4(planDir, event.Cfr === true)) { - return await stitchChunksInS3(s3, event, manifest, started); + return await stitchChunksInS3(s3, event, manifest, started, deps?.lambda); } // `downloadAndMaterializePlanV2` materializes atomically. Audio is // assembler-only and lives at the familiar v1-compatible location. @@ -910,6 +921,7 @@ async function stitchChunksInS3( event: AssembleV2Event, plan: PlanV2Manifest, started: number, + lambda?: LambdaClient, ): Promise { if (event.ChunkS3Uris.length !== plan.chunkCount) { throw new Error( @@ -919,6 +931,7 @@ async function stitchChunksInS3( const stitchStarted = Date.now(); const stitched = await stitchMp4InS3(s3, event.ChunkS3Uris, event.OutputS3Uri, { audioUri: event.AudioS3Uri ?? undefined, + helper: stitchHelper(lambda), }); logEvent({ event: "assemble_stages", @@ -927,6 +940,8 @@ async function stitchChunksInS3( stitchMs: Date.now() - stitchStarted, parts: stitched.parts, downloadedBytes: stitched.downloadedBytes, + helpers: stitched.helpers, + helperFallbacks: stitched.helperFallbacks, }); return { Action: "assemble", @@ -1095,6 +1110,53 @@ function getEventS3Uris( } } +/** + * Parts S3 cannot copy go to the function the stack names: it just ran the chunks, so it is warm, + * and each invocation brings its own network link. + */ +export function stitchHelper(lambda?: LambdaClient): StitchHelper | undefined { + const functionName = process.env.HYPERFRAMES_STITCH_FUNCTION?.trim(); + if (!functionName) return undefined; + const client = lambda ?? new LambdaClient({}); + // This assemble may hold one slot of the reservation itself. + const reserved = reservedConcurrency(); + const maxInvocations = reserved ? Math.max(1, reserved - 1) : Infinity; + const send: StitchHelper["send"] = async (target, parts) => { + const event: StitchPartsEvent = { Action: "stitchParts", ...target, Parts: parts }; + const response = await client.send( + new InvokeCommand({ FunctionName: functionName, Payload: JSON.stringify(event) }), + ); + const body = JSON.parse( + new TextDecoder().decode(response.Payload), + ) as StitchPartsLambdaResult & { + errorMessage?: string; + errorType?: string; + }; + if (response.FunctionError) { + const error = new Error( + `[handler] stitchParts on ${functionName} failed: ${body.errorMessage ?? response.FunctionError}`, + ); + if (body.errorType) error.name = body.errorType; + throw error; + } + return body.ETags; + }; + return { send, maxInvocations }; +} + +async function handleStitchParts( + event: StitchPartsEvent, + deps?: HandlerDeps, +): Promise { + const { Bucket, Key, UploadId, Parts } = event; + assertAllowedBuckets([ + `s3://${Bucket}/${Key}`, + ...Parts.flatMap((part) => part.ranges.map((r) => r.uri)), + ]); + const ETags = await sendRangeParts(deps?.s3 ?? getS3Client(), { Bucket, Key, UploadId }, Parts); + return { Action: "stitchParts", ETags }; +} + /** * Verify every S3 URI in the event resolves to the configured render bucket. * Throws `S3_URI_NOT_ALLOWED` (non-retryable) when a URI targets a different @@ -1106,10 +1168,14 @@ function getEventS3Uris( function validateEventS3Uris( event: PlanEvent | AudioEvent | RenderChunkEvent | AssembleEvent, ): void { + assertAllowedBuckets(getEventS3Uris(event)); +} + +function assertAllowedBuckets(uris: Iterable): void { const allowedBucket = process.env.HYPERFRAMES_RENDER_BUCKET?.trim(); if (!allowedBucket) return; - for (const uri of getEventS3Uris(event)) { + for (const uri of uris) { const { bucket } = parseS3Uri(uri); if (bucket !== allowedBucket) { const err = new Error( diff --git a/packages/aws-lambda/src/s3Stitch.test.ts b/packages/aws-lambda/src/s3Stitch.test.ts index a807ef8f6d..e4b8df2b5d 100644 --- a/packages/aws-lambda/src/s3Stitch.test.ts +++ b/packages/aws-lambda/src/s3Stitch.test.ts @@ -6,7 +6,15 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import type { S3Client } from "@aws-sdk/client-s3"; import { assemble, muxChunkWithAudio } from "@hyperframes/producer/distributed"; -import { type Part, planParts, stitchMp4InS3, uploadParts } from "./s3Stitch.js"; +import { + helperBatches, + type Part, + planParts, + type RangePart, + sendRangeParts, + stitchMp4InS3, + uploadParts, +} from "./s3Stitch.js"; const MiB = 1024 * 1024; const range = (uri: string, start: number, end: number) => ({ uri, start, end }); @@ -45,6 +53,27 @@ describe("planParts", () => { }); }); +describe("helperBatches", () => { + it("hands helpers only the parts made of S3 ranges, about 32 MiB each", () => { + const parts: Part[] = [ + { upload: [{ bytes: Buffer.alloc(10) }, range("s3://b/0", 0, MiB)] }, + { copy: range("s3://b/1", 0, 6 * MiB) }, + ...Array.from({ length: 10 }, (_, i) => ({ + upload: [range(`s3://b/${i + 2}`, 0, 4 * MiB), range(`s3://b/${i + 3}`, 0, 2 * MiB)], + })), + ]; + const numbers = (max?: number) => + helperBatches(parts, max).map((batch) => batch.map((p) => p.partNumber)); + expect(numbers()).toEqual([ + [3, 4, 5, 6, 7, 8], + [9, 10, 11, 12], + ]); + // A reservation of one helper at a time gets one batch, never a throttled second. + expect(numbers(1)).toEqual([[3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]); + expect(numbers(3)).toHaveLength(2); + }); +}); + /** S3 stand-in: ranged reads, and multipart parts concatenated on complete under S3's size rule. */ class FakeS3 { readonly objects = new Map(); @@ -117,6 +146,23 @@ describe("uploadParts", () => { }); expect(s3.sent.some((c) => c.name === "CompleteMultipartUploadCommand")).toBe(false); }); + + it("aborts the upload when a helper fails, after the local parts settle", async () => { + const s3 = new FakeS3(0); + s3.objects.set("s3://bucket/r/chunks/0001.mp4", Buffer.from("abcd")); + const helped: Part[] = [...parts, { upload: [range("s3://bucket/r/chunks/0001.mp4", 0, 4)] }]; + await expect( + uploadParts(s3.asClient(), helped, "s3://bucket/r/out.mp4", "video/mp4", { + send: async () => { + throw new Error("helper failed"); + }, + maxInvocations: Infinity, + }), + ).rejects.toThrow("helper failed"); + const names = s3.sent.map((c) => c.name); + expect(names.filter((n) => n.startsWith("UploadPart"))).toHaveLength(2); + expect(names.at(-1)).toBe("AbortMultipartUploadCommand"); + }); }); const hasFfmpeg = spawnSync("ffmpeg", ["-version"]).status === 0; @@ -260,6 +306,34 @@ describe.skipIf(!hasFfmpeg)("stitchMp4InS3 on real chunks", () => { await assemble(planDir, chunkPaths, null, localPath); expect(readFileSync(localPath).equals(output)).toBe(true); + // The parts S3 cannot copy, built by helper invocations: the same bytes. + const batches: RangePart[][] = []; + await stitchMp4InS3(s3.asClient(), uris, "s3://bucket/helped.mp4", { + minPartBytes: 16 * 1024, + helper: { + send: (target, batch) => { + batches.push(batch); + return sendRangeParts(s3.asClient(), target, batch); + }, + maxInvocations: Infinity, + }, + }); + expect(batches.length).toBeGreaterThan(0); + expect(s3.objects.get("s3://bucket/helped.mp4")!.equals(output)).toBe(true); + + // Every helper throttled (another render holds the reservation): built here, same bytes. + const throttled = await stitchMp4InS3(s3.asClient(), uris, "s3://bucket/throttled.mp4", { + minPartBytes: 16 * 1024, + helper: { + send: async () => { + throw Object.assign(new Error("Rate exceeded"), { name: "TooManyRequestsException" }); + }, + maxInvocations: Infinity, + }, + }); + expect(throttled.helperFallbacks).toBe(batches.length); + expect(s3.objects.get("s3://bucket/throttled.mp4")!.equals(output)).toBe(true); + // Today's join: the video-only chunks plus the plan audio. writeFileSync(join(planDir, "meta", "encoder.json"), "{}"); const referencePath = join(dir, "reference.mp4"); diff --git a/packages/aws-lambda/src/s3Stitch.ts b/packages/aws-lambda/src/s3Stitch.ts index 2d3b23da6c..5f36f6b9e4 100644 --- a/packages/aws-lambda/src/s3Stitch.ts +++ b/packages/aws-lambda/src/s3Stitch.ts @@ -12,6 +12,7 @@ import { buildStitchedHeader, readMp4Layout } from "@hyperframes/producer/distri import { multipartUpload, type MultipartTarget, parseS3Uri } from "./s3Transport.js"; const MIN_PART_BYTES = 5 * 1024 * 1024; +const HELPER_BATCH_BYTES = 32 * 1024 * 1024; /** Bytes `[start, end)` of an S3 object. */ interface RangeSegment { @@ -24,6 +25,16 @@ export type Part = { copy: RangeSegment } | { upload: Segment[] }; const segmentSize = (s: Segment): number => ("bytes" in s ? s.bytes.length : s.end - s.start); +export interface RangePart { + partNumber: number; + ranges: RangeSegment[]; +} +export type SendRangeParts = (target: MultipartTarget, parts: RangePart[]) => Promise; +export interface StitchHelper { + send: SendRangeParts; + maxInvocations: number; +} + // Every part but the last must be >= `min` and a copy reads one object: big ranges are copied, // smaller pieces are downloaded into one upload part that borrows the next range's head if its rest // stays copyable. @@ -85,27 +96,85 @@ async function sendPart(s3: S3Client, part: Part, partNumber: number, target: Mu return result.ETag; } -/** Settles every part before failing, so an abort never races a part still in flight. */ -async function sendParts( - s3: S3Client, - parts: readonly Part[], - target: MultipartTarget, -): Promise { - const settled = await Promise.allSettled( - parts.map(async (part, index) => { - const etag = await sendPart(s3, part, index + 1, target); - if (!etag) throw new Error(`[s3Stitch] part ${index + 1} returned no ETag`); - return etag; - }), - ); +/** Settles every task before failing, so an abort never races a part still in flight. */ +async function settleAll(tasks: readonly Promise[]): Promise { + const settled = await Promise.allSettled(tasks); const failure = settled.find( (result): result is PromiseRejectedResult => result.status === "rejected", ); if (failure) throw failure.reason; - return settled.map((result, index) => ({ - ETag: (result as PromiseFulfilledResult).value, - PartNumber: index + 1, - })); + return settled.map((result) => (result as PromiseFulfilledResult).value); +} + +const rangesOnly = (part: Part): RangeSegment[] | undefined => + "upload" in part && part.upload.every((s) => "uri" in s) + ? (part.upload as RangeSegment[]) + : undefined; + +const rangeBytes = (ranges: readonly RangeSegment[]) => + ranges.reduce((total, r) => total + segmentSize(r), 0); + +// Every closed batch holds at least total / maxBatches bytes, so there are never more than maxBatches. +export function helperBatches(parts: readonly Part[], maxBatches = Infinity): RangePart[][] { + const total = parts.reduce((sum, part) => sum + rangeBytes(rangesOnly(part) ?? []), 0); + const batchBytes = Math.max(HELPER_BATCH_BYTES, Math.ceil(total / maxBatches)); + const batches: RangePart[][] = []; + let batch: RangePart[] = []; + let bytes = 0; + parts.forEach((part, index) => { + const ranges = rangesOnly(part); + if (!ranges) return; + batch.push({ partNumber: index + 1, ranges }); + bytes += rangeBytes(ranges); + if (bytes < batchBytes) return; + batches.push(batch); + batch = []; + bytes = 0; + }); + if (batch.length) batches.push(batch); + return batches; +} + +export async function sendRangeParts( + s3: S3Client, + target: MultipartTarget, + parts: readonly RangePart[], +): Promise { + return settleAll( + parts.map(async ({ partNumber, ranges }) => { + const etag = await sendPart(s3, { upload: ranges }, partNumber, target); + if (!etag) throw new Error(`[s3Stitch] part ${partNumber} returned no ETag`); + return etag; + }), + ); +} + +async function sendParts( + s3: S3Client, + parts: readonly Part[], + target: MultipartTarget, + helper?: StitchHelper, +): Promise { + const batches = helper ? helperBatches(parts, helper.maxInvocations) : []; + const remote = new Set(batches.flat().map((p) => p.partNumber)); + const local = parts.flatMap((part, index) => + remote.has(index + 1) ? [] : [{ partNumber: index + 1, part }], + ); + const done = await settleAll([ + ...local.map(async ({ partNumber, part }) => [ + { PartNumber: partNumber, ETag: await sendPart(s3, part, partNumber, target) }, + ]), + ...batches.map(async (batch) => { + const etags = await helper!.send(target, batch); + return batch.map((p, i) => ({ PartNumber: p.partNumber, ETag: etags[i] })); + }), + ]); + const completed = done.flat().sort((a, b) => a.PartNumber! - b.PartNumber!); + const missing = completed.find((p) => !p.ETag); + if (missing || completed.length !== parts.length) { + throw new Error(`[s3Stitch] part ${missing?.PartNumber ?? "?"} returned no ETag`); + } + return completed; } /** Writes `parts` in order to `destUri`; a failure aborts the upload so no parts linger. */ @@ -114,10 +183,11 @@ export async function uploadParts( parts: readonly Part[], destUri: string, contentType: string, + helper?: StitchHelper, ): Promise { const { bucket, key } = parseS3Uri(destUri); await multipartUpload(s3, { Bucket: bucket, Key: key, ContentType: contentType }, (target) => - sendParts(s3, parts, target), + sendParts(s3, parts, target, helper), ); } @@ -125,6 +195,8 @@ export interface StitchResult { fileSize: number; parts: number; downloadedBytes: number; + helpers: number; + helperFallbacks: number; } /** One faststart mp4 at `destUri` from the stitched mp4 chunks at `chunkUris`, plus `audioUri`'s track when given. */ @@ -132,7 +204,11 @@ export async function stitchMp4InS3( s3: S3Client, chunkUris: readonly string[], destUri: string, - { audioUri, minPartBytes = MIN_PART_BYTES }: { audioUri?: string; minPartBytes?: number } = {}, + { + audioUri, + minPartBytes = MIN_PART_BYTES, + helper, + }: { audioUri?: string; minPartBytes?: number; helper?: StitchHelper } = {}, ): Promise { const sourceUris = audioUri ? [audioUri, ...chunkUris] : chunkUris; const layouts = await Promise.all( @@ -150,11 +226,24 @@ export async function stitchMp4InS3( })), ]; const parts = planParts(segments, minPartBytes); - await uploadParts(s3, parts, destUri, "video/mp4"); + let helperFallbacks = 0; + const helperOrLocal = helper && { + maxInvocations: helper.maxInvocations, + // A throttled helper (another render may hold the reservation) costs time, not the render. + send: (target: MultipartTarget, batch: RangePart[]) => + helper.send(target, batch).catch((error: unknown) => { + helperFallbacks++; + console.warn(`[s3Stitch] helper failed, building ${batch.length} parts here:`, error); + return sendRangeParts(s3, target, batch); + }), + }; + await uploadParts(s3, parts, destUri, "video/mp4", helperOrLocal); const uploaded = parts.flatMap((part) => ("upload" in part ? part.upload : [])); return { fileSize: segments.reduce((total, s) => total + segmentSize(s), 0), parts: parts.length, downloadedBytes: uploaded.reduce((total, s) => total + ("uri" in s ? segmentSize(s) : 0), 0), + helpers: helper ? helperBatches(parts, helper.maxInvocations).length : 0, + helperFallbacks, }; } diff --git a/scripts/package-size-budgets.json b/scripts/package-size-budgets.json index 2e3d2b704a..068dcbaef7 100644 --- a/scripts/package-size-budgets.json +++ b/scripts/package-size-budgets.json @@ -1,7 +1,7 @@ { "@hyperframes/aws-lambda": { - "packed": 73369, - "unpacked": 266805 + "packed": 78851, + "unpacked": 290508 }, "@hyperframes/cli": { "packed": 7902966,