AgentHub/server/lib/user-storage.js

/**
 * @file user-storage.js
 * @description How much disk each server-mode user's HOME folder takes
 * (`homes.homeOf`: their own Linux home under S3, their folder under
 * AGENTHUB_USERS_ROOT otherwise) — the Users page's Storage column
 * (`GET /api/users/storage`). A home is walked in a worker thread, never on
 * the event loop (a member decides how many files their home holds), one home
 * at a time, under a time limit (`AGENTHUB_STORAGE_TIMEOUT_MS`, default 120 s
 * — past it the number is a lower bound, `partial`). The walk is `du`-like and
 * read-only: `lstat` only (no file is opened, a symlink is never followed),
 * stays on the home's file system, counts allocated blocks (a sparse file
 * costs what it really uses) and a hard-linked inode once. Results are cached
 * in memory for `FRESH_MS`; a read returns the cache at once and measures the
 * stale homes in the background. A user whose "home" is only the host
 * account's own (an admin with no folder of their own) is not measured.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

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

const FRESH_MS = 15 * 60_000;

/** Walk `root`: `{ bytes, files, dirs, partial }`. Runs in the worker. */
function walk(root, deadline) {
  const out = { bytes: 0, files: 0, dirs: 0, partial: false };
  let top;
  try {
    top = fs.lstatSync(root, { bigint: false });
  } catch {
    return { ...out, partial: true };
  }
  if (!top.isDirectory()) return out;
  const seen = new Set(); // dev:ino of hard-linked files already counted
  const stack = [root];
  while (stack.length) {
    if (Date.now() > deadline) {
      out.partial = true;
      break;
    }
    const dir = stack.pop();
    let names;
    try {
      names = fs.readdirSync(dir);
    } catch {
      out.partial = true; // unreadable (no CAP_DAC_READ_SEARCH, a race): a lower bound
      continue;
    }
    for (const name of names) {
      const p = path.join(dir, name);
      let st;
      try {
        st = fs.lstatSync(p);
      } catch {
        continue; // gone meanwhile
      }
      if (st.nlink > 1 && !st.isDirectory()) {
        const key = `${st.dev}:${st.ino}`;
        if (seen.has(key)) continue;
        seen.add(key);
      }
      out.bytes += st.blocks * 512;
      if (st.isDirectory()) {
        out.dirs += 1;
        if (st.dev === top.dev) stack.push(p); // one file system, like du -x
      } else {
        out.files += 1;
      }
    }
  }
  out.bytes += top.blocks * 512;
  return out;
}

if (!isMainThread && workerData?.agenthubStorageWorker) {
  const { dir, timeoutMs } = workerData;
  parentPort.postMessage(walk(dir, Date.now() + timeoutMs));
}

function timeoutMs() {
  return Number(process.env.AGENTHUB_STORAGE_TIMEOUT_MS) || 120_000;
}

/** Measure one folder in a fresh worker; never rejects. */
function measure(dir) {
  return new Promise((resolve) => {
    let w;
    const fail = () => resolve({ bytes: 0, files: 0, dirs: 0, partial: true });
    try {
      w = new Worker(__filename, {
        workerData: { agenthubStorageWorker: true, dir, timeoutMs: timeoutMs() },
      });
    } catch {
      return fail();
    }
    // The walk stops itself at the deadline; this only catches a worker stuck
    // inside one syscall (a hung network mount).
    const kill = setTimeout(() => {
      w.terminate().catch(() => {});
      fail();
    }, timeoutMs() + 30_000);
    kill.unref?.();
    w.once("message", (r) => {
      clearTimeout(kill);
      resolve(r);
      w.terminate().catch(() => {});
    });
    w.once("error", () => {
      clearTimeout(kill);
      fail();
    });
  });
}

const cache = new Map(); // dir → { bytes, files, dirs, partial, measuredAt }
const queue = []; // dirs waiting to be measured
let running = null; // dir being measured

function pump() {
  if (running || !queue.length) return;
  running = queue.shift();
  const dir = running;
  measure(dir).then((r) => {
    cache.set(dir, { ...r, measuredAt: new Date().toISOString() });
    running = null;
    pump();
  });
}

function enqueue(dir) {
  if (running === dir || queue.includes(dir)) return;
  queue.push(dir);
  pump();
}

/**
 * The storage of each of `users` (rows from `auth.listUsers()`):
 * `{ storage: { [id]: { bytes, files, partial, measuredAt } | null }, measuring }`.
 * `null` = nothing to measure (no home of their own) or not measured yet —
 * `measuring` says whether more is coming. `refresh` re-measures every home.
 */
function storageFor(users, { refresh = false } = {}) {
  const os = require("node:os");
  const { homeOf } = require("./homes");
  const hostHome = path.resolve(os.homedir());
  const storage = {};
  let measuring = false;
  for (const u of users) {
    let dir = null;
    try {
      dir = homeOf(u);
    } catch {
      dir = null;
    }
    if (!dir || path.resolve(dir) === hostHome) {
      storage[u.id] = null;
      continue;
    }
    dir = path.resolve(dir);
    const hit = cache.get(dir);
    const stale = !hit || refresh || Date.now() - Date.parse(hit.measuredAt) > FRESH_MS;
    if (stale) enqueue(dir);
    if (running === dir || queue.includes(dir)) measuring = true;
    storage[u.id] = hit
      ? { bytes: hit.bytes, files: hit.files, partial: hit.partial, measuredAt: hit.measuredAt }
      : null;
  }
  return { storage, measuring };
}

module.exports = { storageFor, walk, _cache: cache };