UHD_Media_Tool/packages/runtime/src/worker.ts

import { Readable } from "node:stream";
import { type ChildProcess, Command, runtimeWriteSync } from "./index.ts";

/** 本地执行协议。业务类型、重试、权限与超时策略由调用方负责。 */
export interface WorkerRequest {
  id: string;
  task: string;
  input: unknown;
}
export interface WorkerError {
  code: string;
  message: string;
  detail?: Record<string, unknown>;
}
export type WorkerReply =
  | { id: string; type: "result"; result: unknown }
  | { id: string; type: "error"; error: WorkerError };
export type WorkerMessage = WorkerReply | { id: string; type: "progress"; value: unknown };

export class WorkerTransportError extends Error {
  constructor(
    readonly kind: "spawn" | "protocol" | "exit" | "size",
    message: string,
    override readonly cause?: unknown,
  ) {
    super(message);
    this.name = "WorkerTransportError";
  }
}

const encoder = new TextEncoder();
export function encodeWorkerMessage(value: WorkerRequest | WorkerMessage, maxBytes: number): Uint8Array {
  const bytes = encoder.encode(`${JSON.stringify(value)}\n`);
  if (bytes.byteLength > maxBytes) throw new WorkerTransportError("size", `worker 消息超过 ${maxBytes} 字节`);
  return bytes;
}

/** 按字节限制一行的大小；UTF-8 与换行可以跨管道分块。EOF 不能代替消息结束符。 */
export async function* workerLines(
  reader: ReadableStreamDefaultReader<Uint8Array>,
  maxBytes: number,
): AsyncGenerator<string> {
  const decoder = new TextDecoder("utf-8", { fatal: true });
  let parts: string[] = [];
  let length = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) {
      if (length) throw new WorkerTransportError("protocol", "worker 消息未以换行结束");
      return;
    }
    let start = 0;
    for (let i = 0; i < value.length; i++) {
      if (value[i] !== 10) continue;
      length += i - start + 1;
      if (length > maxBytes) throw new WorkerTransportError("size", `worker 消息超过 ${maxBytes} 字节`);
      parts.push(decoder.decode(value.subarray(start, i), { stream: true }), decoder.decode());
      const line = parts.join("");
      parts = [];
      length = 0;
      start = i + 1;
      if (!line) throw new WorkerTransportError("protocol", "worker 返回空消息");
      yield line;
    }
    length += value.length - start;
    if (length > maxBytes) throw new WorkerTransportError("size", `worker 消息超过 ${maxBytes} 字节`);
    parts.push(decoder.decode(value.subarray(start), { stream: true }));
  }
}

export async function readWorkerRequestText(maxBytes: number): Promise<string> {
  const stream = Readable.toWeb(process.stdin) as unknown as ReadableStream<Uint8Array>;
  const reader = stream.getReader();
  try {
    let text: string | undefined;
    for await (const line of workerLines(reader, maxBytes)) {
      if (text !== undefined) throw new WorkerTransportError("protocol", "一次性 worker 只能接收一个请求");
      text = line;
    }
    if (text === undefined) throw new WorkerTransportError("protocol", "worker 没有收到请求");
    return text;
  } finally {
    reader.releaseLock();
  }
}

export function decodeWorkerRequest(text: string): WorkerRequest {
  const req = JSON.parse(text) as WorkerRequest;
  if (!req || typeof req.id !== "string" || !req.id || typeof req.task !== "string" || !("input" in req))
    throw new WorkerTransportError("protocol", "worker 请求格式无效");
  return req;
}

export function writeWorkerMessage(message: WorkerMessage, maxBytes: number): void {
  const bytes = encodeWorkerMessage(message, maxBytes);
  let offset = 0;
  while (offset < bytes.length) offset += runtimeWriteSync(bytes.subarray(offset));
}

function decodeMessage(line: string, id: string): WorkerMessage {
  const value = JSON.parse(line) as WorkerMessage;
  if (!value || value.id !== id) throw new WorkerTransportError("protocol", "worker 响应不属于当前请求");
  if (value.type === "result" && "result" in value) return value;
  if (value.type === "progress" && "value" in value) return value;
  if (value.type === "error" && typeof value.error?.code === "string" && typeof value.error?.message === "string")
    return value;
  throw new WorkerTransportError("protocol", "worker 响应格式无效");
}

/** 一个一次性进程连接。只管通信与回收，不重试、不决定超时、不安装权限策略。 */
export class WorkerConnection {
  readonly pid: number;
  private readonly child: ChildProcess;
  private readonly stdout: ReadableStreamDefaultReader<Uint8Array>;
  private readonly stderrReader: ReadableStreamDefaultReader<Uint8Array>;
  private reject: ((error: unknown) => void) | null = null;
  private terminating: unknown;
  private killTimer: ReturnType<typeof setTimeout> | undefined;
  private settled = false;
  private started = false;
  private stderr = "";

  constructor(command: { cmd: string; args: string[] }, env?: { values: Record<string, string>; clear?: boolean }) {
    try {
      this.child = new Command(command.cmd, {
        args: command.args,
        stdin: "piped",
        stdout: "piped",
        stderr: "piped",
        env: env?.values,
        clearEnv: env?.clear,
      }).spawn();
    } catch (error) {
      throw new WorkerTransportError("spawn", "worker 子进程无法启动", error);
    }
    this.pid = this.child.pid;
    this.stdout = this.child.stdout.getReader();
    this.stderrReader = this.child.stderr.getReader();
  }

  /** graceMs > 0 时先 SIGTERM 供文件操作清理；解析可直接 SIGKILL。等待始终有界。 */
  terminate(error: unknown, graceMs = 0, onForceKill?: () => void): void {
    if (this.settled || this.terminating !== undefined) return;
    this.terminating = error;
    const force = () => {
      try {
        this.child.kill("SIGKILL");
        onForceKill?.();
      } catch {
        // 已退出
      }
      this.child.unref();
      this.fail(error);
    };
    if (graceMs === 0) {
      force();
      return;
    }
    try {
      this.child.kill("SIGTERM");
    } catch {
      // 已退出
    }
    this.killTimer = setTimeout(force, graceMs);
  }

  request(
    req: WorkerRequest,
    options: { maxRequestBytes: number; maxResponseBytes: number; onProgress?: (value: unknown) => void },
  ): Promise<WorkerReply> {
    if (this.started) throw new Error("一次性 worker 连接不能重复请求");
    this.started = true;
    if (this.terminating !== undefined) return Promise.reject(this.terminating);
    if (!this.pid)
      return this.child.status.then(
        () => {
          this.cleanup();
          throw new WorkerTransportError("exit", "worker 未启动");
        },
        (error) => {
          this.cleanup();
          throw new WorkerTransportError("spawn", "worker 子进程无法启动", error);
        },
      );
    return new Promise((resolve, reject) => {
      this.reject = reject;
      const errPump = this.readStderr();
      let final: WorkerReply | undefined;
      const pump = (async () => {
        for await (const line of workerLines(this.stdout, options.maxResponseBytes)) {
          const msg = decodeMessage(line, req.id);
          if (final) throw new WorkerTransportError("protocol", "worker 结果之后仍发送消息");
          if (msg.type === "progress") options.onProgress?.(msg.value);
          else final = msg;
        }
      })();
      pump.catch((error) => this.terminate(this.transportError(error)));
      this.child.status.then(
        async (status) => {
          try {
            await Promise.all([pump, errPump]);
            if (this.settled) return;
            if (this.terminating !== undefined) return this.fail(this.terminating);
            if (!status.success || !final)
              throw new WorkerTransportError(
                "exit",
                `worker 异常退出(code=${status.code} signal=${status.signal}):${this.stderr}`,
              );
            this.cleanup();
            resolve(final);
          } catch (error) {
            this.fail(this.terminating ?? this.transportError(error));
          }
        },
        (error) => this.fail(new WorkerTransportError("spawn", "worker 子进程无法启动", error)),
      );
      (async () => {
        const bytes = encodeWorkerMessage(req, options.maxRequestBytes);
        const writer = this.child.stdin.getWriter();
        try {
          await writer.write(bytes);
          await writer.close();
        } finally {
          writer.releaseLock();
        }
      })().catch((error) => this.terminate(this.transportError(error)));
      if (this.terminating !== undefined) this.fail(this.terminating);
    });
  }

  private transportError(error: unknown): WorkerTransportError {
    return error instanceof WorkerTransportError
      ? error
      : new WorkerTransportError("protocol", `worker 通信失败:${(error as Error).message}`, error);
  }

  private async readStderr(): Promise<void> {
    const decoder = new TextDecoder();
    try {
      for (;;) {
        const { done, value } = await this.stderrReader.read();
        if (done) return;
        this.stderr = (this.stderr + decoder.decode(value, { stream: true })).slice(-2000);
      }
    } catch {
      // 终止时管道可能断开
    }
  }

  private fail(error: unknown): void {
    if (this.settled) return;
    this.cleanup();
    this.reject?.(error);
  }

  private cleanup(): void {
    this.settled = true;
    clearTimeout(this.killTimer);
    void this.stdout.cancel().catch(() => {});
    void this.stderrReader.cancel().catch(() => {});
  }
}