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,