chou18194766xx-comfyui_EncryptPreview/web/loader.mjs

import { MAX_PACKET_BYTES } from "./codec.mjs";

export const abortError = () => new DOMException("Preview cancelled", "AbortError");

export function releaseImage(record) {
  if (!record) return;
  record.image.removeAttribute("src");
  URL.revokeObjectURL(record.url);
}

// Shared slots bound HTTP requests, workers AND image decodes across nodes.
export class PreviewLoader {
  constructor(fetchApi) {
    this.fetchApi = fetchApi;
    this.queue = [];
    this.slots = Array.from({ length: 2 }, () => ({ busy: false, worker: null }));
    this.users = 0;
  }

  acquire() { this.users++; }

  release() {
    if (this.users > 0) this.users--;
    if (!this.users) {
      for (const slot of this.slots) {
        if (!slot.busy) { slot.worker?.terminate(); slot.worker = null; }
      }
    }
  }

  load(batchId, index, kind, signal) {
    return new Promise((resolve, reject) => {
      if (signal.aborted) { reject(abortError()); return; }
      const job = { batchId, index, kind, signal, resolve, reject };
      job.cancel = () => {
        const position = this.queue.indexOf(job);
        if (position !== -1) this.queue.splice(position, 1);
        reject(abortError());
      };
      signal.addEventListener("abort", job.cancel, { once: true });
      if (kind === "original") this.queue.unshift(job);
      else this.queue.push(job);
      this.pump();
    });
  }

  pump() {
    for (const slot of this.slots) {
      if (slot.busy || !this.queue.length) continue;
      const job = this.queue.shift();
      job.signal.removeEventListener("abort", job.cancel);
      slot.busy = true;
      this.run(slot, job).then(job.resolve, job.reject).finally(() => {
        slot.busy = false;
        if (!this.users) { slot.worker?.terminate(); slot.worker = null; }
        this.pump();
      });
    }
  }

  async run(slot, job) {
    const controller = new AbortController();
    const cancel = () => controller.abort();
    job.signal.addEventListener("abort", cancel, { once: true });
    if (job.signal.aborted) controller.abort();
    const timer = setTimeout(cancel, 30_000);
    let record;
    try {
      const response = await this.fetchApi(`/encrypt-preview/v1/${job.batchId}/${job.index}/${job.kind}`, {
        signal: controller.signal, cache: "no-store",
      });
      if (response.status === 410 || response.status === 404) throw new Error("预览已过期，请重新运行节点");
      if (!response.ok) throw new Error(`读取预览失败 (${response.status})`);
      if (Number(response.headers.get("Content-Length")) > MAX_PACKET_BYTES) throw new Error("预览图片过大");
      const buffer = await response.arrayBuffer();
      if (controller.signal.aborted) throw abortError();
      if (buffer.byteLength > MAX_PACKET_BYTES) throw new Error("预览图片过大");
      const decoded = await this.restore(slot, buffer, controller.signal);
      if (controller.signal.aborted) throw abortError();
      const image = new Image();
      image.decoding = "async";
      record = { image, url: URL.createObjectURL(decoded.blob), blob: decoded.blob };
      const stopDecode = () => image.removeAttribute("src");
      controller.signal.addEventListener("abort", stopDecode, { once: true });
      try {
        image.src = record.url;
        // The native widget needs HTMLImageElement. Decode once, asynchronously.
        await image.decode();
      } finally {
        controller.signal.removeEventListener("abort", stopDecode);
      }
      if (controller.signal.aborted) throw abortError();
      if (image.naturalWidth !== decoded.width || image.naturalHeight !== decoded.height) throw new Error("预览图片尺寸不匹配");
      return record;
    } catch (error) {
      releaseImage(record);
      if (controller.signal.aborted) {
        if (job.signal.aborted) throw abortError();
        throw new Error("读取预览超时，请重试");
      }
      throw error;
    } finally {
      clearTimeout(timer);
      job.signal.removeEventListener("abort", cancel);
    }
  }

  restore(slot, buffer, signal) {
    return new Promise((resolve, reject) => {
      let worker;
      let settled = false;
      const finish = (error, result) => {
        if (settled) return;
        settled = true;
        signal.removeEventListener("abort", stop);
        if (worker) { worker.onmessage = null; worker.onerror = null; worker.onmessageerror = null; }
        if (error) {
          slot.worker?.terminate(); slot.worker = null;
          reject(error);
        } else resolve(result);
      };
      const stop = () => finish(abortError());
      try {
        if (signal.aborted) { stop(); return; }
        worker = slot.worker ??= new Worker(new URL("./decode.worker.mjs", import.meta.url), { type: "module" });
        worker.onmessage = ({ data }) => {
          if (data?.error) finish(new Error(data.error));
          else if (!(data?.blob instanceof Blob)) finish(new Error("预览 Worker 返回了无效数据"));
          else finish(null, data);
        };
        worker.onerror = event => { event.preventDefault(); finish(new Error(event.message || "预览 Worker 启动失败")); };
        worker.onmessageerror = () => finish(new Error("预览 Worker 返回了无效数据"));
        signal.addEventListener("abort", stop, { once: true });
        worker.postMessage({ buffer }, [buffer]);
      } catch (error) { finish(error); }
    });
  }
}