cpu-worker.mjs

// Worker harness for the tbdocs build pipeline.  Phase 15: SAB-based task
// metadata (no JS-side taskMeta array), priority-aware claiming, per-chunk
// flush via FIFO _pendingFlush queue, generic payload SAB.
// See PLAN-sab-pull-scheduler.md §Phase 15.

import { promises as fsP } from "node:fs";
import path from "node:path";
import { parentPort, workerData } from "node:worker_threads";
import { compileLightScss, compileDarkScss } from "./scss.mjs";
import { regenerateDot } from "./dot.mjs";
import { captureBuildInfo } from "./build-info.mjs";
import { helpArchiveStep } from "./help-archive-step.mjs";
import { writePdf } from "./pdf.mjs";
import { unpackBookPayload } from "./book-payload.mjs";
import { REPO_ROOT } from "../lib/repo-paths.mjs";

import { createMarkdownIt, renderPhase } from "./render.mjs";
import { renderFullNav, templatePhase, withPartialNav } from "./template.mjs";
import { unpackShared } from "./sab-broadcast.mjs";
import { deriveSearchEntries } from "./search.mjs";
import { computeChunkSeo } from "./seo.mjs";
import { pictureSvgOutputs, pngPictures } from "./theme-pictures.mjs";
import { deriveOfflinePage, isWebsiteOnlyLink, websiteOf } from "./offline-rewrite.mjs";
import { normalizeBaseurl } from "./url.mjs";

import {
  createViews,
  scanAndClaim,
  onTaskDone,
  readTaskMeta,
  HANDLERS,
  READY,
  CLAIMED,
  DONE,
  FAILED,
  F_ON_DEMAND,
  F_RUN_ON_MAIN,
  F_RUN_WHEN_IDLE,
  F_UNIQUE_PER_WORKER,
  MAX_LANES,
} from "./sab-scheduler.mjs";

if (workerData?.spawnTime) parentPort.postMessage({ coldBoot: { start: workerData.spawnTime, end: Date.now() } });

const myLane = workerData?.lane ?? 0;

// ── Mutable state set by init / dynamicData messages ────────────────────────

let views = null; // Int32Array views into the scheduling SAB
let ctx = null; // { srcRoot, destRoot, opts, workerCount }

let _payloadSAB = null; // SharedArrayBuffer with packed per-task payloads
let _sharedSAB = null; // SharedArrayBuffer with packed shared payload
let _shared = null; // _sharedSAB decoded, once per lane

// What dispatch broadcast. A task that reads it runs after dispatch, so the
// SAB is here by then; renderEnvInit, which can be claimed earlier, waits
// for it first.
function sharedData() {
  if (!_sharedSAB) throw new Error("the shared payload was read before dispatch broadcast it");
  _shared ??= unpackShared(_sharedSAB);
  return _shared;
}

// ── Handler table ───────────────────────────────────────────────────────────

const handlers = {
  async warmInit() {
    const { initHighlighter } = await import("./highlight.mjs");
    await initHighlighter();
    return {};
  },

  async renderEnvInit() {
    while (!_sharedSAB) {
      await new Promise((resolve) => setImmediate(resolve));
    }

    const {
      siteData,
      initData,
      linkTablesData,
      staticFilesArr,
      baseurl,
      buildInfo,
      sitePathsArr,
      onlineSitePathsArr,
      skipOffline,
      svgContentsMap,
      checkTrees,
      vendoredVideosObj,
      vendoredImagesObj,
      pictureSizes,
      counts,
      downloads,
    } = sharedData();

    const { initHighlighter } = await import("./highlight.mjs");
    const highlighter = await initHighlighter();
    const linkTables = reconstructLinkTables(linkTablesData);
    const staticFiles = new Set(staticFilesArr);
    const svgContents = new Map(Object.entries(svgContentsMap ?? {}));
    const vendoredVideos = new Map(Object.entries(vendoredVideosObj ?? {}));
    const vendoredImages = new Map(Object.entries(vendoredImagesObj ?? {}));
    const markdown = createMarkdownIt({
      highlighter,
      linkTables,
      baseurl,
      staticFiles,
      pictureSizes,
      svgContents,
      vendoredVideos,
      vendoredImages,
      counts,
      downloads,
    });
    const site = { ...siteData, markdown, buildInfo };

    let offlineBase = null;
    if (!skipOffline) {
      offlineBase = {
        sitePaths: new Set(sitePathsArr),
        website: websiteOf(siteData.config, new Set(onlineSitePathsArr)),
        baseurl: normalizeBaseurl(baseurl),
      };
    }

    // htmlparser2 costs ~23 ms to import, and a build without --check
    // must not pay it on sixteen lanes. Hence the dynamic import here
    // rather than a static one at module scope.
    if (checkTrees) {
      const { checkChunk, treeIndexFor, normalizeBasePath: normBase, TREES } = await import("./check.mjs");
      _checkChunk = checkChunk;
      _checkEnv = {};
      for (const [which, { rels, baseurl: bu }] of Object.entries(checkTrees)) {
        const root = ctx.destRoot + TREES[which].suffix;
        _checkEnv[which] = {
          root,
          tree: TREES[which],
          basePath: normBase(bu),
          index: treeIndexFor(root, rels),
          ...(which === "offline" && offlineBase ? { forbidExempt: (url) => isWebsiteOnlyLink(url, offlineBase) } : {}),
        };
      }
    }

    // An output that shows the pictures' PNGs (picture_svg) gets the renderer's HTML with
    // each picture's SVG named back as its PNG; the others get it as rendered.
    const svgOutputs = pictureSvgOutputs(siteData.config);
    const asPng = (html) => pngPictures(html, staticFiles, baseurl || "");
    _renderEnv = { site, initData, offlineBase, svgOutputs, asPng };
    return {};
  },

  async flush() {
    const items = _pendingFlush.shift() ?? [];
    let written = 0,
      offlineWritten = 0,
      offlineMisses = 0;

    // The check rides along here rather than becoming its own task
    // because this is the one moment both trees' final HTML is already
    // decoded and in this lane's memory. A separate task would have to
    // be handed the strings back, which is the payload cost the whole
    // design exists to avoid.
    let writes = null;
    if (!ctx.opts.dryRun) {
      let next = 0;
      const limit = Math.min(64, items.length || 1);
      // Building the array issues the first `limit` writes before
      // returning -- each async arrow runs synchronously up to its
      // first await -- so constructing it ahead of the check puts the
      // I/O in libuv's threadpool while this thread parses. Measured as
      // a wash, because the lanes are CPU-bound rather than waiting on
      // disk; it is kept because it is the right order and costs
      // nothing.
      const workers = Array.from({ length: limit }, async () => {
        while (next < items.length) {
          const p = items[next++];
          await fsP.writeFile(path.join(ctx.destRoot, p.destPath), p.html, "utf8");
          written++;
          if (p.offlineHtml !== undefined) {
            await fsP.writeFile(path.join(ctx.destRoot + "-offline", p.destPath), p.offlineHtml, "utf8");
            offlineWritten++;
          }
          offlineMisses += p.offlineMisses ?? 0;
        }
      });
      writes = Promise.all(workers);
    }

    const check = _checkEnv ? await runChunkCheck(items) : null;
    if (writes) await writes;

    return { written, offlineWritten, offlineMisses, check };
  },

  async scssLight() {
    const scssLightResult = await compileLightScss(ctx.srcRoot);
    return { scssLightResult };
  },

  async scssDark() {
    const scssDarkResult = await compileDarkScss(ctx.srcRoot);
    return { scssDarkResult };
  },

  async dot() {
    const dotStats = await regenerateDot(ctx.srcRoot);
    return { dotStats };
  },

  async buildInfo() {
    const buildInfo = await captureBuildInfo();
    return { buildInfo };
  },

  // On a worker because it is the main thread's longest task after the
  // pages are written, and it needs nothing from the main thread that
  // dispatch has not already broadcast: it reads the offline tree from disk.
  async helpArchive() {
    const { skipOffline } = sharedData();
    return helpArchiveStep({
      src: path.relative(REPO_ROOT, ctx.srcRoot).replaceAll(path.sep, "/"),
      offlineRoot: skipOffline ? null : `${ctx.destRoot}-offline`,
      serve: !!ctx.opts.serve,
      dryRun: !!ctx.opts.dryRun,
      disabled: !!ctx.opts.skipHelpArchive,
    });
  },

  // The PDF tree, from what bookPayload packed on the main thread. The book
  // itself stays on this lane for checkBook rather than crossing back.
  async writePdf() {
    const book = unpackBookPayload(await taskData("book"));
    _bookCheck = null;
    if (!book) return null;
    const { checkBook, ...counters } = await writePdf(book.pages, book.staticFiles, book.site, ctx.destRoot, {
      tolerateMissingImages: ctx.opts.tolerateMissingImages,
      highlightCss: book.highlightCss,
      check: book.check,
    });
    _bookCheck = checkBook ?? null;
    return counters;
  },

  // The book is one 6.5 MB document with almost entirely internal
  // fragments, written by a different task from a different string.
  // Keeping it its own pass is cheaper than special-casing the chunk
  // path for a single file.
  async checkBook() {
    if (!_bookCheck) return null;
    const { html, rels } = _bookCheck;
    _bookCheck = null;
    const { checkChunk, joinChunks, treeIndexFor, TREES } = await import("./check.mjs");
    const root = ctx.destRoot + TREES.pdf.suffix;
    const env = { root, tree: TREES.pdf, basePath: "", index: treeIndexFor(root, rels) };
    let chunk;
    try {
      chunk = checkChunk([{ destPath: "book.html", html }], env);
    } catch (err) {
      chunk = { error: `book check failed: ${err.message}` };
    }
    return joinChunks([chunk], { root, tree: TREES.pdf, relFiles: ["book.html"] });
  },

  async render(taskIdx) {
    const offset = Atomics.load(views.payloadOffset, taskIdx);
    const length = Atomics.load(views.payloadLength, taskIdx);
    const chunk = JSON.parse(new TextDecoder().decode(new Uint8Array(_payloadSAB, offset, length)));

    const env = _renderEnv;

    await renderPhase(chunk, env.site);
    computeChunkSeo(chunk, env.site.seoSiteTitle, env.site.config, env.site.markdown);
    await templatePhase(chunk, env.site, env.initData);

    if (env.offlineBase) {
      const offlineState = {
        ...env.offlineBase,
        caches: { rawResolution: new Map(), seg: new Map(), result: new Map() },
      };

      // An offline page holds only its part of the nav tree (withPartialNav).
      env.fullNav ??= renderFullNav(env.site);
      const writable = chunk.filter((p) => p.html !== undefined);
      for (const p of writable) {
        const source = env.svgOutputs.offline ? p.html : env.asPng(p.html);
        const offlinePage = { ...p, html: withPartialNav(source, p, env.site, env.fullNav) };
        const { html, misses } = deriveOfflinePage(offlinePage, offlineState);
        p.offlineHtml = html;
        p.offlineMisses = misses;
      }
    }

    // The online page, once the offline one is taken from the HTML as rendered.
    if (!env.svgOutputs.online) for (const p of chunk) if (p.html !== undefined) p.html = env.asPng(p.html);

    // Stash writable pages for the matching flush:i (FIFO; one batch per
    // render:i, drained by exactly one flush:i on the same worker).
    const batch = [];
    for (const p of chunk) {
      if (p.html !== undefined) {
        batch.push({
          destPath: p.destPath,
          html: p.html,
          offlineHtml: p.offlineHtml,
          offlineMisses: p.offlineMisses,
        });
      }
    }
    _pendingFlush.push(batch);

    // Per-chunk search entries; consolidated on main during searchData.
    // Drop `sourcePage` (workers hold cloned page objects, not master
    // refs) and `i` (chunk-local indices are meaningless; main assigns
    // global indices during consolidation).
    const searchEntries = deriveSearchEntries(chunk, env.site).map((e) => ({
      doc: e.doc,
      title: e.title,
      content: e.content,
      url: e.url,
      relUrl: e.relUrl,
      index: e.index,
      index_also: e.index_also,
    }));

    return {
      pages: chunk.map((p) => ({
        destPath: p.destPath,
        renderedContent: p.renderedContent,
        offlineMisses: p.offlineMisses,
      })),
      searchEntries,
    };
  },
};

// Reverse handler table: HANDLERS maps name → integer, handlerById maps
// integer → function.  Built once at module load.
const handlerById = [];
for (const [name, id] of Object.entries(HANDLERS)) handlerById[id] = handlers[name];

let _renderEnv = null;
let _pendingFlush = [];
let _taskData = new Map(); // key → SAB a main-thread task broadcast (broadcastTaskData)
let _bookCheck = null; // writePdf's book for checkBook, which is pinned to the same lane

// The data a main-thread predecessor broadcast for this task. The message
// is queued before the predecessor is marked done, but this lane may still
// be inside the pull loop that claimed the task, so yield until it lands.
async function taskData(key) {
  while (!_taskData.has(key)) await new Promise((resolve) => setImmediate(resolve));
  return _taskData.get(key);
}
let _checkEnv = null; // set by renderEnvInit when --check is on
let _checkChunk = null; // builder/check.mjs's checkChunk, imported with it

// One chunk of pages, checked against every tree it was written to.
//
// A failing check must never take the build's output down with it, so
// everything here is inside a try: a broken link still produces a valid
// site you want on disk to inspect. The error comes back as a finding
// and checkReport decides what to do with it.
async function runChunkCheck(items) {
  const out = {};
  for (const [which, env] of Object.entries(_checkEnv)) {
    // Asserted rather than filtered. The offline pass runs for every
    // page in the lane, so a page without offlineHtml means the rewrite
    // did not happen -- and dropping it here would shrink the chunk
    // quietly, checking fewer pages and still reporting a pass.
    const docs =
      which === "offline"
        ? items.map((p) => {
            if (p.offlineHtml === undefined) {
              throw new Error(
                `${p.destPath} reached the offline check with no offlineHtml; ` +
                  `the chunk would have covered fewer pages than the lane holds`,
              );
            }
            return { destPath: p.destPath, html: p.offlineHtml };
          })
        : items.map((p) => ({ destPath: p.destPath, html: p.html }));
    try {
      out[which] = _checkChunk(docs, env);
    } catch (err) {
      out[which] = { error: `${which} chunk check failed: ${err.message}` };
    }
  }
  return out;
}

// ── Message handler (init, dynamicData, taskData) ───────────────────────────

parentPort.on("message", (msg) => {
  if (msg.init) {
    views = createViews(msg.sab);
    ctx = msg.ctx;
    _payloadSAB = null;
    _sharedSAB = null;
    _shared = null;
    _renderEnv = null;
    _checkEnv = null;
    _pendingFlush = [];
    _taskData = new Map();
    _bookCheck = null;
    pullLoop();
    return;
  }
  if (msg.dynamicData) {
    _payloadSAB = msg.payloadSAB;
    _sharedSAB = msg.sharedSAB;
    _shared = null;
    return;
  }
  if (msg.taskData) {
    _taskData.set(msg.taskData, msg.sab);
    return;
  }
});

// ── Idle-task scan (speculative warmup) ─────────────────────────────────────

function findIdleTask(views, lane) {
  const count = Atomics.load(views.taskCount, 0);
  let bestIdx = -1;
  let bestPri = Infinity;
  for (let i = 0; i < count; i++) {
    if (!(Atomics.load(views.flags, i) & F_RUN_WHEN_IDLE)) continue;
    const meta = readTaskMeta(views, i);
    const pri = meta.idlePriority;
    if (pri >= bestPri) continue;
    if (Atomics.load(views.flags, i) & F_UNIQUE_PER_WORKER) {
      if (Atomics.load(views.perWorkerDone, i * MAX_LANES + lane) !== 0) continue;
      let skip = false;
      for (const predIdx of meta.expectedDeps) {
        if (Atomics.load(views.status, predIdx) !== DONE) {
          skip = true;
          break;
        }
      }
      if (skip) continue;
      for (const depIdx of meta.perWorkerDeps) {
        if (Atomics.load(views.perWorkerDone, depIdx * MAX_LANES + lane) === 0) {
          skip = true;
          break;
        }
      }
      if (skip) continue;
      bestIdx = i;
      bestPri = pri;
    } else {
      if (Atomics.load(views.status, i) !== READY) continue;
      bestIdx = i;
      bestPri = pri;
    }
  }
  // For non-unique_per_worker tasks, CAS-claim at the end.
  if (bestIdx !== -1 && !(Atomics.load(views.flags, bestIdx) & F_UNIQUE_PER_WORKER)) {
    if (Atomics.compareExchange(views.status, bestIdx, READY, CLAIMED) !== READY) return -1;
  }
  return bestIdx;
}

// ── Pull loop ───────────────────────────────────────────────────────────────

// Runs a per-worker task on this lane, marks it done for the lane, and
// reports its timing and output. Returns false once a throw has been
// reported, and the caller then ends the pull loop.
async function runPerWorkerTask(taskIdx, meta) {
  const t0 = Date.now();
  let output;
  try {
    output = await handlerById[meta.handlerIdx]();
  } catch (err) {
    parentPort.postMessage({ taskFailed: taskIdx, message: err.message, stack: err.stack });
    return false;
  }
  const t1 = Date.now();
  Atomics.store(views.perWorkerDone, taskIdx * MAX_LANES + myLane, 1);
  parentPort.postMessage({
    perWorkerTiming: true,
    taskIdx,
    timing: { start: t0, end: t1 },
    lane: myLane,
    output,
  });
  return true;
}

async function pullLoop() {
  while (true) {
    if (Atomics.load(views.buildDone, 0) !== 0) return;

    let taskIdx = scanAndClaim(views, myLane);

    if (taskIdx === -1) {
      // Speculative: run idle-eligible tasks before sleeping.
      const idleTask = findIdleTask(views, myLane);
      if (idleTask !== -1) {
        if (!(await runPerWorkerTask(idleTask, readTaskMeta(views, idleTask)))) return;
        continue;
      }

      const gen = Atomics.load(views.notify, 0);
      // Double-check after reading gen (race: a task may have become
      // READY between the failed scan and this load).
      taskIdx = scanAndClaim(views, myLane);
      if (taskIdx === -1) {
        Atomics.wait(views.notify, 0, gen, 50);
        continue;
      }
    }

    // ── Per-worker deps (unique_per_worker) ──
    const meta = readTaskMeta(views, taskIdx);
    let unsatisfied = null;
    for (const depIdx of meta.perWorkerDeps) {
      if (Atomics.load(views.perWorkerDone, depIdx * MAX_LANES + myLane) === 0) {
        unsatisfied = depIdx;
        break;
      }
    }

    if (unsatisfied !== null) {
      const depFlags = Atomics.load(views.flags, unsatisfied);

      if (depFlags & F_ON_DEMAND && !(depFlags & F_RUN_ON_MAIN)) {
        const depMeta = readTaskMeta(views, unsatisfied);

        // Check the dep's own perWorkerDeps (e.g. renderEnvInit → warmInit).
        let nestedUnsatisfied = null;
        for (const nestedIdx of depMeta.perWorkerDeps) {
          if (Atomics.load(views.perWorkerDone, nestedIdx * MAX_LANES + myLane) === 0) {
            nestedUnsatisfied = nestedIdx;
            break;
          }
        }

        if (nestedUnsatisfied !== null) {
          const nestedFlags = Atomics.load(views.flags, nestedUnsatisfied);
          if (nestedFlags & F_ON_DEMAND && !(nestedFlags & F_RUN_ON_MAIN)) {
            Atomics.store(views.status, taskIdx, READY);
            Atomics.add(views.notify, 0, 1);
            Atomics.notify(views.notify, 0, 1);

            if (!(await runPerWorkerTask(nestedUnsatisfied, readTaskMeta(views, nestedUnsatisfied)))) return;
            continue;
          }
          Atomics.store(views.status, taskIdx, READY);
          Atomics.add(views.notify, 0, 1);
          Atomics.notify(views.notify, 0, 1);
          continue;
        }

        // Check preconditions (expected predecessors on the dep).
        let precondFailed = false;
        for (const predIdx of depMeta.expectedDeps) {
          if (Atomics.load(views.status, predIdx) !== DONE) {
            precondFailed = true;
            break;
          }
        }
        if (precondFailed) {
          Atomics.store(views.status, taskIdx, READY);
          Atomics.add(views.notify, 0, 1);
          Atomics.notify(views.notify, 0, 1);
          continue;
        }

        // All dep's deps satisfied. Release original task, execute the dep.
        Atomics.store(views.status, taskIdx, READY);
        Atomics.add(views.notify, 0, 1);
        Atomics.notify(views.notify, 0, 1);

        if (!(await runPerWorkerTask(unsatisfied, depMeta))) return;
        continue;
      }

      // Other unsatisfied dep types: release and re-scan.
      Atomics.store(views.status, taskIdx, READY);
      Atomics.add(views.notify, 0, 1);
      Atomics.notify(views.notify, 0, 1);
      continue;
    }

    // ── Execute task ──
    const handler = handlerById[meta.handlerIdx];
    if (!handler) {
      parentPort.postMessage({ taskFailed: taskIdx, message: `unknown handlerIdx: ${meta.handlerIdx}`, stack: "" });
      return;
    }

    const t0 = Date.now();
    let result;
    try {
      result = await handler(taskIdx);
    } catch (err) {
      parentPort.postMessage({ taskFailed: taskIdx, message: err.message, stack: err.stack });
      Atomics.store(views.status, taskIdx, FAILED);
      return;
    }
    const t1 = Date.now();

    // Post the output BEFORE the SAB update, so the result is at least
    // QUEUED on the main thread before any successor's dep count drops.
    //
    // That is an ordering of the two operations, NOT a guarantee that
    // the merge has happened: the message still has to be delivered and
    // the matching submit() still has to run. A successor can reach a
    // dep count of zero with the result sitting unread in the main
    // thread's queue -- which is exactly how renderJoin lost pages from
    // the search index. What closes that window is the barrier's
    // `expected` list (see registerBarrier in tbdocs.mjs), not this
    // ordering.
    parentPort.postMessage({
      done: taskIdx,
      output: result,
      timing: { start: t0, end: t1 },
      lane: myLane,
    });

    const { readyCount, wakeMain } = onTaskDone(views, taskIdx, myLane);
    if (readyCount > 0) {
      Atomics.add(views.notify, 0, 1);
      Atomics.notify(views.notify, 0, readyCount);
    }
    if (wakeMain) {
      parentPort.postMessage({ mainTaskReady: true });
    }
  }
}

function reconstructLinkTables({ byPath, byUrl, byRedirect, byShortcut }) {
  const make = (pairs) => new Map(pairs.map(([k, pl]) => [k, { permalink: pl }]));
  return { byPath: make(byPath), byUrl: make(byUrl), byRedirect: make(byRedirect), byShortcut: make(byShortcut) };
}