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 (and, when the caller asks, reused for a short window), 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. The server's periodic WAL checkpoint
 * also runs here (`walCheckpoint`). 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
    ),
  // Not a read: the worker runs it on its own writable connection (below), so
  // the checkpoint's copy into the database file never blocks the event loop.
  walCheckpoint: (conn) => conn.pragma("wal_checkpoint(PASSIVE)", { simple: false }),
  analyticsEvents: (conn, a) =>
    require("./scoped-stats").eventAggregates(conn, a.scope, a.range, a.tzModifier),
};

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
const recent = new Map(); // job+args → { at, result } for callers passing maxAgeMs

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; with `maxAgeMs`, a result that young is reused too,
 * so N open dashboards cost one query per window instead of N. Args are part
 * of the key, so a result is only reused for the very same query.
 *
 * @param {keyof JOBS} job
 * @param {object} [args] structured-cloneable
 * @param {{maxAgeMs?: number}} [opts]
 * @returns {Promise<any>}
 */
function runRead(job, args = {}, { maxAgeMs = 0 } = {}) {
  if (!JOBS[job]) return Promise.reject(codedError(`unknown read job ${job}`, "EBADJOB"));
  const key = `${job}:${JSON.stringify(args)}`;
  const now = Date.now();
  const fresh = recent.get(key);
  if (fresh && now - fresh.at < maxAgeMs) return Promise.resolve(fresh.result);
  const hit = inflight.get(key);
  if (hit) return hit;
  const p = dispatch(job, args)
    .then((result) => {
      if (maxAgeMs > 0) {
        // ponytail: drop-expired-on-write keeps this bounded by the keys used
        // within one window; an LRU if a caller ever varies args per request.
        for (const [k, v] of recent) if (now - v.at > 60_000) recent.delete(k);
        recent.set(key, { at: Date.now(), result });
      }
      return result;
    })
    .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.
  const open = () => {
    try {
      const Database = require("better-sqlite3");
      return new Database(workerData.dbPath);
    } catch {
      const Database = require("../compat-sqlite");
      return new Database(workerData.dbPath);
    }
  };
  const conn = open();
  let writer = null; // walCheckpoint's connection, opened on first use
  conn.pragma("busy_timeout = 5000");
  conn.pragma("query_only = ON"); // this connection can never write
  conn.pragma("cache_size = -65536"); // same read tuning as server/db.js
  conn.pragma("mmap_size = 268435456");
  parentPort.on("message", ({ id, job, args }) => {
    try {
      const c = job === "walCheckpoint" ? (writer ??= open()) : conn;
      parentPort.postMessage({ id, result: JOBS[job](c, args) });
    } catch (err) {
      parentPort.postMessage({ id, error: { message: err.message, code: err.code } });
    }
  });
}

module.exports = { runRead };