diff --git a/examples/aws-lambda/template.yaml b/examples/aws-lambda/template.yaml index 8e47da3eb6..45475d7fe8 100644 --- a/examples/aws-lambda/template.yaml +++ b/examples/aws-lambda/template.yaml @@ -162,6 +162,10 @@ Resources: Policies: - S3CrudPolicy: BucketName: !Ref RenderBucket + - Statement: + - Effect: Allow + Action: s3:AbortMultipartUpload + Resource: !Sub "arn:${AWS::Partition}:s3:::${RenderBucket}/*" # CloudWatch Logs perms are covered by SAM's default # AWSLambdaBasicExecutionRole — explicit `CloudWatchLogsFullAccess` # would be overscope (`logs:*` on `*`, including DeleteLogGroup + diff --git a/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts b/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts index ff4cc4ffaa..b9cc767b39 100644 --- a/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts +++ b/packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts @@ -279,6 +279,18 @@ describe("HyperframesRenderStack — snapshot", () => { } }); + it("lets the render function abort the multipart uploads its S3 stitch starts", () => { + const cdkPolicies = JSON.stringify(SYNTHED.template.findResources("AWS::IAM::Policy")); + expect(cdkPolicies).toContain("s3:Abort*"); + const samFunction = requireRecordProperty(readSamResources(), "RenderFunction", "SAM function"); + const samProperties = requireRecordProperty( + samFunction, + "Properties", + "SAM function properties", + ); + expect(JSON.stringify(samProperties.Policies)).toContain("s3:AbortMultipartUpload"); + }); + it("keeps same-named Lambda tasks in different branches apart", () => { const task = (format: string) => ({ Type: "Task", @@ -433,7 +445,7 @@ function getV1TaskStates(definition: { }; } -function readSamDefinition(): { States: Record } { +function readSamResources(): Record { const source = readFileSync( new URL("../../../../examples/aws-lambda/template.yaml", import.meta.url), "utf8", @@ -443,7 +455,11 @@ function readSamDefinition(): { States: Record } { // warnings for the intentionally unresolved `!Ref`/`!GetAtt` tags. const parsed: unknown = parseYaml(source, { logLevel: "silent" }); const root = requireRecord(parsed, "SAM template"); - const resources = requireRecordProperty(root, "Resources", "SAM resources"); + return requireRecordProperty(root, "Resources", "SAM resources"); +} + +function readSamDefinition(): { States: Record } { + const resources = readSamResources(); const stateMachine = requireRecordProperty( resources, "RenderStateMachine", diff --git a/packages/aws-lambda/src/handler.test.ts b/packages/aws-lambda/src/handler.test.ts index 9160725c19..c2c46f441c 100644 --- a/packages/aws-lambda/src/handler.test.ts +++ b/packages/aws-lambda/src/handler.test.ts @@ -16,6 +16,7 @@ */ import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from "bun:test"; +import { spawnSync } from "node:child_process"; import { createHash } from "node:crypto"; import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; @@ -35,7 +36,7 @@ import type { AssembleEvent, LambdaEvent, PlanEvent, RenderChunkEvent } from "./ import { handler, unwrapEvent } from "./handler.js"; interface FakeS3Op { - kind: "download" | "upload"; + kind: "download" | "upload" | "copy"; uri: string; bytes?: number; } @@ -63,12 +64,19 @@ class FakeS3Client { const cmdName = op.constructor?.name ?? ""; const uri = `s3://${op.input.Bucket}/${op.input.Key}`; if (cmdName === "GetObjectCommand") { - const bytes = this.objects.get(uri) ?? Buffer.alloc(0); + const bytes = this.sliceRange( + this.objects.get(uri) ?? Buffer.alloc(0), + (command as { input: { Range?: string } }).input.Range, + ); this.ops.push({ kind: "download", uri, bytes: bytes.length }); // Mock the AWS SDK stream contract just enough for pipeline() to // pump bytes into a write stream. const { Readable } = await import("node:stream"); - return { Body: Readable.from([bytes]) }; + return { + Body: Object.assign(Readable.from([bytes]), { + transformToByteArray: async () => new Uint8Array(bytes), + }), + }; } if (cmdName === "HeadObjectCommand") { const bytes = this.objects.get(uri); @@ -104,8 +112,37 @@ class FakeS3Client { if (metadata) this.metadata.set(uri, metadata); return {}; } + if (cmdName === "CreateMultipartUploadCommand") return { UploadId: "upload-1" }; + if (cmdName === "UploadPartCommand") { + const input = (command as { input: { Body: Buffer; PartNumber: number } }).input; + this.parts.set(input.PartNumber, input.Body); + return { ETag: `"part-${input.PartNumber}"` }; + } + if (cmdName === "UploadPartCopyCommand") { + const input = ( + command as { input: { CopySource: string; CopySourceRange?: string; PartNumber: number } } + ).input; + const source = `s3://${decodeURI(input.CopySource)}`; + this.ops.push({ kind: "copy", uri: source }); + const bytes = this.objects.get(source) ?? Buffer.alloc(0); + this.parts.set(input.PartNumber, this.sliceRange(bytes, input.CopySourceRange)); + return { CopyPartResult: { ETag: `"part-${input.PartNumber}"` } }; + } + if (cmdName === "CompleteMultipartUploadCommand") { + const parts = (command as { input: { MultipartUpload: { Parts: { PartNumber: number }[] } } }) + .input.MultipartUpload.Parts; + this.objects.set(uri, Buffer.concat(parts.map((part) => this.parts.get(part.PartNumber)!))); + return {}; + } return {}; } + + parts = new Map(); + + private sliceRange(bytes: Buffer, range?: string): Buffer { + const match = /^bytes=(\d+)-(\d+)$/.exec(range ?? ""); + return match ? bytes.subarray(Number(match[1]), Number(match[2]) + 1) : bytes; + } } const tmpDirs: string[] = []; @@ -239,6 +276,7 @@ describe("handler dispatch", () => { expect(result.PlanHash).toBe("fakehash"); expect(result.ChunkCount).toBe(4); expect(planMock).toHaveBeenCalledTimes(1); + expect((planMock.mock.calls[0] as unknown[])[3]).toEqual({ chunksCarryAudio: true }); expect(renderChunkMock).not.toHaveBeenCalled(); expect(assembleMock).not.toHaveBeenCalled(); // Plan should have downloaded the project zip and uploaded the plan tar. @@ -730,6 +768,90 @@ describe("handler dispatch", () => { ).toBe(true); }); + async function planStitched(chunkCount: number) { + const tmpRoot = makeTmpRoot(); + const s3 = new FakeS3Client(); + s3.objects.set("s3://bucket/project.tar.gz", await makeMinimalProjectTar()); + const assembleMock = mock( + async (_p: string, _c: readonly string[], _a: string | null, outputPath: string) => { + writeFileSync(outputPath, "LOCAL-JOIN"); + return { outputPath, durationMs: 1, framesEncoded: 30 * chunkCount, fileSize: 10 }; + }, + ); + const planV2WithPublisherMock = mock( + async (_dir: string, _config: unknown, publisher: PlanV2ArtifactPublisher) => { + const v1Dir = join(tmpRoot, "v1"); + makeMinimalV1PlanDir(v1Dir, true, { chunkCount, chunksCarryAudio: true }); + return publishPlanV2FromExecutionPlan(v1Dir, publisher); + }, + ); + const deps = { + s3: s3 as unknown as import("@aws-sdk/client-s3").S3Client, + primitives: { + planV2WithPublisher: + planV2WithPublisherMock as unknown as typeof import("@hyperframes/producer/distributed").planV2WithPublisher, + assemble: + assembleMock as unknown as typeof import("@hyperframes/producer/distributed").assemble, + }, + tmpRoot, + skipChromeResolution: true, + }; + const planned = await handler( + { + Action: "plan", + ProjectS3Uri: "s3://bucket/project.tar.gz", + PlanOutputS3Prefix: "s3://bucket/renders/f/", + Config: { fps: 30, width: 640, height: 360, format: "mp4" }, + }, + deps, + ); + if (!("PlanProtocol" in planned) || planned.PlanProtocol !== "v2") + throw new Error("expected v2"); + const assemble = (chunkS3Uris: string[]) => + handler( + { + Action: "assemble", + PlanV2ManifestS3Uri: planned.PlanV2ManifestS3Uri, + PlanV2ArtifactS3Prefix: planned.PlanV2ArtifactS3Prefix, + PlanHash: planned.PlanHash, + ChunkS3Uris: chunkS3Uris, + AudioS3Uri: null, + OutputS3Uri: "s3://bucket/renders/f/output.mp4", + Format: "mp4", + }, + deps, + ); + return { s3, assembleMock, assemble }; + } + + it.skipIf(spawnSync("ffmpeg", ["-version"]).status !== 0)( + "stitches mp4 chunks inside S3 instead of joining them locally", + async () => { + const { s3, assembleMock, assemble } = await planStitched(2); + const chunks = [0, 1].map((index) => { + const path = join(makeTmpRoot(), "chunk.mp4"); + const lavfi = ["testsrc2=size=64x48:rate=30:duration=1", "sine=duration=1"]; + const muxed = spawnSync("ffmpeg", [ + ...["-nostdin", "-loglevel", "error", "-f", "lavfi", "-i", lavfi[0]!], + ...["-f", "lavfi", "-i", lavfi[1]!, "-c:v", "libx264", "-bf", "0", "-c:a", "aac"], + ...["-shortest", "-y", path], + ]); + if (muxed.status !== 0) throw new Error(String(muxed.stderr)); + const uri = `s3://bucket/renders/f/chunks/000${index}.mp4`; + s3.objects.set(uri, readFileSync(path)); + return uri; + }); + const result = await assemble(chunks); + const output = s3.objects.get("s3://bucket/renders/f/output.mp4")!; + expect(output.toString("latin1", 4, 8)).toBe("ftyp"); + expect( + output.toString("latin1", output.readUInt32BE(0) + 4, output.readUInt32BE(0) + 8), + ).toBe("moov"); + expect(result).toMatchObject({ FramesEncoded: 60, FileSize: output.length }); + expect(assembleMock).not.toHaveBeenCalled(); + }, + ); + it("rejects unknown plan protocol values", async () => { const tmpRoot = makeTmpRoot(); const s3 = new FakeS3Client(); @@ -903,14 +1025,19 @@ async function makeMinimalPlanTar(): Promise { return rf(tarPath); } -function makeMinimalV1PlanDir(dir: string, withAudio: boolean): void { +function makeMinimalV1PlanDir( + dir: string, + withAudio: boolean, + options: { chunkCount?: number; chunksCarryAudio?: boolean } = {}, +): void { + const chunkCount = options.chunkCount ?? 1; mkdirSync(join(dir, "meta"), { recursive: true }); mkdirSync(join(dir, "compiled"), { recursive: true }); writeFileSync(join(dir, "compiled", "index.html"), "aws v2 fixture"); const planJson = { planHash: "a".repeat(64), - chunkCount: 1, - totalFrames: 30, + chunkCount, + totalFrames: 30 * chunkCount, dimensions: { fpsNum: 30, fpsDen: 1, width: 640, height: 360, format: "mp4" }, ffmpegVersion: "6.0", producerVersion: "test", @@ -919,9 +1046,18 @@ function makeMinimalV1PlanDir(dir: string, withAudio: boolean): void { writeFileSync(join(dir, "plan.json"), JSON.stringify(planJson)); writeFileSync( join(dir, "meta", "chunks.json"), - JSON.stringify([{ index: 0, startFrame: 0, endFrame: 30 }]), + JSON.stringify( + Array.from({ length: chunkCount }, (_, index) => ({ + index, + startFrame: index * 30, + endFrame: (index + 1) * 30, + })), + ), + ); + writeFileSync( + join(dir, "meta", "encoder.json"), + options.chunksCarryAudio ? JSON.stringify({ chunksCarryAudio: true }) : "{}", ); - writeFileSync(join(dir, "meta", "encoder.json"), "{}"); if (withAudio) writeFileSync(join(dir, "audio.m4a"), "AAC"); planJson.planHash = recomputePlanHashFromPlanDir(dir); writeFileSync(join(dir, "plan.json"), JSON.stringify(planJson)); diff --git a/packages/aws-lambda/src/handler.ts b/packages/aws-lambda/src/handler.ts index 4c8cbb5e0a..f3f6e4fecd 100644 --- a/packages/aws-lambda/src/handler.ts +++ b/packages/aws-lambda/src/handler.ts @@ -17,6 +17,7 @@ import { basename, join } from "node:path"; import { S3Client } from "@aws-sdk/client-s3"; import { assemble, + chunksAreStitchedMp4, type AssembleResult, type ChunkRenderer, type ChunkResult, @@ -36,6 +37,7 @@ import { renderChunk, } from "@hyperframes/producer/distributed"; import { resolveChromeExecutablePath } from "./chromium.js"; +import { stitchMp4InS3 } from "./s3Stitch.js"; import { type DistributedFormat, formatExtension } from "./formatExtension.js"; import type { AssembleEvent, @@ -365,7 +367,10 @@ async function handlePlan(event: PlanEvent, deps?: HandlerDeps): Promise { + if (event.ChunkS3Uris.length !== plan.chunkCount) { + throw new Error( + `[handler] ${event.ChunkS3Uris.length} chunk URIs for a plan of ${plan.chunkCount} chunks`, + ); + } + const stitchStarted = Date.now(); + const stitched = await stitchMp4InS3(s3, event.ChunkS3Uris, event.OutputS3Uri); + logEvent({ + event: "assemble_stages", + mode: "s3-stitch", + planMs: stitchStarted - started, + stitchMs: Date.now() - stitchStarted, + parts: stitched.parts, + downloadedBytes: stitched.downloadedBytes, + }); + return { + Action: "assemble", + OutputS3Uri: event.OutputS3Uri, + FramesEncoded: plan.totalFrames, + FileSize: stitched.fileSize, + DurationMs: Date.now() - started, + }; +} + async function downloadAndMaterializePlanV2( s3: S3Client, event: { @@ -727,7 +771,7 @@ async function downloadAndMaterializePlanV2( }, target: PlanV2MaterializationTarget, work: string, -): Promise { +): Promise<{ planDir: string; manifest: PlanV2Manifest }> { const transportDir = join(work, "plan-v2"); mkdirSync(transportDir, { recursive: true }); await downloadS3ObjectToFile(s3, event.PlanV2ManifestS3Uri, join(transportDir, "plan.json")); @@ -744,7 +788,7 @@ async function downloadAndMaterializePlanV2( }); const planDir = join(work, "plan"); materializePlanV2Target(transportDir, target, planDir); - return planDir; + return { planDir, manifest }; } async function downloadPlanV2Artifact( diff --git a/packages/aws-lambda/src/s3Stitch.test.ts b/packages/aws-lambda/src/s3Stitch.test.ts new file mode 100644 index 0000000000..74c5adf61e --- /dev/null +++ b/packages/aws-lambda/src/s3Stitch.test.ts @@ -0,0 +1,284 @@ +// fallow-ignore-file code-duplication +import { spawnSync } from "node:child_process"; +import { afterAll, beforeAll, describe, expect, it } from "bun:test"; +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +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"; + +const MiB = 1024 * 1024; +const range = (uri: string, start: number, end: number) => ({ uri, start, end }); + +describe("planParts", () => { + it("copies big ranges and merges the header into a part that borrows from chunk 0", () => { + const header = { bytes: Buffer.alloc(MiB) }; + const parts = planParts([ + header, + range("s3://b/0", 40, 40 + 10 * MiB), + range("s3://b/1", 40, 40 + 6 * MiB), + range("s3://b/2", 40, 1040), + ]); + expect(parts).toEqual([ + { upload: [header, range("s3://b/0", 40, 40 + 4 * MiB)] }, + { copy: range("s3://b/0", 40 + 4 * MiB, 40 + 10 * MiB) }, + { copy: range("s3://b/1", 40, 40 + 6 * MiB) }, + { upload: [range("s3://b/2", 40, 1040)] }, + ]); + }); + + it("downloads a chunk whole when borrowing would leave its rest under the minimum", () => { + const header = { bytes: Buffer.alloc(MiB) }; + const small = range("s3://b/1", 0, 2 * MiB); + const parts = planParts([ + header, + range("s3://b/0", 0, 8 * MiB), + small, + range("s3://b/2", 0, 9 * MiB), + ]); + expect(parts).toEqual([ + { upload: [header, range("s3://b/0", 0, 8 * MiB)] }, + { upload: [small, range("s3://b/2", 0, 3 * MiB)] }, + { copy: range("s3://b/2", 3 * MiB, 9 * MiB) }, + ]); + }); +}); + +/** S3 stand-in: ranged reads, and multipart parts concatenated on complete under S3's size rule. */ +class FakeS3 { + readonly objects = new Map(); + readonly sent: { name: string; input: Record }[] = []; + private readonly parts = new Map(); + constructor( + private readonly minPartBytes: number, + private readonly failPart?: number, + ) {} + + private object(bucket: unknown, key: unknown, byteRange?: unknown): Buffer { + const bytes = this.objects.get(`s3://${bucket}/${key}`); + if (!bytes) throw new Error(`NoSuchKey ${bucket}/${key}`); + const match = /^bytes=(\d+)-(\d+)$/.exec(String(byteRange ?? "")); + return match ? bytes.subarray(Number(match[1]), Number(match[2]) + 1) : bytes; + } + + // fallow-ignore-next-line complexity + async send(command: { constructor: { name: string }; input: Record }) { + const { name } = command.constructor; + const input = command.input; + this.sent.push({ name, input }); + if (this.failPart && input.PartNumber === this.failPart) throw new Error("part failed"); + if (name === "GetObjectCommand") { + const bytes = this.object(input.Bucket, input.Key, input.Range); + return { Body: { transformToByteArray: async () => new Uint8Array(bytes) } }; + } + if (name === "CreateMultipartUploadCommand") return { UploadId: "up-1" }; + if (name === "UploadPartCommand") { + this.parts.set(input.PartNumber as number, input.Body as Buffer); + return { ETag: `"p${input.PartNumber}"` }; + } + if (name === "UploadPartCopyCommand") { + const [bucket, ...key] = decodeURI(String(input.CopySource)).split("/"); + const bytes = this.object(bucket, key.join("/"), input.CopySourceRange); + this.parts.set(input.PartNumber as number, bytes); + return { CopyPartResult: { ETag: `"p${input.PartNumber}"` } }; + } + if (name === "CompleteMultipartUploadCommand") { + const listed = (input.MultipartUpload as { Parts: { PartNumber: number }[] }).Parts; + const parts = listed.map((part) => this.parts.get(part.PartNumber)!); + if (parts.slice(0, -1).some((part) => part.length < this.minPartBytes)) { + throw new Error("EntityTooSmall"); + } + this.objects.set(`s3://${input.Bucket}/${input.Key}`, Buffer.concat(parts)); + } + return {}; + } + + asClient(): S3Client { + return this as unknown as S3Client; + } +} + +describe("uploadParts", () => { + const parts: Part[] = [ + { upload: [{ bytes: Buffer.from("header") }] }, + { copy: range("s3://bucket/r/chunks/0001.mp4", 0, 4) }, + ]; + + it("aborts the upload when a part fails", async () => { + const s3 = new FakeS3(0, 2); + s3.objects.set("s3://bucket/r/chunks/0001.mp4", Buffer.from("abcd")); + await expect( + uploadParts(s3.asClient(), parts, "s3://bucket/r/out.mp4", "video/mp4"), + ).rejects.toThrow("part failed"); + expect(s3.sent.at(-1)).toEqual({ + name: "AbortMultipartUploadCommand", + input: { Bucket: "bucket", Key: "r/out.mp4", UploadId: "up-1" }, + }); + expect(s3.sent.some((c) => c.name === "CompleteMultipartUploadCommand")).toBe(false); + }); +}); + +const hasFfmpeg = spawnSync("ffmpeg", ["-version"]).status === 0; + +type Slice = { index: number; startFrame: number; endFrame: number }; +const slices = (edges: readonly number[]): Slice[] => + edges.slice(1).map((endFrame, index) => ({ index, startFrame: edges[index]!, endFrame })); + +describe.skipIf(!hasFfmpeg)("stitchMp4InS3 on real chunks", () => { + let root: string; + + beforeAll(() => { + root = mkdtempSync(join(tmpdir(), "hf-s3-stitch-")); + }); + afterAll(() => { + rmSync(root, { recursive: true, force: true }); + }); + + function run(cmd: string, args: string[]): string { + const result = spawnSync(cmd, args, { maxBuffer: 1 << 28, encoding: "latin1" }); + if (result.status !== 0) throw new Error(`${cmd} ${args.join(" ")}: ${result.stderr}`); + return result.stdout; + } + const ffmpeg = (args: string[]) => run("ffmpeg", ["-nostdin", "-loglevel", "error", ...args]); + + /** Closed-GOP chunk shaped like renderChunk's: IDR at frame 0, no B-frames. */ + function encodeChunk(path: string, slice: Slice, rate: number): void { + const frames = String(slice.endFrame - slice.startFrame); + ffmpeg( + `-f lavfi -i testsrc2=size=160x120:rate=${rate},trim=start_frame=${slice.startFrame}:end_frame=${slice.endFrame} -c:v libx264 -preset veryfast -g ${frames} -keyint_min ${frames} -sc_threshold 0 -bf 0 -x264-params scenecut=0:open-gop=0:repeat-headers=1 -pix_fmt yuv420p -avoid_negative_ts make_zero -y` + .split(" ") + .concat(path), + ); + } + + const frames = (path: string) => + ffmpeg(["-i", path, "-map", "0:v", "-f", "framemd5", "-"]) + .split("\n") + .filter((line) => line && !line.startsWith("#")); + const pcm = (path: string) => ffmpeg(["-i", path, "-map", "0:a", "-f", "s16le", "-"]); + const probe = (path: string) => + JSON.parse( + run("ffprobe", [ + "-v", + "error", + "-count_packets", + "-show_entries", + "stream=codec_type,nb_read_packets,duration", + "-of", + "json", + "--", + path, + ]), + ).streams as { codec_type: string; nb_read_packets: string; duration: string }[]; + + // Uneven chunks and a 1-frame tail; at 60 fps that tail is shorter than an AAC packet and owns + // none, so its chunk has no audio track. + it.each([ + ["30", 30, [0, 47, 90, 91]], + ["60", 60, [0, 50, 51]], + ] as const)( + "assembles a faststart mp4 at %s fps with the frames and audio of today's join", + async (name, rate, edges) => { + const FPS = { num: rate, den: 1 }; + const CHUNKS = slices(edges); + const TOTAL_FRAMES = edges.at(-1)!; + const dir = join(root, name); + mkdirSync(dir, { recursive: true }); + const audioPath = join(dir, "audio.m4a"); + ffmpeg( + `-f lavfi -i aevalsrc=0.4*sin(2*PI*(220+40*t)*t):s=48000:d=${TOTAL_FRAMES / rate} -c:a aac -b:a 192k -y` + .split(" ") + .concat(audioPath), + ); + const planDir = join(dir, "plan"); + mkdirSync(join(planDir, "meta"), { recursive: true }); + writeFileSync( + join(planDir, "plan.json"), + JSON.stringify({ + planHash: "fake", + totalFrames: TOTAL_FRAMES, + hasAudio: true, + dimensions: { fpsNum: rate, fpsDen: 1, width: 160, height: 120, format: "mp4" }, + }), + ); + writeFileSync(join(planDir, "meta", "chunks.json"), JSON.stringify(CHUNKS)); + + const s3 = new FakeS3(16 * 1024); + const videoPaths: string[] = []; + const chunkPaths: string[] = []; + for (const slice of CHUNKS) { + videoPaths.push(join(dir, `video-${slice.index}.mp4`)); + chunkPaths.push(join(dir, `chunk-${slice.index}.mp4`)); + encodeChunk(videoPaths[slice.index]!, slice, rate); + await muxChunkWithAudio({ + videoPath: videoPaths[slice.index]!, + audioPath, + outputPath: chunkPaths[slice.index]!, + chunks: CHUNKS, + chunkIndex: slice.index, + fps: FPS, + timeoutMs: 60_000, + }); + s3.objects.set( + `s3://bucket/chunks/${slice.index}.mp4`, + readFileSync(chunkPaths[slice.index]!), + ); + } + // The 60 fps tail owns no AAC packet, so its chunk carries video alone. + const tailAudio = probe(chunkPaths.at(-1)!).some((s) => s.codec_type === "audio"); + expect(tailAudio).toBe(rate === 30); + const uris = CHUNKS.map((slice) => `s3://bucket/chunks/${slice.index}.mp4`); + // A 16 KiB part minimum makes these small chunks take every path S3 takes at 5 MiB. + const result = await stitchMp4InS3(s3.asClient(), uris, "s3://bucket/out.mp4", 16 * 1024); + const outputPath = join(dir, "out.mp4"); + const output = s3.objects.get("s3://bucket/out.mp4")!; + writeFileSync(outputPath, output); + expect(result.fileSize).toBe(output.length); + const sent = s3.sent.map((c) => c.name); + expect(sent).toContain("UploadPartCopyCommand"); + expect(sent).toContain("UploadPartCommand"); + + const check = spawnSync("ffprobe", ["-v", "error", "--", outputPath], { encoding: "utf-8" }); + expect([check.status, check.stderr]).toEqual([0, ""]); + const topLevel: string[] = []; + for (let at = 0; at < output.length; ) { + topLevel.push(output.toString("latin1", at + 4, at + 8)); + const size = output.readUInt32BE(at); + at += size === 1 ? Number(output.readBigUInt64BE(at + 8)) : size; + } + expect(topLevel).toEqual(["ftyp", "moov", "mdat"]); + + // The same chunks joined on disk give the same bytes. + writeFileSync( + join(planDir, "meta", "encoder.json"), + JSON.stringify({ chunksCarryAudio: true }), + ); + const localPath = join(dir, "local.mp4"); + await assemble(planDir, chunkPaths, null, localPath); + expect(readFileSync(localPath).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"); + await assemble(planDir, videoPaths, audioPath, referencePath); + expect(frames(outputPath)).toEqual(frames(referencePath)); + expect(pcm(outputPath) === pcm(referencePath)).toBe(true); + + const [video, sound] = probe(outputPath); + const audioPackets = probe(audioPath)[0]!.nb_read_packets; + expect([video!.nb_read_packets, sound!.nb_read_packets]).toEqual([ + String(TOTAL_FRAMES), + audioPackets, + ]); + // Same lengths as today's mux: the edit lists trim only the AAC priming (both round + // through the 1000-tick movie timescale). + const reference = probe(referencePath); + expect([video!.duration, sound!.duration]).toEqual([ + reference[0]!.duration, + reference[1]!.duration, + ]); + }, + 60_000, + ); +}); diff --git a/packages/aws-lambda/src/s3Stitch.ts b/packages/aws-lambda/src/s3Stitch.ts new file mode 100644 index 0000000000..fc4f074a23 --- /dev/null +++ b/packages/aws-lambda/src/s3Stitch.ts @@ -0,0 +1,172 @@ +// Builds the final mp4 inside S3: part 1 carries the stitched header, the other parts copy each +// chunk's mdat payload with UploadPartCopy, so almost no video byte passes through the function. + +import { + AbortMultipartUploadCommand, + CompleteMultipartUploadCommand, + CreateMultipartUploadCommand, + GetObjectCommand, + type S3Client, + UploadPartCommand, + UploadPartCopyCommand, +} from "@aws-sdk/client-s3"; +import { buildStitchedHeader, readMp4Layout } from "@hyperframes/producer/distributed"; +import { parseS3Uri } from "./s3Transport.js"; + +const MIN_PART_BYTES = 5 * 1024 * 1024; + +/** Bytes `[start, end)` of an S3 object. */ +interface RangeSegment { + uri: string; + start: number; + end: number; +} +type Segment = { bytes: Buffer } | RangeSegment; +export type Part = { copy: RangeSegment } | { upload: Segment[] }; + +const segmentSize = (s: Segment): number => ("bytes" in s ? s.bytes.length : s.end - s.start); + +// 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. +// fallow-ignore-next-line complexity +export function planParts(segments: readonly Segment[], min = MIN_PART_BYTES): Part[] { + const parts: Part[] = []; + let pending: Segment[] = []; + let pendingSize = 0; + const flush = () => { + if (pending.length) parts.push({ upload: pending }); + pending = []; + pendingSize = 0; + }; + for (const segment of segments) { + const borrow = pendingSize ? min - pendingSize : 0; + if ("uri" in segment && segmentSize(segment) - borrow >= min) { + if (borrow) pending.push({ ...segment, end: segment.start + borrow }); + flush(); + parts.push({ copy: { ...segment, start: segment.start + borrow } }); + continue; + } + pending.push(segment); + pendingSize += segmentSize(segment); + if (pendingSize >= min) flush(); + } + flush(); + return parts; +} + +async function getRange(s3: S3Client, uri: string, start: number, end: number): Promise { + const { bucket, key } = parseS3Uri(uri); + const response = await s3.send( + new GetObjectCommand({ Bucket: bucket, Key: key, Range: `bytes=${start}-${end - 1}` }), + ); + if (!response.Body) throw new Error(`[s3Stitch] empty body for ${uri}`); + return Buffer.from(await response.Body.transformToByteArray()); +} + +type UploadTarget = { Bucket: string; Key: string; UploadId: string }; + +async function sendPart(s3: S3Client, part: Part, partNumber: number, target: UploadTarget) { + if ("copy" in part) { + const { uri, start, end } = part.copy; + const source = parseS3Uri(uri); + const result = await s3.send( + new UploadPartCopyCommand({ + ...target, + PartNumber: partNumber, + CopySource: encodeURI(`${source.bucket}/${source.key}`), + CopySourceRange: `bytes=${start}-${end - 1}`, + }), + ); + return result.CopyPartResult?.ETag; + } + const body = await Promise.all( + part.upload.map((s) => ("bytes" in s ? s.bytes : getRange(s3, s.uri, s.start, s.end))), + ); + const result = await s3.send( + new UploadPartCommand({ ...target, PartNumber: partNumber, Body: Buffer.concat(body) }), + ); + 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: UploadTarget) { + 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; + }), + ); + const failure = settled.find( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + if (failure) throw failure.reason; + return settled.map((result) => (result as PromiseFulfilledResult).value); +} + +/** Writes `parts` in order to `destUri`; a failure aborts the upload so no parts linger. */ +export async function uploadParts( + s3: S3Client, + parts: readonly Part[], + destUri: string, + contentType: string, +): Promise { + const { bucket, key } = parseS3Uri(destUri); + const created = await s3.send( + new CreateMultipartUploadCommand({ Bucket: bucket, Key: key, ContentType: contentType }), + ); + if (!created.UploadId) throw new Error(`[s3Stitch] no upload id for ${destUri}`); + const target = { Bucket: bucket, Key: key, UploadId: created.UploadId }; + try { + const etags = await sendParts(s3, parts, target); + await s3.send( + new CompleteMultipartUploadCommand({ + ...target, + MultipartUpload: { Parts: etags.map((ETag, index) => ({ ETag, PartNumber: index + 1 })) }, + }), + ); + } catch (error) { + await s3.send(new AbortMultipartUploadCommand(target)).catch((abortError: unknown) => { + console.error( + `[s3Stitch] abort failed for ${destUri}; parts remain until a lifecycle rule removes them`, + abortError, + ); + }); + throw error; + } +} + +export interface StitchResult { + fileSize: number; + parts: number; + downloadedBytes: number; +} + +/** One faststart mp4 at `destUri` from the stitched mp4 chunks at `chunkUris`, in order. */ +export async function stitchMp4InS3( + s3: S3Client, + chunkUris: readonly string[], + destUri: string, + minPartBytes = MIN_PART_BYTES, +): Promise { + const layouts = await Promise.all( + chunkUris.map((uri) => readMp4Layout((start, end) => getRange(s3, uri, start, end))), + ); + const segments: Segment[] = [ + { bytes: buildStitchedHeader(layouts) }, + ...layouts.map((layout, i) => ({ + uri: chunkUris[i]!, + start: layout.payloadStart, + end: layout.payloadEnd, + })), + ]; + const parts = planParts(segments, minPartBytes); + await uploadParts(s3, parts, destUri, "video/mp4"); + 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), + }; +} diff --git a/packages/engine/src/services/frameCapture-firstBeginFrame.test.ts b/packages/engine/src/services/frameCapture-firstBeginFrame.test.ts new file mode 100644 index 0000000000..5d9d442f2b --- /dev/null +++ b/packages/engine/src/services/frameCapture-firstBeginFrame.test.ts @@ -0,0 +1,92 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { captureFrameToBuffer, createCaptureSession } from "./frameCapture.js"; +import { beginFrameCapture } from "./screenshotService.js"; + +const calls: string[] = []; + +vi.mock("./browserManager.js", async (importOriginal) => ({ + ...(await importOriginal()), + resolveBrowserGpuMode: async () => "software", + acquireBrowser: async () => ({ + captureMode: "beginframe", + browser: { + version: async () => "HeadlessChrome/150.0.0.0", + newPage: async () => ({ + mainFrame: () => ({}), + evaluateOnNewDocument: async () => {}, + setViewport: async () => {}, + evaluate: async () => undefined, + createCDPSession: async () => ({ + send: async ( + method: string, + params: { frameTimeTicks?: number; noDisplayUpdates?: boolean }, + ) => { + calls.push( + `${method} ticks=${params.frameTimeTicks} noDisplayUpdates=${params.noDisplayUpdates}`, + ); + return {}; + }, + }), + }), + }, + }), +})); +vi.mock("./screenshotService.js", async (importOriginal) => ({ + ...(await importOriginal()), + beginFrameCapture: vi.fn(async (_page: unknown, _options: unknown, frameTimeTicks: number) => { + calls.push(`capture ticks=${frameTimeTicks}`); + return { buffer: Buffer.alloc(8), hasDamage: true }; + }), +})); + +const dirs: string[] = []; +afterEach(() => { + calls.length = 0; + vi.mocked(beginFrameCapture).mockClear(); + for (const dir of dirs.splice(0)) rmSync(dir, { recursive: true, force: true }); +}); + +async function beginFrameSession(baseTicks: number) { + const dir = mkdtempSync(join(tmpdir(), "hf-first-beginframe-")); + dirs.push(dir); + const session = await createCaptureSession( + "http://localhost:0", + dir, + { width: 1920, height: 1080, fps: { num: 30, den: 1 }, format: "jpeg" }, + null, + {}, + ); + session.isInitialized = true; + session.captureMode = "beginframe"; + session.beginFrameTimeTicks = baseTicks; + return session; +} + +describe("first BeginFrame capture of a session", () => { + // Measured on Linux headless-shell: without this, a chunk's first frame came back as the + // frame composited during init (the composition at t=0) in about half the runs. + it("displays the seek once before the first capture, between the probe tick and the capture tick", async () => { + const session = await beginFrameSession(2000); + const interval = session.beginFrameIntervalMs; + + await captureFrameToBuffer(session, 0, 70 / 30); + await captureFrameToBuffer(session, 1, 71 / 30); + + expect(calls).toEqual([ + `HeadlessExperimental.beginFrame ticks=${2000 - interval / 2} noDisplayUpdates=false`, + "capture ticks=2000", + `capture ticks=${2000 + interval}`, + ]); + // The producer's liveness probe runs at base - 5 intervals; ticks must stay monotonic. + expect(2000 - interval / 2).toBeGreaterThan(2000 - 5 * interval); + }); + + it("does it again on a fresh session (worker or retry page)", async () => { + await captureFrameToBuffer(await beginFrameSession(2000), 0, 0); + await captureFrameToBuffer(await beginFrameSession(2000), 0, 0); + expect(calls.filter((c) => c.startsWith("HeadlessExperimental.beginFrame"))).toHaveLength(2); + }); +}); diff --git a/packages/engine/src/services/frameCapture.ts b/packages/engine/src/services/frameCapture.ts index 99b0aef478..36eadfef87 100644 --- a/packages/engine/src/services/frameCapture.ts +++ b/packages/engine/src/services/frameCapture.ts @@ -236,6 +236,7 @@ export interface CaptureSession { beginFrameIntervalMs: number; beginFrameHasDamageCount: number; beginFrameNoDamageCount: number; + beginFrameSeekPresented?: boolean; /** Optional producer config — when set, overrides module-level env var constants. */ config?: Partial; /** True if running on SwiftShader (detected at init). Undefined before init. */ @@ -4025,6 +4026,17 @@ async function captureFrameSurface( if (session.captureMode === "beginframe") { const frameTimeTicks = session.beginFrameTimeTicks + frameIndex * session.beginFrameIntervalMs; + if (!session.beginFrameSeekPresented) { + // A session's first capture can return the frame composited during init instead of + // this seek, so display the seek once at a tick between the probe and this capture. + const client = await getCdpSession(page); + await client.send("HeadlessExperimental.beginFrame", { + frameTimeTicks: frameTimeTicks - session.beginFrameIntervalMs / 2, + interval: session.beginFrameIntervalMs, + noDisplayUpdates: false, + }); + session.beginFrameSeekPresented = true; + } const result = await beginFrameCapture( page, options, diff --git a/packages/producer/src/distributed.ts b/packages/producer/src/distributed.ts index ec2fbba1d9..2abb295a93 100644 --- a/packages/producer/src/distributed.ts +++ b/packages/producer/src/distributed.ts @@ -105,7 +105,18 @@ export { } from "./services/distributed/renderChunk.js"; // ── Assemble (Activity C) ─────────────────────────────────────────────────── -export { assemble, type AssembleResult } from "./services/distributed/assemble.js"; +export { + assemble, + chunksAreStitchedMp4, + type AssembleResult, +} from "./services/distributed/assemble.js"; +export { muxChunkWithAudio } from "./services/distributed/chunkAudio.js"; +export { + buildStitchedHeader, + readMp4Layout, + type Mp4Layout, + type ReadRange, +} from "./services/distributed/mp4Stitch.js"; // ── Cloud-agnostic adapter helpers ────────────────────────────────────────── // Shared by the distributed-render adapters (aws-lambda, gcp-cloud-run, …) so diff --git a/packages/producer/src/regression-harness-lambda-local.ts b/packages/producer/src/regression-harness-lambda-local.ts index 88a2da4127..84dd83e553 100644 --- a/packages/producer/src/regression-harness-lambda-local.ts +++ b/packages/producer/src/regression-harness-lambda-local.ts @@ -23,6 +23,7 @@ */ import { + type Dirent, createWriteStream, existsSync, mkdirSync, @@ -274,14 +275,20 @@ function requireV2PlanResult( // branches would make this small test-harness utility harder to audit. // fallow-ignore-next-line complexity function measureDirectoryBytes(path: string): number { - if (!existsSync(path)) return 0; + let entries: Dirent[]; + try { + entries = readdirSync(path, { withFileTypes: true }); + } catch (err) { + if ((err as NodeJS.ErrnoException).code === "ENOENT") return 0; + throw err; + } let bytes = 0; - for (const entry of readdirSync(path, { withFileTypes: true })) { + for (const entry of entries) { const child = join(path, entry.name); if (entry.isDirectory()) { bytes += measureDirectoryBytes(child); } else if (entry.isFile()) { - bytes += statSync(child).size; + bytes += statSync(child, { throwIfNoEntry: false })?.size ?? 0; } } return bytes; diff --git a/packages/producer/src/services/distributed/assemble.ts b/packages/producer/src/services/distributed/assemble.ts index 09431e404d..1b44d8192f 100644 --- a/packages/producer/src/services/distributed/assemble.ts +++ b/packages/producer/src/services/distributed/assemble.ts @@ -52,6 +52,7 @@ import { import { encoderFailureError } from "../render/encoderInterruption.js"; import type { ChunkSliceJson, LockedRenderConfig } from "../render/stages/freezePlan.js"; import { DISTRIBUTED_RENDER_CAPABILITIES, readPlanProtocolV1 } from "./planProtocol.js"; +import { writeStitchedMp4 } from "./mp4Stitch.js"; import { validatePlanV2MaterializedTarget } from "./planV2.js"; import type { DistributedFormat } from "./shared.js"; @@ -141,6 +142,16 @@ export async function assemble( return mergePngFrameDirs(chunkPaths, outputPath, plan.totalFrames, audioPath, start); } + if (chunksAreStitchedMp4(planDir, cfr)) { + await writeStitchedMp4(chunkPaths, outputPath); + return { + outputPath, + durationMs: Date.now() - start, + framesEncoded: plan.totalFrames, + fileSize: statSync(outputPath).size, + }; + } + // ── 2b. mp4 / mov / webm: join + audio mux + faststart in one pass ───── const format = plan.dimensions.format; const cfrEncoder = cfr ? readCfrEncoder(planDir, format) : null; @@ -155,9 +166,6 @@ export async function assemble( format, fps, fpsArg: fpsToFfmpegArg(fps), - chunkSeconds: chunksCarryAudio(planDir) - ? plan.chunks.map((chunk) => ((chunk.endFrame - chunk.startFrame) * fps.den) / fps.num) - : null, signal: options?.abortSignal, timeout: resolveConfig().ffmpegProcessTimeout, log, @@ -200,10 +208,7 @@ export async function assemble( }; } -function readAssemblePlan( - planDir: string, - chunkPaths: readonly string[], -): PlanJsonForAssemble & { chunks: ChunkSliceJson[] } { +function readAssemblePlan(planDir: string, chunkPaths: readonly string[]): PlanJsonForAssemble { // ── 1. Validate planDir manifest matches chunkPaths shape ────────────── const planJsonPath = join(planDir, "plan.json"); const chunksJsonPath = join(planDir, "meta", "chunks.json"); @@ -230,7 +235,7 @@ function readAssemblePlan( } } - return { ...plan, chunks }; + return plan; } interface JoinContext { @@ -242,7 +247,6 @@ interface JoinContext { timeout: number; log: ProducerLogger; timed: (stage: string, run: () => Promise) => Promise; - chunkSeconds: number[] | null; } // A one-entry concat list ignores the input-side `-r` (observed `359/12` in the @@ -250,7 +254,7 @@ interface JoinContext { function videoInputArgs(ctx: JoinContext, videoSources: readonly string[]): string[] { return videoSources.length === 1 ? ["-i", videoSources[0]!] - : concatDemuxerInput(videoSources, ctx.workDir, ctx.fpsArg, ctx.chunkSeconds); + : concatDemuxerInput(videoSources, ctx.workDir, ctx.fpsArg); } // ── 2c. Optional exact-CFR re-encode ────────────────────────────────── @@ -347,18 +351,11 @@ function concatDemuxerInput( chunkPaths: readonly string[], workDir: string, fpsArg: string, - chunkSeconds: readonly number[] | null, ): string[] { // ffmpeg's concat list escapes a single quote as `'\''`. const listPath = join(workDir, "concat-list.txt"); - // Chunks that carry audio: read only their video (track 1), and end each chunk at its last frame - // rather than after its audio, which starts with the AAC priming and can run one packet longer. - const head = chunkSeconds ? ["stream", "exact_stream_id 1"] : []; - const body = chunkPaths.map((path, i) => { - const line = `file '${path.replace(/'/g, "'\\''")}'`; - return chunkSeconds ? `${line}\nduration ${chunkSeconds[i]!.toFixed(6)}` : line; - }); - writeFileSync(listPath, `${[...head, ...body].join("\n")}\n`, "utf-8"); + const body = chunkPaths.map((path) => `file '${path.replace(/'/g, "'\\''")}'`).join("\n"); + writeFileSync(listPath, `${body}\n`, "utf-8"); return ["-r", fpsArg, "-f", "concat", "-safe", "0", "-i", listPath]; } @@ -393,6 +390,15 @@ async function runAssembleFfmpeg( } } +/** Chunks that carry their audio are stitched under one index, the same bytes as the S3 stitch. */ +export function chunksAreStitchedMp4(planDir: string, cfr: boolean): boolean { + if (!chunksCarryAudio(planDir)) return false; + if (cfr) { + throw new Error("[assemble] cfr=true cannot re-encode a plan whose mp4 chunks carry audio"); + } + return true; +} + function chunksCarryAudio(planDir: string): boolean { const encoderJsonPath = join(planDir, "meta", "encoder.json"); if (!existsSync(encoderJsonPath)) return false; diff --git a/packages/producer/src/services/distributed/chunkAudio.test.ts b/packages/producer/src/services/distributed/chunkAudio.test.ts index 832e3a0906..ac263f0a8b 100644 --- a/packages/producer/src/services/distributed/chunkAudio.test.ts +++ b/packages/producer/src/services/distributed/chunkAudio.test.ts @@ -38,7 +38,6 @@ describe("audioCutForChunk", () => { const hasFfmpeg = spawnSync("ffmpeg", ["-version"]).status === 0; describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { - const TOTAL_FRAMES = 91; let root: string; beforeAll(() => { @@ -74,13 +73,18 @@ describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { ]), ).streams as { codec_type: string; nb_read_packets: string; duration: string }[]; - // Uneven chunks with a 1-frame tail, and a single chunk (read without a concat list today). + // Uneven chunks with a 1-frame tail, and a single chunk (read without a concat list today). At + // 60 fps a 1-frame tail is shorter than an AAC packet: 51 frames are 40800 samples, 41 packets + // after the priming, and packet 41 would be the tail's first, so the tail owns none. it.each([ - ["three", [0, 47, 90, 91]], - ["one", [0, 91]], + ["three", 30, [0, 47, 90, 91], true], + ["one", 30, [0, 91], true], + ["60fps-silent-tail", 60, [0, 50, 51], false], ] as const)( - "%s chunk(s) join through today's assemble like video-only chunks", - async (name, edges) => { + "%s chunk(s) at %i fps assemble into the frames and audio of today's video-only join", + async (name, rate, edges, tailHasAudio) => { + const fps = { num: rate, den: 1 }; + const totalFrames = edges.at(-1)!; const CHUNKS = edges.slice(1).map((end, index) => ({ index, startFrame: edges[index]!, @@ -90,7 +94,7 @@ describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { mkdirSync(dir, { recursive: true }); const audioPath = join(dir, "audio.m4a"); ffmpeg( - `-f lavfi -i aevalsrc=0.4*sin(2*PI*(220+40*t)*t):s=48000:d=${TOTAL_FRAMES / 30} -c:a aac -b:a 192k -y` + `-f lavfi -i aevalsrc=0.4*sin(2*PI*(220+40*t)*t):s=48000:d=${totalFrames / rate} -c:a aac -b:a 192k -y` .split(" ") .concat(audioPath), ); @@ -100,9 +104,9 @@ describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { join(planDir, "plan.json"), JSON.stringify({ planHash: "fake", - totalFrames: TOTAL_FRAMES, + totalFrames, hasAudio: true, - dimensions: { fpsNum: 30, fpsDen: 1, width: 160, height: 120, format: "mp4" }, + dimensions: { fpsNum: rate, fpsDen: 1, width: 160, height: 120, format: "mp4" }, }), ); writeFileSync(join(planDir, "meta", "chunks.json"), JSON.stringify(CHUNKS)); @@ -115,7 +119,7 @@ describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { chunkPaths.push(join(dir, `chunk-${slice.index}.mp4`)); // Closed GOP, no B-frames, headers repeated: shaped like renderChunk's encode. ffmpeg( - `-f lavfi -i testsrc2=size=160x120:rate=30,trim=start_frame=${slice.startFrame}:end_frame=${slice.endFrame} -c:v libx264 -preset veryfast -g ${n} -keyint_min ${n} -sc_threshold 0 -bf 0 -x264-params scenecut=0:open-gop=0:repeat-headers=1 -pix_fmt yuv420p -avoid_negative_ts make_zero -y` + `-f lavfi -i testsrc2=size=160x120:rate=${rate},trim=start_frame=${slice.startFrame}:end_frame=${slice.endFrame} -c:v libx264 -preset veryfast -g ${n} -keyint_min ${n} -sc_threshold 0 -bf 0 -x264-params scenecut=0:open-gop=0:repeat-headers=1 -pix_fmt yuv420p -avoid_negative_ts make_zero -y` .split(" ") .concat(videoPaths[slice.index]!), ); @@ -125,14 +129,16 @@ describe.skipIf(!hasFfmpeg)("chunks that carry their audio", () => { outputPath: chunkPaths[slice.index]!, chunks: CHUNKS, chunkIndex: slice.index, - fps: FPS, + fps, timeoutMs: 60_000, }); } - // Every plan audio packet lands in exactly one chunk. - const audioPackets = Number(streams(audioPath)[0]!.nb_read_packets); - const chunkPackets = chunkPaths.map((path) => Number(streams(path)[1]!.nb_read_packets)); + // Every plan audio packet lands in exactly one chunk; a tail that owns none has no audio track. + const audioOf = (path: string) => streams(path).find((s) => s.codec_type === "audio"); + const audioPackets = Number(audioOf(audioPath)!.nb_read_packets); + const chunkPackets = chunkPaths.map((path) => Number(audioOf(path)?.nb_read_packets ?? 0)); expect(chunkPackets.reduce((a, b) => a + b, 0)).toBe(audioPackets); + expect(audioOf(chunkPaths.at(-1)!) !== undefined).toBe(tailHasAudio); const withAudio = join(dir, "with-audio.mp4"); const videoOnly = join(dir, "video-only.mp4"); diff --git a/packages/producer/src/services/distributed/mp4Stitch.test.ts b/packages/producer/src/services/distributed/mp4Stitch.test.ts new file mode 100644 index 0000000000..0f58da9ffa --- /dev/null +++ b/packages/producer/src/services/distributed/mp4Stitch.test.ts @@ -0,0 +1,271 @@ +import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import { mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { buildStitchedHeader, readMp4Layout, readTracks, writeStitchedMp4 } from "./mp4Stitch.js"; + +function box(type: string, ...body: Buffer[]): Buffer { + const header = Buffer.alloc(8); + header.writeUInt32BE(8 + body.reduce((n, part) => n + part.length, 0)); + header.write(type, 4, "latin1"); + return Buffer.concat([header, ...body]); +} + +function u32(...values: number[]): Buffer { + const buf = Buffer.alloc(values.length * 4); + values.forEach((value, i) => buf.writeUInt32BE(value, i * 4)); + return buf; +} +const full = (type: string, ...body: Buffer[]) => box(type, u32(0), ...body); +// Version 0 mvhd/tkhd/mdhd with only the fields the stitcher reads filled in. +const header = (type: string, timescaleOrTrack: number, duration: number) => + type === "tkhd" + ? full(type, u32(0, 0, timescaleOrTrack, 0, duration), Buffer.alloc(60)) + : full(type, u32(0, 0, timescaleOrTrack, duration), Buffer.alloc(type === "mvhd" ? 80 : 4)); + +interface FakeTrack { + handler: "vide" | "soun"; + timescale: number; + mediaTime: number; + durations: number[]; + sizes: number[]; + sync?: number[]; + codec?: string; +} + +/** A chunk file shaped like ffmpeg's: ftyp, mdat with samples interleaved one per mp4 chunk, moov + * (or the moov first, as a faststart file has it). */ +function fakeChunk( + tag: number, + tracks: FakeTrack[], + moovFirst = false, +): { file: Buffer; samples: Buffer[][] } { + const samples = tracks.map((track, t) => + track.sizes.map((size, k) => Buffer.alloc(size, (tag * 61 + t * 17 + k * 5 + 1) & 255)), + ); + const ftyp = box("ftyp", Buffer.from("isom\0\0\x02\0isom", "latin1")); + const moovAt = (payloadAt: number) => { + const offsets = tracks.map(() => [] as number[]); + const payload: Buffer[] = []; + let at = payloadAt; + for (let k = 0; k < Math.max(...tracks.map((track) => track.sizes.length)); k++) { + samples.forEach((track, t) => { + if (!track[k]) return; + offsets[t]!.push(at); + payload.push(track[k]!); + at += track[k]!.length; + }); + } + const traks = tracks.map((track, t) => { + const duration = track.durations.reduce((n, d) => n + d, 0); + const entry = track.codec ?? (track.handler === "vide" ? "avc1" : "mp4a"); + const stbl = box( + "stbl", + full("stsd", u32(1), box(entry, Buffer.alloc(8, t))), + full("stts", u32(track.durations.length, ...track.durations.flatMap((d) => [1, d]))), + ...(track.sync ? [full("stss", u32(track.sync.length, ...track.sync))] : []), + full("stsc", u32(1, 1, 1, 1)), + full("stsz", u32(0, track.sizes.length, ...track.sizes)), + full("stco", u32(offsets[t]!.length, ...offsets[t]!)), + ); + return box( + "trak", + header("tkhd", t + 1, 0), + box("edts", full("elst", u32(1, 0, track.mediaTime, 0x10000))), + box( + "mdia", + header("mdhd", track.timescale, duration), + full("hdlr", u32(0), Buffer.from(track.handler), Buffer.alloc(13)), + box("minf", stbl), + ), + ); + }); + return { moov: box("moov", header("mvhd", 1000, 0), ...traks), mdat: box("mdat", ...payload) }; + }; + if (!moovFirst) { + const { moov, mdat } = moovAt(ftyp.length + 8); + return { file: Buffer.concat([ftyp, mdat, moov]), samples }; + } + // The moov's size does not depend on the offsets it holds. + const { moov, mdat } = moovAt(ftyp.length + moovAt(0).moov.length + 8); + return { file: Buffer.concat([ftyp, moov, mdat]), samples }; +} + +const memoryRange = (file: Buffer) => async (start: number, end: number) => + file.subarray(start, end); + +/** Every sample of `track`, read from `file` where its tables place it. */ +function samplesOf(file: Buffer, track: ReturnType[number]): Buffer[] { + const read: Buffer[] = []; + let k = 0; + track.chunkOffsets.forEach((offset, c) => { + for (let s = 0, at = offset; s < track.samplesPerChunk[c]!; s++) { + read.push(file.subarray(at, at + track.sizes[k]!)); + at += track.sizes[k++]!; + } + }); + return read; +} + +/** The stitched file: the header, then each chunk's mdat payload in order. */ +async function stitchInMemory(files: Buffer[]): Promise { + const layouts = await Promise.all(files.map((file) => readMp4Layout(memoryRange(file)))); + return Buffer.concat([ + buildStitchedHeader(layouts), + ...layouts.map((l, i) => files[i]!.subarray(l.payloadStart, l.payloadEnd)), + ]); +} + +describe("buildStitchedHeader", () => { + const video = (sizes: number[], mediaTime = 0): FakeTrack => ({ + handler: "vide", + timescale: 15360, + mediaTime, + durations: sizes.map(() => 512), + sizes, + sync: [1], + }); + const audio = (durations: number[], mediaTime: number): FakeTrack => ({ + handler: "soun", + timescale: 48000, + mediaTime, + durations, + sizes: durations.map((_, k) => 40 + k), + }); + + it("indexes every chunk's samples at their place after the header", async () => { + const chunks = [ + fakeChunk(0, [video([500, 120, 130]), audio([1024, 1024, 1024, 1024], 1024)]), + fakeChunk(1, [video([480, 110]), audio([1024, 1024, 640], 0)]), + ]; + const layouts = await Promise.all(chunks.map((c) => readMp4Layout(memoryRange(c.file)))); + const header = buildStitchedHeader(layouts); + const file = Buffer.concat([ + header, + ...layouts.map((l, i) => chunks[i]!.file.subarray(l.payloadStart, l.payloadEnd)), + ]); + + const stitched = await readMp4Layout(memoryRange(file)); + expect(stitched.payloadStart).toBe(header.length); + expect(stitched.payloadEnd).toBe(file.length); + const tracks = readTracks(stitched.moov); + tracks.forEach((track, t) => { + expect(samplesOf(file, track)).toEqual(chunks.flatMap((chunk) => chunk.samples[t]!)); + }); + expect(tracks.map((track) => [track.handler, track.mediaTime])).toEqual([ + ["vide", 0], + ["soun", 1024], + ]); + expect(tracks[0]!.sync).toEqual([true, false, false, true, false]); + expect(tracks[1]!.durations).toEqual([1024, 1024, 1024, 1024, 1024, 1024, 640]); + + const moov = stitched.moov; + const after = (type: string, from = 0) => moov.indexOf(type, from, "latin1") + 4; + // Only audio needs an edit list (version 1): it shows its media minus the 1024-sample + // priming, (6784 - 1024) / 48000 s = 120 ms in the 1000-tick movie timescale. + const audioElst = after("elst"); + expect(moov.indexOf("elst", audioElst, "latin1")).toBe(-1); + expect(audioElst).toBeGreaterThan(after("vide")); + expect(Number(moov.readBigUInt64BE(audioElst + 8))).toBe(120); + expect(Number(moov.readBigInt64BE(audioElst + 16))).toBe(1024); + expect(moov.readUInt32BE(after("mdhd", audioElst) + 16)).toBe(6784); + expect(moov.readUInt32BE(after("mvhd") + 16)).toBe(167); + }); + + it("refuses a chunk whose video starts late, as B-frames make it", async () => { + const chunks = [fakeChunk(0, [video([500, 120])]), fakeChunk(1, [video([480, 110], 512)])]; + const layouts = await Promise.all(chunks.map((c) => readMp4Layout(memoryRange(c.file)))); + expect(() => buildStitchedHeader(layouts)).toThrow("chunk 1 video starts at media time 512"); + }); + + // A tail shorter than one AAC packet owns no packet (the one before it spans the boundary), so + // its chunk is muxed video-only; the plan's audio track still holds every packet once. + it("keeps the plan's audio when a chunk owns no audio packets", async () => { + const chunks = [ + fakeChunk(0, [video([500, 120, 130]), audio([1024, 1024, 1024, 640], 1024)]), + fakeChunk(1, [video([90])]), + ]; + const file = await stitchInMemory(chunks.map((chunk) => chunk.file)); + const tracks = readTracks((await readMp4Layout(memoryRange(file))).moov); + expect(tracks.map((track) => [track.handler, track.mediaTime])).toEqual([ + ["vide", 0], + ["soun", 1024], + ]); + expect(samplesOf(file, tracks[0]!)).toEqual([ + ...chunks[0]!.samples[0]!, + ...chunks[1]!.samples[0]!, + ]); + expect(samplesOf(file, tracks[1]!)).toEqual(chunks[0]!.samples[1]!); + expect(tracks[1]!.durations).toEqual([1024, 1024, 1024, 640]); + }); + + it.each([ + [ + "an audio codec chunk 0 does not use", + [video([500]), audio([1024], 1024)], + [video([90]), { ...audio([1024], 0), codec: "Opus" }], + ], + ["no video", [video([500]), audio([1024], 1024)], [audio([1024], 0)]], + ["audio chunk 0 does not carry", [video([500])], [video([90]), audio([1024], 0)]], + ] as const)("refuses a chunk with %s", async (_, first, second) => { + const layouts = await Promise.all( + [fakeChunk(0, [...first]), fakeChunk(1, [...second])].map((c) => + readMp4Layout(memoryRange(c.file)), + ), + ); + expect(() => buildStitchedHeader(layouts)).toThrow("chunk 1 tracks differ from chunk 0"); + }); +}); + +describe("writeStitchedMp4", () => { + const video = (sizes: number[], mediaTime = 0): FakeTrack => ({ + handler: "vide", + timescale: 15360, + mediaTime, + durations: sizes.map(() => 512), + sizes, + sync: [1], + }); + let dir: string; + beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), "hf-mp4-stitch-")); + }); + afterEach(() => { + rmSync(dir, { recursive: true, force: true }); + }); + const write = (name: string, file: Buffer) => { + writeFileSync(join(dir, name), file); + return join(dir, name); + }; + + it("writes the header and every payload, replacing an earlier output", async () => { + const files = [fakeChunk(0, [video([500, 120])]).file, fakeChunk(1, [video([480])]).file]; + const output = write("out.mp4", Buffer.from("stale")); + await writeStitchedMp4([write("a.mp4", files[0]!), write("b.mp4", files[1]!)], output); + expect(readFileSync(output).equals(await stitchInMemory(files))).toBe(true); + expect(readdirSync(dir).sort()).toEqual(["a.mp4", "b.mp4", "out.mp4"]); + }); + + it("opens no output for chunks it refuses", async () => { + const paths = [ + write("a.mp4", fakeChunk(0, [video([500])]).file), + write("b.mp4", fakeChunk(1, [video([480], 512)]).file), + ]; + await expect(writeStitchedMp4(paths, join(dir, "out.mp4"))).rejects.toThrow( + "chunk 1 video starts at media time 512", + ); + expect(readdirSync(dir).sort()).toEqual(["a.mp4", "b.mp4"]); + }); + + it("fails a chunk whose mdat is cut short mid-copy and leaves no output", async () => { + const whole = fakeChunk(1, [video([480, 300])], true).file; + const paths = [ + write("a.mp4", fakeChunk(0, [video([500])]).file), + write("b.mp4", whole.subarray(0, whole.length - 200)), + ]; + await expect(writeStitchedMp4(paths, join(dir, "out.mp4"))).rejects.toThrow( + "chunk 1 mdat holds 580 of 780 payload bytes", + ); + expect(readdirSync(dir).sort()).toEqual(["a.mp4", "b.mp4"]); + }); +}); diff --git a/packages/producer/src/services/distributed/mp4Stitch.ts b/packages/producer/src/services/distributed/mp4Stitch.ts new file mode 100644 index 0000000000..e84533b61d --- /dev/null +++ b/packages/producer/src/services/distributed/mp4Stitch.ts @@ -0,0 +1,459 @@ +// Assemble for chunks that carry their audio: one faststart header (ftyp, a moov indexing every +// chunk's samples, a 64-bit mdat header); the chunks' mdat payloads follow it unchanged and in +// order, so storage can copy them into place. + +import { + createReadStream, + createWriteStream, + mkdirSync, + mkdtempSync, + renameSync, + rmSync, +} from "node:fs"; +import { open } from "node:fs/promises"; +import { dirname, join } from "node:path"; +import type { Writable } from "node:stream"; +import { finished, pipeline } from "node:stream/promises"; + +// ── Box reading ───────────────────────────────────────────────────────────── + +interface Box { + type: string; + start: number; + bodyStart: number; + end: number; +} + +/** The box whose header sits at `off` in `src`, which holds the bytes at file offset `at`. */ +function boxAt(src: Buffer, off: number, at: number, limit: number): Box { + let size = src.readUInt32BE(off); + let bodyStart = at + 8; + if (size === 1) { + size = Number(src.readBigUInt64BE(off + 8)); + bodyStart = at + 16; + } + if (size < 8 || at + size > limit) throw new Error(`[mp4Stitch] malformed box at ${at}`); + return { type: src.toString("latin1", off + 4, off + 8), start: at, bodyStart, end: at + size }; +} + +function childBoxes(buf: Buffer, parent: Box): Box[] { + const boxes: Box[] = []; + for (let at = parent.bodyStart; at < parent.end; at = boxes.at(-1)!.end) { + boxes.push(boxAt(buf, at, at, parent.end)); + } + return boxes; +} + +function findChild(buf: Buffer, parent: Box, type: string): Box | undefined { + return childBoxes(buf, parent).find((box) => box.type === type); +} + +function child(buf: Buffer, parent: Box, type: string): Box { + const found = findChild(buf, parent, type); + if (!found) throw new Error(`[mp4Stitch] ${parent.type} has no ${type}`); + return found; +} + +/** Reads bytes `[start, end)`; returns fewer past the end of the object. */ +export type ReadRange = (start: number, end: number) => Promise; + +export interface Mp4Layout { + ftyp: Buffer; + moov: Buffer; + /** File offsets of the mdat payload, `[payloadStart, payloadEnd)`. */ + payloadStart: number; + payloadEnd: number; +} + +const SCAN_BYTES = 64 * 1024; + +/** Reads a chunk's ftyp and moov and locates its mdat payload without reading the payload. */ +export async function readMp4Layout(read: ReadRange): Promise { + const found = new Map(); + let payload: Box | undefined; + let buf: Buffer = Buffer.alloc(0); + let bufAt = 0; + // Holds `[start, end)` in `buf` when the object has those bytes. + const load = async (start: number, end: number) => { + if (start >= bufAt && end <= bufAt + buf.length) return; + buf = await read(start, Math.max(end, start + SCAN_BYTES)); + bufAt = start; + }; + for (let at = 0; !(found.has("ftyp") && found.has("moov") && payload); ) { + await load(at, at + 16); + if (at + 8 > bufAt + buf.length) { + throw new Error(`[mp4Stitch] chunk needs ftyp, moov and mdat; found ${[...found.keys()]}`); + } + const box = boxAt(buf, at - bufAt, at, Number.MAX_SAFE_INTEGER); + // Samples in a later mdat fail buildStitchedHeader's offset check. + if (box.type === "mdat") payload ??= box; + if (box.type === "ftyp" || box.type === "moov") { + await load(at, box.end); + found.set(box.type, buf.subarray(at - bufAt, box.end - bufAt)); + } + at = box.end; + } + const [ftyp, moov] = [found.get("ftyp")!, found.get("moov")!]; + if (moov.length !== moovSize(moov)) throw new Error("[mp4Stitch] moov is cut off"); + return { ftyp, moov, payloadStart: payload.bodyStart, payloadEnd: payload.end }; +} + +const moovSize = (moov: Buffer) => boxAt(moov, 0, 0, Number.MAX_SAFE_INTEGER).end; + +export interface TrackTable { + trak: Box; + handler: string; + timescale: number; + /** Media time the edit list starts presenting at (the AAC priming on audio). */ + mediaTime: number; + stsd: Buffer; + durations: number[]; + sizes: number[]; + /** Per-sample keyframe flags; `null` when every sample is one (no stss). */ + sync: boolean[] | null; + samplesPerChunk: number[]; + chunkOffsets: number[]; +} + +/** Each track's sample tables, expanded to one entry per sample (and per chunk). */ +export function readTracks(moov: Buffer): TrackTable[] { + const root = boxAt(moov, 0, 0, moov.length); + return childBoxes(moov, root) + .filter((box) => box.type === "trak") + .map((trak) => readTrack(moov, trak)); +} + +/** `[count, rows...]` tables: returns `columns` u32 values per row. */ +function rows(buf: Buffer, table: Box, columns: number, countAt = 4): number[][] { + const count = buf.readUInt32BE(table.bodyStart + countAt); + const first = table.bodyStart + countAt + 4; + return Array.from({ length: count }, (_, row) => + Array.from({ length: columns }, (_, col) => + buf.readUInt32BE(first + (row * columns + col) * 4), + ), + ); +} + +function readMediaTime(moov: Buffer, trak: Box): number { + const edts = findChild(moov, trak, "edts"); + if (!edts) return 0; + const elst = child(moov, edts, "elst"); + const wide = moov[elst.bodyStart] === 1; + const entries = moov.readUInt32BE(elst.bodyStart + 4); + for (let i = 0, at = elst.bodyStart + 8; i < entries; i++, at += wide ? 20 : 12) { + const offset = at + (wide ? 8 : 4); + const mediaTime = wide ? Number(moov.readBigInt64BE(offset)) : moov.readInt32BE(offset); + if (mediaTime !== -1) return mediaTime; + } + return 0; +} + +function readSizes(moov: Buffer, stsz: Box): number[] { + const fixedSize = moov.readUInt32BE(stsz.bodyStart + 4); + if (!fixedSize) return rows(moov, stsz, 1, 8).map(([size]) => size!); + return Array(moov.readUInt32BE(stsz.bodyStart + 8)).fill(fixedSize); +} + +function readSync(moov: Buffer, stbl: Box, sampleCount: number): boolean[] | null { + const stss = findChild(moov, stbl, "stss"); + if (!stss) return null; + const sync = Array(sampleCount).fill(false); + for (const [n] of rows(moov, stss, 1)) sync[n! - 1] = true; + return sync; +} + +function readChunkOffsets(moov: Buffer, stbl: Box): number[] { + const stco = findChild(moov, stbl, "stco"); + if (stco) return rows(moov, stco, 1).map(([offset]) => offset!); + return rows(moov, child(moov, stbl, "co64"), 2).map(([hi, lo]) => hi! * 2 ** 32 + lo!); +} + +function readSamplesPerChunk(moov: Buffer, stbl: Box, chunkCount: number): number[] { + const stsc = rows(moov, child(moov, stbl, "stsc"), 3); + return stsc.flatMap(([firstChunk, perChunk, description], i) => { + if (description !== 1) throw new Error("[mp4Stitch] chunk uses more than one sample entry"); + const nextFirst = stsc[i + 1]?.[0] ?? chunkCount + 1; + return Array(nextFirst - firstChunk!).fill(perChunk!); + }); +} + +function readTrack(moov: Buffer, trak: Box): TrackTable { + const mdia = child(moov, trak, "mdia"); + const hdlr = child(moov, mdia, "hdlr"); + const stbl = child(moov, child(moov, mdia, "minf"), "stbl"); + // Chunks encode without B-frames; composition offsets would shift every chunk differently. + if (findChild(moov, stbl, "ctts")) throw new Error("[mp4Stitch] chunk has ctts (B-frames)"); + const stts = rows(moov, child(moov, stbl, "stts"), 2); + const durations = stts.flatMap(([n, delta]) => Array(n!).fill(delta!)); + const sizes = readSizes(moov, child(moov, stbl, "stsz")); + const chunkOffsets = readChunkOffsets(moov, stbl); + const samplesPerChunk = readSamplesPerChunk(moov, stbl, chunkOffsets.length); + if (durations.length !== sizes.length || sum(samplesPerChunk) !== sizes.length) { + throw new Error("[mp4Stitch] sample tables disagree on the sample count"); + } + const stsd = child(moov, stbl, "stsd"); + return { + trak, + handler: moov.toString("latin1", hdlr.bodyStart + 8, hdlr.bodyStart + 12), + timescale: timescaleOf(moov, child(moov, mdia, "mdhd")), + mediaTime: readMediaTime(moov, trak), + stsd: moov.subarray(stsd.start, stsd.end), + durations, + sizes, + sync: readSync(moov, stbl, sizes.length), + samplesPerChunk, + chunkOffsets, + }; +} + +function sum(values: readonly number[]): number { + return values.reduce((total, value) => total + value, 0); +} + +// ── Box writing ───────────────────────────────────────────────────────────── + +function box(type: string, ...body: Buffer[]): Buffer { + const header = Buffer.alloc(8); + header.writeUInt32BE(8 + sum(body.map((part) => part.length))); + header.write(type, 4, "latin1"); + return Buffer.concat([header, ...body]); +} + +function fullBox(type: string, version: number, body: Buffer): Buffer { + return box(type, Buffer.from([version, 0, 0, 0]), body); +} + +function u32s(values: readonly number[]): Buffer { + const buf = Buffer.alloc(values.length * 4); + values.forEach((value, i) => buf.writeUInt32BE(value, i * 4)); + return buf; +} + +/** Run-length pairs `[count, value]` of consecutive equal values. */ +function runs(values: readonly number[]): [number, number][] { + const out: [number, number][] = []; + for (const value of values) { + const last = out.at(-1); + if (last && last[1] === value) last[0]++; + else out.push([1, value]); + } + return out; +} + +/** Rebuilds `parent` from its children, replacing the ones `replace` names. */ +function rebuild(buf: Buffer, parent: Box, replace: Record Buffer>): Buffer { + const parts = childBoxes(buf, parent).map( + (box) => replace[box.type]?.(box) ?? buf.subarray(box.start, box.end), + ); + return box(parent.type, ...parts); +} + +// Duration field offset in the box body, for version 0 and 1. +const DURATION_AT: Record = { + mvhd: [16, 24], + mdhd: [16, 24], + tkhd: [20, 28], +}; + +function withDuration(buf: Buffer, target: Box, duration: number): Buffer { + const copy = Buffer.from(buf.subarray(target.start, target.end)); + const body = target.bodyStart - target.start; + const wide = copy[body] === 1; + const at = body + DURATION_AT[target.type]![wide ? 1 : 0]; + if (wide) copy.writeBigUInt64BE(BigInt(duration), at); + else if (duration <= 0xffffffff) copy.writeUInt32BE(duration, at); + else throw new Error(`[mp4Stitch] ${target.type} duration ${duration} needs a version 1 box`); + return copy; +} + +function timescaleOf(buf: Buffer, header: Box): number { + return buf.readUInt32BE(header.bodyStart + (buf[header.bodyStart] === 1 ? 20 : 12)); +} + +interface MergedTrack { + durations: number[]; + sizes: number[]; + sync: boolean[] | null; + samplesPerChunk: number[]; + chunkOffsets: number[]; +} + +function buildStbl(stsd: Buffer, track: MergedTrack): Buffer { + const stts = runs(track.durations); + const stsc: number[][] = []; + let firstChunk = 1; + for (const [count, perChunk] of runs(track.samplesPerChunk)) { + stsc.push([firstChunk, perChunk, 1]); + firstChunk += count; + } + const co64 = Buffer.alloc(4 + track.chunkOffsets.length * 8); + co64.writeUInt32BE(track.chunkOffsets.length); + track.chunkOffsets.forEach((offset, i) => co64.writeBigUInt64BE(BigInt(offset), 4 + i * 8)); + const syncSamples = track.sync?.flatMap((isSync, i) => (isSync ? [i + 1] : [])) ?? null; + return box( + "stbl", + stsd, + fullBox("stts", 0, u32s([stts.length, ...stts.flat()])), + ...(syncSamples ? [fullBox("stss", 0, u32s([syncSamples.length, ...syncSamples]))] : []), + fullBox("stsc", 0, u32s([stsc.length, ...stsc.flat()])), + fullBox("stsz", 0, u32s([0, track.sizes.length, ...track.sizes])), + fullBox("co64", 0, co64), + ); +} + +// Only chunk 0's sample entry is kept. Chunks repeat SPS/PPS in-band (`repeat-headers=1`), and +// their entries legitimately differ in per-file fields (x264 frame_num width, esds bitrates). +const codecOf = (track: TrackTable) => track.stsd.toString("latin1", 20, 24); + +const sameTrack = (track: TrackTable, plan: TrackTable): boolean => + track.handler === plan.handler && + track.timescale === plan.timescale && + codecOf(track) === codecOf(plan); + +/** Chunk `index`'s track for each of chunk 0's, in order; null where the chunk owns no audio + * packets (a chunk shorter than one AAC packet is muxed video-only). Anything else must match. */ +function planTracksOf(tracks: readonly TrackTable[], plan: readonly TrackTable[], index: number) { + let next = 0; + const matched = plan.map((planTrack) => { + const track = tracks[next]; + if (track && sameTrack(track, planTrack)) return tracks[next++]!; + if (planTrack.handler === "soun") return null; + throw new Error(`[mp4Stitch] chunk ${index} tracks differ from chunk 0`); + }); + if (next !== tracks.length) + throw new Error(`[mp4Stitch] chunk ${index} tracks differ from chunk 0`); + for (const track of tracks) { + if (track.handler === "vide" && track.mediaTime !== 0) { + throw new Error(`[mp4Stitch] chunk ${index} video starts at media time ${track.mediaTime}`); + } + } + return matched; +} + +// ftyp + moov + mdat header for the chunks' payloads placed in order after it. The moov is chunk +// 0's with merged sample tables; chunk 0's audio priming skip stays in an edit list. +export function buildStitchedHeader(layouts: readonly Mp4Layout[]): Buffer { + const plan = readTracks(layouts[0]!.moov); + const chunks = layouts.map((layout, i) => planTracksOf(readTracks(layout.moov), plan, i)); + const payloadBases: number[] = []; + let payloadBytes = 0; + for (const layout of layouts) { + payloadBases.push(payloadBytes); + payloadBytes += layout.payloadEnd - layout.payloadStart; + } + const template = layouts[0]!.moov; + const root = boxAt(template, 0, 0, template.length); + const movieTimescale = timescaleOf(template, child(template, root, "mvhd")); + + const buildMoov = (dataStart: number): Buffer => { + const traks = plan.map((track, t) => { + // The chunks that carry this track, with their place in the file. + const parts = chunks.flatMap((tracks, chunk) => { + const part = tracks[t]; + return part ? [{ ...part, chunk }] : []; + }); + const merged: MergedTrack = { + durations: parts.flatMap((part) => part.durations), + sizes: parts.flatMap((part) => part.sizes), + sync: parts.every((part) => !part.sync) + ? null + : parts.flatMap((part) => part.sync ?? part.sizes.map(() => true)), + samplesPerChunk: parts.flatMap((part) => part.samplesPerChunk), + chunkOffsets: parts.flatMap((part) => + part.chunkOffsets.map((offset) => { + const i = part.chunk; + const { payloadStart, payloadEnd } = layouts[i]!; + if (offset < payloadStart || offset >= payloadEnd) { + throw new Error(`[mp4Stitch] chunk ${i} has samples outside its mdat`); + } + return offset - payloadStart + payloadBases[i]! + dataStart; + }), + ), + }; + const mediaDuration = sum(merged.durations); + const shown = Math.round( + ((mediaDuration - track.mediaTime) * movieTimescale) / track.timescale, + ); + const elst = Buffer.alloc(24); + elst.writeUInt32BE(1); + elst.writeBigUInt64BE(BigInt(shown), 4); + elst.writeBigInt64BE(BigInt(track.mediaTime), 12); + elst.writeUInt32BE(0x00010000, 20); + // An edit list rounds the shown length to the movie timescale, so write one only to skip + // leading samples (the AAC priming). + const edts = track.mediaTime ? box("edts", fullBox("elst", 1, elst)) : Buffer.alloc(0); + const trak = rebuild(template, track.trak, { + tkhd: (tkhd) => Buffer.concat([withDuration(template, tkhd, shown), edts]), + edts: () => Buffer.alloc(0), + mdia: (mdia) => + rebuild(template, mdia, { + mdhd: (mdhd) => withDuration(template, mdhd, mediaDuration), + minf: (minf) => rebuild(template, minf, { stbl: () => buildStbl(track.stsd, merged) }), + }), + }); + return { trak, shown }; + }); + let next = 0; + return rebuild(template, root, { + mvhd: (mvhd) => withDuration(template, mvhd, Math.max(...traks.map((t) => t.shown))), + trak: () => traks[next++]!.trak, + }); + }; + + const ftyp = layouts[0]!.ftyp; + const dataStart = ftyp.length + buildMoov(0).length + 16; + const mdat = Buffer.alloc(16); + mdat.writeUInt32BE(1); + mdat.write("mdat", 4, "latin1"); + mdat.writeBigUInt64BE(BigInt(16 + payloadBytes), 8); + return Buffer.concat([ftyp, buildMoov(dataStart), mdat]); +} + +function fileRange(path: string): ReadRange { + return async (start, end) => { + // Read-only, so the mode never applies; an explicit one keeps CodeQL's temp-file rule quiet. + const file = await open(path, "r", 0o600); + try { + const buf = Buffer.alloc(end - start); + const { bytesRead } = await file.read(buf, 0, buf.length, start); + return buf.subarray(0, bytesRead); + } finally { + await file.close(); + } + }; +} + +/** Copies chunk `index`'s mdat payload to `out`; a payload cut short fails, as an S3 copy would. */ +async function copyPayload(path: string, layout: Mp4Layout, index: number, out: Writable) { + const { payloadStart, payloadEnd } = layout; + const payload = createReadStream(path, { start: payloadStart, end: payloadEnd - 1 }); + await pipeline(payload, out, { end: false }); + if (payload.bytesRead !== payloadEnd - payloadStart) { + throw new Error( + `[mp4Stitch] chunk ${index} mdat holds ${payload.bytesRead} of ${payloadEnd - payloadStart} payload bytes`, + ); + } +} + +/** Local join: the stitched header, then each chunk's mdat payload in order. Like the ffmpeg + * join, it writes in a work folder beside `outputPath` and renames only a finished file there. */ +export async function writeStitchedMp4( + chunkPaths: readonly string[], + outputPath: string, +): Promise { + const layouts = await Promise.all(chunkPaths.map((path) => readMp4Layout(fileRange(path)))); + const header = buildStitchedHeader(layouts); + mkdirSync(dirname(outputPath), { recursive: true }); + const workDir = mkdtempSync(`${outputPath}.stitch-`); + const out = createWriteStream(join(workDir, "out.mp4")); + try { + out.write(header); + for (const [i, path] of chunkPaths.entries()) await copyPayload(path, layouts[i]!, i, out); + out.end(); + await finished(out); + renameSync(join(workDir, "out.mp4"), outputPath); + } finally { + out.destroy(); + rmSync(workDir, { recursive: true, force: true }); + } +} diff --git a/packages/producer/src/services/distributed/plan.test.ts b/packages/producer/src/services/distributed/plan.test.ts index d7322cc5bf..da9c7cfc4a 100644 --- a/packages/producer/src/services/distributed/plan.test.ts +++ b/packages/producer/src/services/distributed/plan.test.ts @@ -17,10 +17,23 @@ import { afterAll, beforeAll, describe, expect, it } from "bun:test"; import { spawnSync } from "node:child_process"; -import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { + chmodSync, + existsSync, + mkdirSync, + mkdtempSync, + readFileSync, + rmSync, + writeFileSync, +} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { applyConcreteGpuScreenshotClamp, buildChromeArgs } from "@hyperframes/engine"; +import { + applyConcreteGpuScreenshotClamp, + buildChromeArgs, + FFMPEG_PATH_ENV, + getFfmpegBinary, +} from "@hyperframes/engine"; import { recomputePlanHashFromPlanDir } from "../render/stages/freezePlan.js"; import { defaultLogger } from "../../logger.js"; import { RenderQualityError } from "../renderOrchestrator.js"; @@ -889,7 +902,7 @@ describe("plan() — golden planDir + planHash determinism", () => { }); describe.skipIf(spawnSync("ffmpeg", ["-version"]).status !== 0)("stitched mp4 plans", () => { - function audioProject(name: string): string { + function audioProject(name: string, amplitude = 0.99, duration = 1): string { const dir = join(runRoot, name); mkdirSync(dir, { recursive: true }); const tone = spawnSync("ffmpeg", [ @@ -899,7 +912,7 @@ describe.skipIf(spawnSync("ffmpeg", ["-version"]).status !== 0)("stitched mp4 pl "-f", "lavfi", "-i", - "aevalsrc=0.99*sin(2*PI*440*t):s=48000:d=0.6", + `aevalsrc=${amplitude}*sin(2*PI*440*t):s=48000:d=0.6`, "-c:a", "pcm_s16le", join(dir, "tone.wav"), @@ -908,7 +921,7 @@ describe.skipIf(spawnSync("ffmpeg", ["-version"]).status !== 0)("stitched mp4 pl writeFileSync( join(dir, "index.html"), ` -
+
`, ); @@ -950,6 +963,54 @@ describe.skipIf(spawnSync("ffmpeg", ["-version"]).status !== 0)("stitched mp4 pl const cfrEncoder = JSON.parse(readFileSync(join(withCfr, "meta", "encoder.json"), "utf-8")); expect(cfrEncoder.chunksCarryAudio).toBeUndefined(); }, 60_000); + + it("encodes the audio once, at the frames' exact length", async () => { + // 1.01 s at 30 fps is 31 frames: 49,600 samples, 48.4375 AAC frames of 1,024. + const project = audioProject("project-one-encode", 0.3, 1.01); + const shimDir = join(runRoot, "ffmpeg-shim"); + mkdirSync(shimDir, { recursive: true }); + const shim = join(shimDir, "ffmpeg"); + const calls = join(shimDir, "calls.txt"); + writeFileSync(shim, `#!/bin/sh\necho "$*" >> "${calls}"\nexec "${getFfmpegBinary()}" "$@"\n`); + chmodSync(shim, 0o755); + const planDir = join(runRoot, "plan-one-encode"); + const previous = process.env[FFMPEG_PATH_ENV]; + process.env[FFMPEG_PATH_ENV] = shim; + try { + await buildLocalExecutionPlan( + project, + { fps: 30, width: 320, height: 240, format: "mp4" }, + planDir, + { chunksCarryAudio: true }, + ); + } finally { + if (previous === undefined) delete process.env[FFMPEG_PATH_ENV]; + else process.env[FFMPEG_PATH_ENV] = previous; + } + const aacEncodes = readFileSync(calls, "utf-8") + .split("\n") + .filter((line) => / aac -b:a /.test(line)); + expect(aacEncodes).toHaveLength(1); + const planAudio = join(planDir, "audio.m4a"); + const shown = spawnSync( + "ffprobe", + [ + "-v", + "error", + "-select_streams", + "a:0", + "-show_entries", + "stream=duration_ts", + "-of", + "csv=p=0", + "--", + planAudio, + ], + { encoding: "utf-8" }, + ); + // The edit list counts in ms, so 49,600 samples (1.0333 s) show as 1.033 s. + expect(Math.abs(Number(shown.stdout) - 49_600)).toBeLessThan(48); + }, 60_000); }); describe("plan() — codec knob", () => { diff --git a/packages/producer/src/services/distributed/plan.ts b/packages/producer/src/services/distributed/plan.ts index 441083ac51..ec62cea24b 100644 --- a/packages/producer/src/services/distributed/plan.ts +++ b/packages/producer/src/services/distributed/plan.ts @@ -54,7 +54,7 @@ import { type RenderStrictness, } from "../renderOrchestrator.js"; import { closeFileServerSafely } from "../fileServer.js"; -import { padOrTrimAudioToVideoFrameCount } from "../render/audioPadTrim.js"; +import { constrainAacTruePeak } from "../render/audioPadTrim.js"; import { encoderFailureError } from "../render/encoderInterruption.js"; import { runAudioStage } from "../render/stages/audioStage.js"; import { runCompileStage } from "../render/stages/compileStage.js"; @@ -880,32 +880,6 @@ function carriesChunkAudio( return options.chunksCarryAudio === true && config.format === "mp4" && config.cfr !== true; } -/** - * Chunks that carry their own audio cut it from the plan track, so the length and true-peak fix - * the assembler applies to a joined file happens here, once, before any chunk reads it. - */ -async function fitPlanAudioToFrames( - audioPath: string, - workDir: string, - totalFrames: number, - fps: number, - run: { signal?: AbortSignal; timeoutMs: number }, -): Promise { - const fittedPath = join(workDir, "audio-fitted.m4a"); - const result = await padOrTrimAudioToVideoFrameCount({ - videoPath: audioPath, - probeVideoFrameInfo: async () => ({ frameCount: totalFrames, fpsNum: fps, fpsDen: 1 }), - audioPath, - outputPath: fittedPath, - signal: run.signal, - timeoutMs: run.timeoutMs, - }); - if (!result.success) { - throw encoderFailureError("[plan] audio pad/trim failed", result); - } - renameSync(fittedPath, audioPath); -} - /** * Resolve the encoder + pixel-format + preset triple for a distributed * render. Distributed mode is SDR-only: H.264 or H.265 8-bit for mp4, @@ -1373,13 +1347,17 @@ export async function buildLocalExecutionPlan( // Only the assembler reads the mix, so hashing and publishing the final // compiled and video trees run beside it instead of after it. const digests: PlanFileDigests = options.digests ?? new Map(); + const chunksCarryAudio = carriesChunkAudio(config, options); + // Chunks cut their audio from this track, so it is mixed to the frames' exact length in its one + // encode; the assembler pads or trims the track of any other plan. + const audioSeconds = chunksCarryAudio ? totalFrames / config.fps : probeResult.duration; stageStarted = Date.now(); const [audioOutcome, overlapOutcome] = await Promise.allSettled([ runAudioStage({ projectDir, workDir, compiledDir: finalCompiledDir, - duration: probeResult.duration, + duration: audioSeconds, ffmpegProcessTimeout: cfg.ffmpegProcessTimeout, audioGain: cfg.audioGain, audios: composition.audios, @@ -1410,12 +1388,16 @@ export async function buildLocalExecutionPlan( if (audioResult.hasAudio && existsSync(audioResult.audioOutputPath)) { renameSync(audioResult.audioOutputPath, planAudioPath); } - const chunksCarryAudio = carriesChunkAudio(config, options); if (chunksCarryAudio && existsSync(planAudioPath)) { - await fitPlanAudioToFrames(planAudioPath, workDir, totalFrames, config.fps, { + stageStarted = Date.now(); + const peak = await constrainAacTruePeak({ + audioPath: planAudioPath, + targetDurationSeconds: audioSeconds, signal: abortSignal, timeoutMs: cfg.ffmpegProcessTimeout, }); + if (!peak.success) throw encoderFailureError("[plan] audio true-peak check failed", peak); + logPlanStage(log, "audio-peak", stageStarted); } // ── Locked config ── @@ -1543,8 +1525,9 @@ export async function plan( projectDir: string, config: DistributedRenderConfig, planDir: string, + options: Pick = {}, ): Promise { - const executionPlan = await buildLocalExecutionPlan(projectDir, config, planDir); + const executionPlan = await buildLocalExecutionPlan(projectDir, config, planDir, options); return { planDir: executionPlan.executionPlanDir, planProtocol: PLAN_PROTOCOL_V1, diff --git a/packages/producer/src/services/distributed/planV2.test.ts b/packages/producer/src/services/distributed/planV2.test.ts index 6e49af0f3f..114b130be8 100644 --- a/packages/producer/src/services/distributed/planV2.test.ts +++ b/packages/producer/src/services/distributed/planV2.test.ts @@ -421,7 +421,7 @@ describe("Plan v2 manifest", () => { expect(assembler.some((artifact) => artifact.path === "compiled/index.html")).toBe(false); }); - it("also sends audio to every chunk when mp4 chunks carry their audio", () => { + it("sends audio to every chunk and not the assembler when mp4 chunks carry their audio", () => { const root = tempPath("hf-plan-v2-chunk-audio-targets-"); const v1 = createV1Plan(root, { audio: true, chunksCarryAudio: true }); const manifest = readPlanV2Manifest(createPlanV2FromV1(v1, join(root, "v2")).planDir); @@ -430,7 +430,7 @@ describe("Plan v2 manifest", () => { expect(chunk.some((artifact) => artifact.path === "audio.m4a")).toBe(true); } const assembler = listPlanV2ArtifactsForTarget(manifest, { role: "assembler" }); - expect(assembler.some((artifact) => artifact.path === "audio.m4a")).toBe(true); + expect(assembler.some((artifact) => artifact.path === "audio.m4a")).toBe(false); }); it("uses the runtime frame lookup to select exact video frames per chunk", () => { diff --git a/packages/producer/src/services/distributed/planV2.ts b/packages/producer/src/services/distributed/planV2.ts index 282b1d81d2..2f83170bfd 100644 --- a/packages/producer/src/services/distributed/planV2.ts +++ b/packages/producer/src/services/distributed/planV2.ts @@ -271,7 +271,7 @@ function artifactTargets( chunksCarryAudio: boolean, ): Pick { if (isPlanAudioArtifactPath(path)) { - return { chunks: chunksCarryAudio ? "all" : [], assembler: true }; + return chunksCarryAudio ? { chunks: "all", assembler: false } : { chunks: [], assembler: true }; } if (path === "plan.json" || path === "meta/chunks.json" || path === "meta/encoder.json") { return { chunks: "all", assembler: true }; diff --git a/packages/producer/src/services/render/audioPadTrim.ts b/packages/producer/src/services/render/audioPadTrim.ts index ef1bdd45d3..2ebe8f7337 100644 --- a/packages/producer/src/services/render/audioPadTrim.ts +++ b/packages/producer/src/services/render/audioPadTrim.ts @@ -428,6 +428,32 @@ export async function padOrTrimAudioToVideoFrameCount( }; } +/** + * The true-peak half of {@link padOrTrimAudioToVideoFrameCount}, for AAC already encoded at its + * exact length: measures once and re-encodes only when the peak is over the ceiling. + */ +export async function constrainAacTruePeak(input: { + audioPath: string; + targetDurationSeconds: number; + signal?: AbortSignal; + timeoutMs?: number; +}): Promise<{ + success: boolean; + error?: string; + failureReason?: "external_interruption"; + loweredDb?: number; +}> { + const timeoutMs = input.timeoutMs ?? DEFAULT_CONFIG.ffmpegProcessTimeout; + return enforceAacTruePeak({ + audioPath: input.audioPath, + targetDurationSeconds: input.targetDurationSeconds, + signal: input.signal, + probeTruePeak: (audioPath, signal) => + defaultProbeAudioTruePeakDbfs(audioPath, timeoutMs, signal), + runner: (args) => defaultRunFfmpeg(args, timeoutMs, input.signal), + }); +} + interface EnforceAacTruePeakInput { audioPath: string; targetDurationSeconds: number; diff --git a/scripts/package-size-budgets.json b/scripts/package-size-budgets.json index 17426986f9..8f459fa934 100644 --- a/scripts/package-size-budgets.json +++ b/scripts/package-size-budgets.json @@ -1,5 +1,5 @@ { - "@hyperframes/aws-lambda": { "packed": 63198, "unpacked": 230499 }, + "@hyperframes/aws-lambda": { "packed": 66677, "unpacked": 242569 }, "@hyperframes/cli": { "packed": 7517641, "unpacked": 22401400 }, "@hyperframes/core": { "packed": 1636442, "unpacked": 3903211 }, "@hyperframes/engine": { "packed": 428347, "unpacked": 1436293 },