T-HUB/apps/media-worker/src/compositor/program.ts

import { CompositorLagError, FrameGate, type GateImpl } from "../gate/frame-gate";
import type { GateEvent, GateStats } from "../gate/protocol";
import { spawnManaged, type ManagedProcess, type StderrRedactor } from "../relay/ffmpeg";
import { programTapGeometry } from "../preview/argv";
import { buildEncoderArgs, type Canvas, type EncoderOutput, type EncoderProgramTap } from "./command-builder";
import { CompositorInstance, type ComposeRequest } from "./compositor";
import { SourceBus, type BusSubscription } from "./source-bus";

export interface ProgramChainOptions {
  canvas: Canvas;
  scratchDir: string;
  videoBitrateKbps: number;
  /**
   * flvPipe: the push relay's inherited fd (WP9/WP13); null: x264 into `-f null` (CPU smoke);
   * external: no encoder is spawned, the caller has opened the gate's encoder FIFOs (verifying sinks).
   */
  output: EncoderOutput | { kind: "external" };
  redactor: StderrRedactor;
  onStderr(message: string): void;
  /** Encoder or gate died: the RTMP session is gone and the supervisor must restart the channel. */
  onFatal(reason: string): void;
  /** The on-air compositor died (gate is repeating its last frame); the chain respawns it. */
  onCompositorLost?(reason: string): void;
  gateImpl?: GateImpl;
  /** Swap must go on air within this long, or B is abandoned and A stays (default 10 s). */
  swapTimeoutMs?: number;
  /**
   * Program preview tap (§2.1): the encoder also emits a 360p copy of what it encodes into
   * `programBus`, which preview taps subscribe to. Needs an encoder (not `external`).
   */
  programTap?: boolean;
}

/** Child fd of the push relay's pipe in the encoder (flvPipe), and of the program tap outputs after it. */
const ENCODER_OUTPUT_FD = 4;

const LAG_RETRIES = 2;

/**
 * The persistent part of a channel's program path (§1.4): frame gate + encoder,
 * started once, plus the A/B compositor swap. `swap()` starts B for the new
 * scene, waits until the gate splices it in at a frame boundary, then stops A.
 * The encoder and its push session are never restarted by a swap.
 */
export class ProgramChain {
  readonly gate: FrameGate;
  /** The program tap's bus (360p of the encoder input), when `programTap` is set. */
  readonly programBus: SourceBus | null;
  private encoder: ManagedProcess | null = null;
  private active: CompositorInstance | null = null;
  private activeRequest: ComposeRequest | null = null;
  private queue: Promise<unknown> = Promise.resolve();
  private stopping = false;
  swaps = 0;
  /** Swap-in compositors replaced after a gate `LAG` report. */
  lagRetries = 0;

  constructor(private readonly options: ProgramChainOptions) {
    // The Node reference gate has neither A/V pair lockstep nor origin alignment (s2a-report §5.3):
    // swaps through it can step backwards and skew A/V. Spike/contract comparison only (gate/spike.ts, gate.test.ts).
    if (options.gateImpl === "node") throw new Error("the Node reference gate is spike-only; a program chain needs the Rust gate");
    this.gate = new FrameGate({
      geometry: options.canvas, scratchDir: options.scratchDir, impl: options.gateImpl ?? "rust",
      redactor: options.redactor, onStderr: options.onStderr,
      onEvent: (event) => this.onGateEvent(event),
      onExit: (reason) => { if (!this.stopping) options.onFatal(reason); },
    });
    if (options.programTap && options.output.kind === "external") throw new Error("the program tap needs an encoder");
    this.programBus = options.programTap ? new SourceBus({
      geometry: programTapGeometry(options.canvas), scratchDir: options.scratchDir, redactor: options.redactor, onStderr: options.onStderr,
      // Losing the tap bus loses previews, never the program: the encoder's tee drops that slave.
      onExit: (reason) => options.onStderr(`program tap bus: ${reason}`),
    }) : null;
  }

  get activeCompositor(): CompositorInstance | null { return this.active; }
  get encoderPid(): number | undefined { return this.encoder?.child.pid; }

  async start(first: ComposeRequest): Promise<void> {
    const { canvas, output } = this.options;
    if (output.kind === "external") {
      await this.gate.start();
      await this.swap(first);
      return;
    }
    const inputs = this.gate.encoderInputs();
    const stdio: Array<number | "pipe" | "ignore"> = [inputs.videoFd, "ignore", "pipe", inputs.audioFd];
    // The push relay's fd is inherited as child fd 4, whatever its number in the worker.
    let encoderOutput: EncoderOutput = output;
    if (output.kind === "flvPipe") { stdio.push(output.fd); encoderOutput = { kind: "flvPipe", fd: ENCODER_OUTPUT_FD }; }
    let programTap: EncoderProgramTap | undefined;
    if (this.programBus) {
      const tap = await this.programBus.start();
      const first = stdio.length;
      stdio.push(tap.videoFd, tap.audioFd);
      programTap = { videoFd: first, audioFd: first + 1, size: programTapGeometry(canvas) };
    }
    let encoder: ManagedProcess;
    try {
      encoder = spawnManaged("ffmpeg", buildEncoderArgs({ canvas, videoBitrateKbps: this.options.videoBitrateKbps, audioBitrateKbps: 128, output: encoderOutput, programTap }),
        stdio, this.options.redactor, this.options.onStderr);
    } finally {
      this.programBus?.releaseIngestWriters();
    }
    this.encoder = encoder;
    void encoder.exited.then(({ code, signal }) => { if (!this.stopping) this.options.onFatal(`encoder exited code=${code} signal=${signal}`); });
    await encoder.started;
    await this.gate.start();
    await this.swap(first);
  }

  /** Serialized A/B swap. Resolves with the gate tick at which the new compositor went on air. */
  swap(request: ComposeRequest): Promise<number> {
    const run = this.queue.catch(() => undefined).then(() => this.swapNow(request));
    this.queue = run;
    return run;
  }

  /**
   * A B the gate reports as lagging (still behind A's on-air source position, typically a spawn
   * starved of CPU while draining its startup backlog) is replaced by a fresh one; A stays on air.
   */
  private async swapNow(request: ComposeRequest): Promise<number> {
    for (let attempt = 0; ; attempt += 1) {
      try {
        return await this.swapOnce(request);
      } catch (error) {
        if (!(error instanceof CompositorLagError) || attempt >= LAG_RETRIES) throw error;
        this.lagRetries += 1;
      }
    }
  }

  private async swapOnce(request: ComposeRequest): Promise<number> {
    const next = new CompositorInstance(this.gate, request, this.options.redactor, this.options.onStderr);
    const current = this.active;
    await next.start((subs) => (current ? this.originAfter(current, subs) : 0));
    let tick: number;
    try {
      tick = await this.gate.waitActive(next.slotId, this.options.swapTimeoutMs ?? 10_000);
    } catch (error) {
      try { this.gate.detach(next.slotId); } catch { /* gate gone */ }
      await next.stop();
      throw error;
    }
    const old = this.active;
    this.active = next;
    this.activeRequest = request;
    this.swaps += 1;
    if (old) void old.stop();
    return tick;
  }

  /**
   * Origin of B relative to on-air A, in A's numbering: through a bus both use
   * (exact: pair counters, minus pairs A had lost to backlog when B subscribed), else by wall clock.
   * Pairs B loses after attach only make B *ahead*: a forward skip, never backwards.
   */
  private originAfter(a: CompositorInstance, subs: BusSubscription[]): number | undefined {
    if (a.origin === undefined) return undefined;
    for (const sb of subs) {
      const sa = a.subscriptions.find((x) => x.bus === sb.bus);
      if (sa) return a.origin + sb.startPair - sa.startPair - (sb.droppedBefore(sa) ?? sa.dropped);
    }
    const sa = a.subscriptions[0];
    const sb = subs[0];
    if (!sa || !sb) return undefined;
    return a.origin + Math.round((Number(sb.startedAt - sa.startedAt) / 1e9) * this.options.canvas.fps);
  }

  private onGateEvent(event: GateEvent): void {
    if (event.type !== "EOF" || this.stopping || !this.active || event.id !== this.active.slotId || !this.activeRequest) return;
    // The on-air compositor died (crash, kill -9). The gate is already holding
    // the last frame + silence, so the encoder never notices; respawn the scene.
    const lost = this.active;
    this.active = null;
    void lost.stop();
    this.options.onCompositorLost?.(`compositor ${event.id} ended on air`);
    void this.swap(this.activeRequest).catch((error: unknown) => this.options.onCompositorLost?.(`respawn failed: ${String(error)}`));
  }

  stats(): Promise<GateStats> { return this.gate.stats(); }

  async stop(): Promise<void> {
    this.stopping = true;
    await this.queue.catch(() => undefined);
    await this.active?.stop();
    this.active = null;
    await this.gate.stop();
    await this.encoder?.terminate();
    this.encoder = null;
    await this.programBus?.stop();
  }
}