AgentHub/server/lib/read-pool.js

/**
 * @file Runs the dashboard's heavy read-only aggregations in a worker thread
 * on its own SQLite connection, so a slow one never blocks the event loop that
 * serves every other page, hook and websocket. Identical requests in flight are
 * shared, a job over `AGENTHUB_READ_TIMEOUT_MS` (default 60 s) is answered with
 * `EREADTIMEOUT` and its worker replaced, and `AGENTHUB_READ_WORKER=0` — or a
 * worker that cannot start — runs the job inline as before. This file is also
 * the worker's entry point.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { Worker, isMainThread, workerData, parentPort } = require("node:worker_threads");

/** Read-only jobs by name: (connection, args) → plain data. */
const JOBS = {
  workflows: (conn, a) =>
    require("./workflow-stats").computeWorkflowData(
      conn,
      { status: a.status || null, tool: a.tool || null, range: a.range || null },
      a.sections || null
    ),
};

function timeoutMs() {
  return Number(process.env.AGENTHUB_READ_TIMEOUT_MS) || 60_000;
}

function codedError(message, code) {
  return Object.assign(new Error(message), { code });
}

// ponytail: one worker, so heavy reads queue behind each other (the event loop
// stays free either way); a small pool if concurrent heavy reads ever pile up.
let worker = null;
let seq = 0;
const pending = new Map(); // id → { resolve, reject, timer, w }
const inflight = new Map(); // job+args → Promise

function spawn(dbPath) {
  const w = new Worker(__filename, { workerData: { agenthubReadWorker: true, dbPath } });
  w.on("message", ({ id, result, error }) => {
    const p = pending.get(id);
    if (!p) return;
    pending.delete(id);
    clearTimeout(p.timer);
    if (error) p.reject(codedError(error.message, error.code));
    else p.resolve(result);
  });
  const fail = (err) => {
    if (worker === w) worker = null;
    for (const [id, p] of pending) {
      if (p.w !== w) continue;
      pending.delete(id);
      clearTimeout(p.timer);
      p.reject(err);
    }
  };
  w.on("error", fail);
  w.on("exit", (code) => fail(codedError(`read worker exited (${code})`, "EREADWORKER")));
  w.unref(); // after the listeners: adding a "message" listener refs it again
  return w;
}

function runInline(job, args) {
  const { db } = require("../db");
  return Promise.resolve().then(() => JOBS[job](db, args));
}

function dispatch(job, args) {
  if (process.env.AGENTHUB_READ_WORKER === "0") return runInline(job, args);
  try {
    if (!worker) worker = spawn(require("../db").DB_PATH);
  } catch {
    return runInline(job, args);
  }
  const w = worker;
  return new Promise((resolve, reject) => {
    const id = ++seq;
    const ms = timeoutMs();
    const timer = setTimeout(() => {
      // The stuck query keeps its thread until SQLite returns; the next job
      // gets a fresh worker. Jobs queued behind it are rejected by `exit`.
      pending.delete(id);
      if (worker === w) worker = null;
      w.terminate();
      reject(codedError(`read query took over ${ms} ms`, "EREADTIMEOUT"));
    }, ms);
    timer.unref?.();
    pending.set(id, { resolve, reject, timer, w });
    w.postMessage({ id, job, args });
  });
}

/**
 * Run a named read job off the event loop. Concurrent calls with the same job
 * and args share one run.
 *
 * @param {keyof JOBS} job
 * @param {object} [args] structured-cloneable
 * @returns {Promise<any>}
 */
function runRead(job, args = {}) {
  if (!JOBS[job]) return Promise.reject(codedError(`unknown read job ${job}`, "EBADJOB"));
  const key = `${job}:${JSON.stringify(args)}`;
  const hit = inflight.get(key);
  if (hit) return hit;
  const p = dispatch(job, args).finally(() => inflight.delete(key));
  inflight.set(key, p);
  return p;
}

if (!isMainThread && workerData?.agenthubReadWorker) {
  // Same driver order as server/db.js: better-sqlite3, else node:sqlite.
  let conn;
  try {
    const Database = require("better-sqlite3");
    conn = new Database(workerData.dbPath);
  } catch {
    const Database = require("../compat-sqlite");
    conn = new Database(workerData.dbPath);
  }
  conn.pragma("busy_timeout = 5000");
  conn.pragma("query_only = ON"); // this connection can never write
  parentPort.on("message", ({ id, job, args }) => {
    try {
      parentPort.postMessage({ id, result: JOBS[job](conn, args) });
    } catch (err) {
      parentPort.postMessage({ id, error: { message: err.message, code: err.code } });
    }
  });
}

module.exports = { runRead };