diff --git a/packages/cli/src/commands/snapshot.test.ts b/packages/cli/src/commands/snapshot.test.ts index 545350acca..177f14d4b4 100644 --- a/packages/cli/src/commands/snapshot.test.ts +++ b/packages/cli/src/commands/snapshot.test.ts @@ -22,6 +22,19 @@ const snapshotState = vi.hoisted(() => ({ noFfmpeg: false, })); +const spawned = vi.hoisted(() => ({ files: [] as string[] })); + +vi.mock("../utils/cancellableProcess.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + runCancellableProcess: (...args: Parameters) => { + spawned.files.push(String(args[0])); + return actual.runCancellableProcess(...args); + }, + }; +}); + vi.mock("../browser/ffmpeg.js", async (importOriginal) => { const actual = await importOriginal(); return { @@ -64,6 +77,7 @@ import snapshotCommand, { decodeFramesInOnePass, extractVideoFrameOrWarn, probeContainingSourceFrameSeek, + createSourceProbeCache, containingSourceFrameIndex, computeSnapshotTimes, formatSnapshotTimestamp, @@ -708,6 +722,47 @@ describe("extractVideoFrameToBuffer", () => { }, ); + it.skipIf(noProbe)( + "probes a clip once for all the times asked of it, and plans each time as it would alone", + async () => { + const dir = mkdtempSync(join(tmpdir(), "hf-snapshot-probe-cache-")); + try { + const clip = join(dir, "clip.mp4"); + execFileSync(ffmpeg!, [ + ...["-hide_banner", "-loglevel", "error", "-f", "lavfi"], + ...["-i", "testsrc2=d=20:r=25:s=320x180", "-c:v", "libx264", "-g", "50"], + ...["-pix_fmt", "yuv420p", clip], + ]); + const times = Array.from({ length: 40 }, (_, i) => 0.4 + i * 0.47); + spawned.files.length = 0; + const alone = []; + for (const time of times) + alone.push(await probeContainingSourceFrameSeek(clip, time, false)); + const spawnedAlone = spawned.files.length; + + spawned.files.length = 0; + const cache = createSourceProbeCache(); + const horizon = times.at(-1)!; + const cached = await Promise.all( + times.map((time) => + probeContainingSourceFrameSeek(clip, time, false, { cache, horizon }), + ), + ); + expect(cached).toEqual(alone); + const probes = spawned.files.filter((file) => file === ffprobe).length; + const landings = spawned.files.length - probes; + // One stream read and one packet read; a seek landing per keyframe (a keyframe every 2 s), not per time. + expect(probes).toBe(2); + expect(landings).toBeLessThanOrEqual(2 * Math.ceil(20 / 2)); + expect(spawned.files.length * 3, "far fewer processes than a probe per time").toBeLessThan( + spawnedAlone, + ); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, + ); + it.skipIf(noProbe)( "gives two times inside one source frame the same frame, without decoding on", async () => { diff --git a/packages/cli/src/commands/snapshot.ts b/packages/cli/src/commands/snapshot.ts index a2ab3a9842..1dc210c8c8 100644 --- a/packages/cli/src/commands/snapshot.ts +++ b/packages/cli/src/commands/snapshot.ts @@ -401,11 +401,24 @@ export async function remoteVideoInput(url: string, download = downloadToTemp): } } +export interface SourceProbeCache { + streams: Map | null>>; + packets: Map | null> }>; + landings: Map>; +} + +export const createSourceProbeCache = (): SourceProbeCache => ({ + streams: new Map(), + packets: new Map(), + landings: new Map(), +}); + /** Keyframe seek plus the decoded-time threshold that selects the frame containing `time`, or null. */ export async function probeContainingSourceFrameSeek( videoPath: string, time: number, holdLastFrame: boolean, + { cache, horizon = time }: { cache?: SourceProbeCache; horizon?: number } = {}, ): Promise<{ seek?: number; select?: number; decodesFrom?: number } | null> { const ffprobe = findFFprobe(); // Without ffprobe, fall back to a plain seek: the first frame at or after `time`. @@ -434,7 +447,10 @@ export async function probeContainingSourceFrameSeek( ); } }; - const data = await probe("%+#1", "stream=time_base,start_pts:format=start_time,duration"); + const streamQuery = () => probe("%+#1", "stream=time_base,start_pts:format=start_time,duration"); + const data = await (cache + ? (cache.streams.get(videoPath) ?? cache.streams.set(videoPath, streamQuery()).get(videoPath)!) + : streamQuery()); if (data === null || !Array.isArray(data.streams)) return null; const stream: unknown = data.streams[0]; if ( @@ -464,7 +480,18 @@ export async function probeContainingSourceFrameSeek( const startTicks = Math.abs(startPts * timeBase - formatStart) <= 1e-6 ? startPts : formatStart / timeBase; // An absolute end: edit-list preroll packets start before the video, so a relative window can stop short. - const packetData = await probe(`%${startTicks * timeBase + time + 1}`, "packet=pts,dts,flags"); + const packetEnd = Math.max(time, horizon) + 1; + const packetQuery = () => probe(`%${startTicks * timeBase + packetEnd}`, "packet=pts,dts,flags"); + const kept = cache?.packets.get(videoPath); + const packetData = await (kept && kept.end >= packetEnd + ? kept.data + : cache + ? (() => { + const data = packetQuery(); + cache.packets.set(videoPath, { end: packetEnd, data }); + return data; + })() + : packetQuery()); if (packetData === null || !Array.isArray(packetData.packets)) return null; const packets: { pts: number; dts: number; key: boolean }[] = []; for (const packet of packetData.packets) { @@ -487,7 +514,12 @@ export async function probeContainingSourceFrameSeek( // seek lands; a seek to exactly that keyframe trims its leading frames, so it is `start_t`. const ffmpeg = findFFmpeg(); const landing = async (seek: number) => { - const raw = ffmpeg ? await probeSeekLanding(ffmpeg, videoPath, seek) : null; + const key = `${videoPath}\0${seek}`; + const asked = + cache?.landings.get(key) ?? + (ffmpeg ? probeSeekLanding(ffmpeg, videoPath, seek) : Promise.resolve(null)); + cache?.landings.set(key, asked); + const raw = await asked; return raw === null ? null : raw - startTicks * timeBase; }; const candidates = packets.filter( @@ -1101,8 +1133,13 @@ async function captureSnapshots( console.warn(` ${c.warn("⚠")} No player API — seeks will be skipped`); } + const sourceProbes = createSourceProbeCache(); const frameBank = createFrameBank({ - probe: probeContainingSourceFrameSeek, + probe: (path, time, holdLastFrame, horizon) => + probeContainingSourceFrameSeek(path, time, holdLastFrame, { + cache: sourceProbes, + horizon, + }), decode: async (path, frames) => decodeFramesInOnePass( path, diff --git a/packages/cli/src/commands/snapshotFrameBank.ts b/packages/cli/src/commands/snapshotFrameBank.ts index 45854d1405..608d210c0d 100644 --- a/packages/cli/src/commands/snapshotFrameBank.ts +++ b/packages/cli/src/commands/snapshotFrameBank.ts @@ -19,7 +19,12 @@ export interface PlannedFrame extends FrameWant { } interface FrameBankDeps { - probe: (path: string, time: number, holdLastFrame: boolean) => Promise; + probe: ( + path: string, + time: number, + holdLastFrame: boolean, + horizon: number, + ) => Promise; /** One ffmpeg pass over `frames`, which share a displayed size; throws when the pass fails. */ decode: (path: string, frames: PlannedFrame[]) => Promise; } @@ -62,8 +67,9 @@ async function probeFrames( wants: FrameWant[], ): Promise { try { + const horizon = Math.max(...wants.map((want) => want.time)); const plans = await Promise.all( - wants.map((want) => probe(path, want.time, want.holdLastFrame)), + wants.map((want) => probe(path, want.time, want.holdLastFrame, horizon)), ); return plans.flatMap((plan, i) => plan?.seek !== undefined && plan.select !== undefined