mirror of
https://github.com/heygen-com/hyperframes.git
synced 2026-09-11 06:30:03 +00:00
feat(producer): add shaderTransitionWorkerPool (hf#732 PR 3/5) (#758)
## Summary PR 3 of 5 in the hf#732 decomposition stack. Adds a `worker_threads`-based pool that runs the shader-transition blend (one of 15 transition shaders) on a fixed-size worker pool. **No production wiring yet** — the pool stands alone; PR 4 wires it. The shader blend is a hot inner loop over every pixel of every transition frame at 16bpc. Moving it off the main event loop removes the JS-event-loop ceiling that capped throughput in earlier hf#732 iterations. ### New files - `packages/producer/src/services/shaderTransitionWorker.ts` — worker entry. Imports from `@hyperframes/engine/shader-transitions` (zero-import TS source). - `packages/producer/src/services/shaderTransitionWorkerPool.ts` — fixed-size pool. Uses `transferList` so the 16bpc HDR `from`/`to`/`out` buffers move by ownership. - `packages/producer/src/services/shaderTransitionWorkerPool.test.ts` — 6 vitest tests pinning byte-equivalence across all 15 shaders, transferList correctness, pool lifecycle. All pass. ### Build wiring - `packages/cli/tsup.config.ts`: third tsup entry emits `dist/shaderTransitionWorker.js`. - `packages/producer/build.mjs`: fourth esbuild entry for direct producer consumers. - `packages/engine/package.json`: adds `./shader-transitions` subpath export. ## Stack Stacked on top of #757 (PR 2: pngDecodeBlit pool). No behavior change in any render. ## Test plan - [x] 6 pool tests pass - [x] Producer + engine typecheck clean - [x] oxlint clean — Vai
This commit is contained in:
Vance Ingalls
@@ -0,0 +1,382 @@
|
||||
/**
|
||||
* Pool of Node `worker_threads` Workers for off-main-thread shader-blend
|
||||
* execution. See `shaderTransitionWorker.ts` for the per-worker contract and
|
||||
* the hf#677 follow-up rationale (closing the JS event-loop ceiling on the
|
||||
* layered transition path).
|
||||
*
|
||||
* Pool shape:
|
||||
*
|
||||
* - Spawned once at the start of a layered render and terminated in the
|
||||
* `finally`. Worker spawn cost is ~10–50 ms each; amortized over the
|
||||
* full transition phase (typically 100+ frames) it's negligible.
|
||||
* - Pool size is sized to `min(layeredWorkerCount, cpuCount)`. We don't
|
||||
* spawn more workers than DOM sessions (no benefit — at most N DOM
|
||||
* sessions can be dispatching to us at any moment) and we don't oversubscribe
|
||||
* beyond physical cores.
|
||||
* - Each Worker holds zero per-frame state. Pool simply dispatches one
|
||||
* shader-blend per Worker at a time; ordering within the pool doesn't
|
||||
* matter because each frame's output is gated by the encoder's
|
||||
* `FrameReorderBuffer` upstream.
|
||||
*
|
||||
* API:
|
||||
*
|
||||
* const pool = await createShaderTransitionWorkerPool({ size, log });
|
||||
* const result = await pool.run({
|
||||
* shader, bufferA, bufferB, output, width, height, progress,
|
||||
* });
|
||||
* // result.bufferA / result.bufferB / result.output are the same memory,
|
||||
* // now re-attached to the main thread.
|
||||
* await pool.terminate();
|
||||
*
|
||||
* Buffer transfer contract: `run` takes Node Buffers, transfers their
|
||||
* underlying ArrayBuffers to the worker, and returns NEW Buffer views over
|
||||
* the transferred-back ArrayBuffers. The caller is responsible for
|
||||
* swapping its Buffer references — the *original* Buffers passed in are
|
||||
* detached (their `.length` becomes 0 / accessing throws) after `run` resolves.
|
||||
*/
|
||||
|
||||
import { Worker } from "node:worker_threads";
|
||||
import { fileURLToPath, pathToFileURL } from "node:url";
|
||||
import { dirname, join } from "node:path";
|
||||
import { createRequire } from "node:module";
|
||||
import { existsSync } from "node:fs";
|
||||
import { cpus } from "node:os";
|
||||
|
||||
interface PoolLogger {
|
||||
info?: (msg: string, meta?: Record<string, unknown>) => void;
|
||||
warn?: (msg: string, meta?: Record<string, unknown>) => void;
|
||||
error?: (msg: string, meta?: Record<string, unknown>) => void;
|
||||
}
|
||||
|
||||
export interface ShaderTransitionPoolOptions {
|
||||
/** Number of worker threads. Clamped to [1, cpus().length]. */
|
||||
size: number;
|
||||
/** Optional logger; falls back to no-op. */
|
||||
log?: PoolLogger;
|
||||
/**
|
||||
* Absolute filesystem path to the worker entry module. When provided, the
|
||||
* pool spawns workers from this exact path and skips the fallback
|
||||
* `import.meta.url`-based resolver entirely. Required by callers that
|
||||
* bundle the worker via a separate build (e.g. the CLI's tsup bundle):
|
||||
* `import.meta.url` inside the bundled pool resolves to the bundle's own
|
||||
* location, NOT the bundled worker entry's location, so the heuristic
|
||||
* resolver below cannot find the worker. Path extension determines the
|
||||
* loader behaviour (`.ts` → tsx/esm loader is appended to execArgv).
|
||||
*/
|
||||
workerEntryPath?: string;
|
||||
}
|
||||
|
||||
export interface ShaderBlendRequest {
|
||||
shader: string;
|
||||
bufferA: Buffer;
|
||||
bufferB: Buffer;
|
||||
output: Buffer;
|
||||
width: number;
|
||||
height: number;
|
||||
progress: number;
|
||||
}
|
||||
|
||||
export interface ShaderBlendResult {
|
||||
/** Re-attached buffer A (zero-copy view over the transferred-back ArrayBuffer). */
|
||||
bufferA: Buffer;
|
||||
/** Re-attached buffer B. */
|
||||
bufferB: Buffer;
|
||||
/** Re-attached output buffer holding the shader-blended frame. */
|
||||
output: Buffer;
|
||||
}
|
||||
|
||||
interface PendingTask {
|
||||
req: ShaderBlendRequest;
|
||||
resolve: (r: ShaderBlendResult) => void;
|
||||
reject: (err: Error) => void;
|
||||
/** Set when `HF_SHADER_POOL_TRACE=1`; used to log dispatch latency. */
|
||||
enqueuedAtMs?: number;
|
||||
/** Set when `HF_SHADER_POOL_TRACE=1`; assigned at dispatch. */
|
||||
traceId?: number;
|
||||
}
|
||||
|
||||
interface WorkerSlot {
|
||||
worker: Worker;
|
||||
busy: boolean;
|
||||
current: PendingTask | null;
|
||||
}
|
||||
|
||||
interface WorkerReply {
|
||||
ok: boolean;
|
||||
error?: string;
|
||||
bufferA: ArrayBuffer;
|
||||
bufferB: ArrayBuffer;
|
||||
output: ArrayBuffer;
|
||||
}
|
||||
|
||||
export interface ShaderTransitionWorkerPool {
|
||||
readonly size: number;
|
||||
run(req: ShaderBlendRequest): Promise<ShaderBlendResult>;
|
||||
terminate(): Promise<void>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the path to the compiled worker module.
|
||||
*
|
||||
* Resolution order (first match wins):
|
||||
* 1. Explicit `workerEntryPath` factory option — callers that bundle the
|
||||
* worker via a separate build pipeline (e.g. the CLI's tsup bundle that
|
||||
* emits `shaderTransitionWorker.js` next to `cli.js`) must use this.
|
||||
* The bundled-CLI case is the *only* one where the fallback below
|
||||
* cannot find the worker: `import.meta.url` inside the inlined pool
|
||||
* resolves to the bundle path, not the worker's emitted path, so the
|
||||
* sibling probe lands in the wrong directory.
|
||||
* 2. `HF_SHADER_WORKER_ENTRY` env var — test/dev infra override (file
|
||||
* path or `file://` URL).
|
||||
* 3. Same-directory `.js` sibling — works when both pool source and
|
||||
* worker source compile into the same `dist/services/` directory
|
||||
* (in-tree dev builds and the colocated tsc emit).
|
||||
* 4. Same-directory `.ts` sibling — vitest/bun raw-TS execution path.
|
||||
*/
|
||||
function resolveWorkerEntry(explicit: string | undefined): { path: string; isTs: boolean } {
|
||||
if (explicit && explicit.length > 0) {
|
||||
return { path: explicit, isTs: explicit.endsWith(".ts") };
|
||||
}
|
||||
const override = process.env.HF_SHADER_WORKER_ENTRY;
|
||||
if (override && override.length > 0) {
|
||||
const isTs = override.endsWith(".ts");
|
||||
return { path: override, isTs };
|
||||
}
|
||||
const moduleDir = dirname(fileURLToPath(import.meta.url));
|
||||
const jsPath = join(moduleDir, "shaderTransitionWorker.js");
|
||||
if (existsSync(jsPath)) return { path: jsPath, isTs: false };
|
||||
const tsPath = join(moduleDir, "shaderTransitionWorker.ts");
|
||||
return { path: tsPath, isTs: true };
|
||||
}
|
||||
|
||||
/**
|
||||
* Probe whether the parent process already has a TS loader registered
|
||||
* (tsx, ts-node, esm-loader). Worker threads inherit the parent's loader
|
||||
* only if we copy `process.execArgv` AND the relevant flag is present.
|
||||
* Vitest runs its own transformer and does NOT register a loader on
|
||||
* `process.execArgv`, so when the resolved entry is `.ts` and no loader
|
||||
* is detected we try to inject `tsx/esm` so `new Worker(<.ts file>)`
|
||||
* loads correctly.
|
||||
*
|
||||
* This is best-effort: if `tsx/esm` can't be resolved (e.g. minimal prod
|
||||
* install), we fall back to plain `process.execArgv` and the Worker will
|
||||
* surface a clear "cannot find module" error rather than silently
|
||||
* misbehaving.
|
||||
*/
|
||||
function buildExecArgv(entryIsTs: boolean): string[] {
|
||||
const inherited = [...process.execArgv];
|
||||
if (!entryIsTs) return inherited;
|
||||
const hasLoader = inherited.some(
|
||||
(a) => a.includes("tsx/esm") || a.includes("ts-node/esm") || a.includes("--import"),
|
||||
);
|
||||
if (hasLoader) return inherited;
|
||||
try {
|
||||
const require = createRequire(import.meta.url);
|
||||
const tsxEsm = require.resolve("tsx/esm");
|
||||
inherited.push("--import", pathToFileURL(tsxEsm).href);
|
||||
} catch {
|
||||
// tsx not installed (prod) — leave execArgv as-is. The caller will
|
||||
// get a clear error if the .ts entry can't be loaded.
|
||||
}
|
||||
return inherited;
|
||||
}
|
||||
|
||||
/**
|
||||
* Spawn a worker pool ready to run shader-blends. The returned pool is
|
||||
* usable as soon as the function resolves. If any worker fails to spawn,
|
||||
* all already-spawned workers are terminated and the error is propagated.
|
||||
*/
|
||||
export async function createShaderTransitionWorkerPool(
|
||||
opts: ShaderTransitionPoolOptions,
|
||||
): Promise<ShaderTransitionWorkerPool> {
|
||||
const cpuCount = Math.max(1, cpus().length);
|
||||
const size = Math.max(1, Math.min(opts.size, cpuCount));
|
||||
const log = opts.log ?? {};
|
||||
const { path: entry, isTs: entryIsTs } = resolveWorkerEntry(opts.workerEntryPath);
|
||||
|
||||
const slots: WorkerSlot[] = [];
|
||||
const queue: PendingTask[] = [];
|
||||
let terminated = false;
|
||||
|
||||
// hf#732 follow-up: instrumentation flag to log per-task dispatch /
|
||||
// completion timestamps so we can confirm the pool actually runs blends
|
||||
// concurrently when N DOM workers each dispatch K tasks. Enabled by
|
||||
// setting `HF_SHADER_POOL_TRACE=1`. Off by default — the per-task log
|
||||
// line is high-volume on long shader-transition renders.
|
||||
const traceEnabled = process.env.HF_SHADER_POOL_TRACE === "1";
|
||||
let nextTaskId = 0;
|
||||
|
||||
// Bind the parent's execArgv (e.g. tsx's `--import tsx/esm` loader) into
|
||||
// every Worker so a `.ts` entry point loads under tsx in dev without a
|
||||
// separate loader registration step. In the bundled prod build the
|
||||
// entry is `.js` and execArgv is typically empty — passing it is a no-op.
|
||||
// Under vitest the parent has no tsx loader on execArgv; `buildExecArgv`
|
||||
// appends one so the `.ts` worker entry still loads.
|
||||
const execArgv = buildExecArgv(entryIsTs);
|
||||
|
||||
const dispatchNext = (slot: WorkerSlot): void => {
|
||||
if (terminated || slot.busy) return;
|
||||
const task = queue.shift();
|
||||
if (!task) return;
|
||||
slot.busy = true;
|
||||
slot.current = task;
|
||||
if (traceEnabled) {
|
||||
const slotIdx = slots.indexOf(slot);
|
||||
const waitMs = task.enqueuedAtMs ? Date.now() - task.enqueuedAtMs : 0;
|
||||
const busyCount = slots.filter((s) => s.busy).length;
|
||||
log.info?.("[shaderPool] dispatch", {
|
||||
task: task.traceId,
|
||||
slot: slotIdx,
|
||||
shader: task.req.shader,
|
||||
waitMs,
|
||||
busyCount,
|
||||
queueDepth: queue.length,
|
||||
});
|
||||
}
|
||||
const { bufferA, bufferB, output, shader, width, height, progress } = task.req;
|
||||
// `Buffer.alloc` always returns a Buffer over a plain ArrayBuffer (not
|
||||
// SharedArrayBuffer) at runtime — TS narrows `.buffer` to the union
|
||||
// `ArrayBuffer | SharedArrayBuffer`, so cast at the boundary. The pool
|
||||
// would not work with SharedArrayBuffer-backed Buffers anyway because
|
||||
// transferList rejects them.
|
||||
const abA = bufferA.buffer as ArrayBuffer;
|
||||
const abB = bufferB.buffer as ArrayBuffer;
|
||||
const abOut = output.buffer as ArrayBuffer;
|
||||
try {
|
||||
slot.worker.postMessage(
|
||||
{
|
||||
shader,
|
||||
bufferA: abA,
|
||||
bufferB: abB,
|
||||
output: abOut,
|
||||
width,
|
||||
height,
|
||||
progress,
|
||||
},
|
||||
[abA, abB, abOut],
|
||||
);
|
||||
} catch (err) {
|
||||
// postMessage can throw if the ArrayBuffer was already detached
|
||||
// (e.g. caller reused a buffer mid-flight). Surface clearly.
|
||||
slot.busy = false;
|
||||
slot.current = null;
|
||||
task.reject(err instanceof Error ? err : new Error(String(err)));
|
||||
}
|
||||
};
|
||||
|
||||
const onWorkerMessage = (slot: WorkerSlot, reply: WorkerReply): void => {
|
||||
const task = slot.current;
|
||||
slot.current = null;
|
||||
slot.busy = false;
|
||||
if (!task) {
|
||||
// Spurious message; nothing to resolve. Drain queue anyway.
|
||||
dispatchNext(slot);
|
||||
return;
|
||||
}
|
||||
if (!reply.ok) {
|
||||
task.reject(new Error(reply.error ?? "shader-blend worker failed"));
|
||||
} else {
|
||||
task.resolve({
|
||||
bufferA: Buffer.from(reply.bufferA),
|
||||
bufferB: Buffer.from(reply.bufferB),
|
||||
output: Buffer.from(reply.output),
|
||||
});
|
||||
}
|
||||
dispatchNext(slot);
|
||||
};
|
||||
|
||||
const onWorkerError = (slot: WorkerSlot, err: Error): void => {
|
||||
const task = slot.current;
|
||||
slot.current = null;
|
||||
slot.busy = false;
|
||||
if (task) {
|
||||
// The in-flight task's buffers were transferred to the worker. They're
|
||||
// lost on the worker crash — the caller's original Buffers are
|
||||
// already detached. Reject so the render fails fast rather than
|
||||
// continuing with corrupted state.
|
||||
task.reject(new Error(`shader-blend worker crashed mid-task: ${err.message}; buffers lost`));
|
||||
}
|
||||
log.warn?.("[shaderTransitionWorkerPool] worker errored", { err: err.message });
|
||||
};
|
||||
|
||||
const onWorkerExit = (slot: WorkerSlot, code: number): void => {
|
||||
if (terminated) return;
|
||||
// Unexpected exit — fail any in-flight task and drop the slot. We don't
|
||||
// auto-respawn in the middle of a render because the lost transferList
|
||||
// buffers can't be reconstructed, and silently shrinking the pool would
|
||||
// mask the real failure. Pool teardown handles graceful shutdown.
|
||||
if (slot.current) {
|
||||
slot.current.reject(new Error(`shader-blend worker exited (code=${code}) mid-task`));
|
||||
slot.current = null;
|
||||
slot.busy = false;
|
||||
}
|
||||
log.warn?.("[shaderTransitionWorkerPool] worker exited unexpectedly", { code });
|
||||
};
|
||||
|
||||
// Spawn workers. If any throws synchronously we still want to terminate
|
||||
// the partially-spawned set before rejecting.
|
||||
try {
|
||||
for (let i = 0; i < size; i++) {
|
||||
const worker = new Worker(entry, { execArgv });
|
||||
const slot: WorkerSlot = { worker, busy: false, current: null };
|
||||
worker.on("message", (msg: WorkerReply) => onWorkerMessage(slot, msg));
|
||||
worker.on("error", (err: unknown) =>
|
||||
onWorkerError(slot, err instanceof Error ? err : new Error(String(err))),
|
||||
);
|
||||
worker.on("exit", (code) => onWorkerExit(slot, code));
|
||||
slots.push(slot);
|
||||
}
|
||||
} catch (err) {
|
||||
terminated = true;
|
||||
await Promise.all(slots.map((s) => s.worker.terminate().catch(() => undefined)));
|
||||
throw err;
|
||||
}
|
||||
|
||||
log.info?.("[shaderTransitionWorkerPool] spawned", { size, entry });
|
||||
|
||||
return {
|
||||
size,
|
||||
async run(req: ShaderBlendRequest): Promise<ShaderBlendResult> {
|
||||
if (terminated) {
|
||||
throw new Error("shader-blend pool already terminated");
|
||||
}
|
||||
return new Promise<ShaderBlendResult>((resolve, reject) => {
|
||||
const task: PendingTask = traceEnabled
|
||||
? { req, resolve, reject, enqueuedAtMs: Date.now(), traceId: ++nextTaskId }
|
||||
: { req, resolve, reject };
|
||||
// Find an idle slot; otherwise queue.
|
||||
const idle = slots.find((s) => !s.busy);
|
||||
if (idle) {
|
||||
queue.unshift(task);
|
||||
dispatchNext(idle);
|
||||
} else {
|
||||
queue.push(task);
|
||||
}
|
||||
});
|
||||
},
|
||||
async terminate(): Promise<void> {
|
||||
if (terminated) return;
|
||||
terminated = true;
|
||||
// Reject any queued (not-yet-dispatched) tasks. Their buffers are
|
||||
// still attached on the main thread — caller can recover.
|
||||
while (queue.length > 0) {
|
||||
const t = queue.shift();
|
||||
if (t) t.reject(new Error("shader-blend pool terminated before task ran"));
|
||||
}
|
||||
// Reject any in-flight tasks before worker.terminate() races with
|
||||
// the message reply. Calling Worker.terminate() forcefully stops
|
||||
// the worker; if a task was mid-execution its parentPort.postMessage
|
||||
// never lands, so the `current` task promise would otherwise leak.
|
||||
for (const slot of slots) {
|
||||
const t = slot.current;
|
||||
if (t) {
|
||||
slot.current = null;
|
||||
slot.busy = false;
|
||||
t.reject(new Error("shader-blend pool terminated mid-task"));
|
||||
}
|
||||
}
|
||||
await Promise.all(slots.map((s) => s.worker.terminate().catch(() => undefined)));
|
||||
log.info?.("[shaderTransitionWorkerPool] terminated", { size });
|
||||
},
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user