Merge issue 776 follow-up: shared worker coordination on Bun 1.3, per-kind pool keys, startup watchdog at dispatch

Co-authored-by: Alfonso <alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/plugin/src/hooks/magic-context/raw-index-worker-pool.test.ts

@@ -0,0 +1,114 @@
+import { expect, test } from "bun:test";
+import { join } from "node:path";
+import { initializeDatabase } from "../../features/magic-context/storage-db";
+import { resetReusableWorkersForTests } from "../../shared/reusable-worker";
+import { Database } from "../../shared/sqlite";
+import { createTestTempDir } from "../../shared/test-temp-dir";
+import { resetRawSessionOrdinalWarmupsForTest, runSessionWarmupWorker } from "./raw-ordinal-warmup";
+
+test("queued indexing does not spend its startup deadline before dispatch", async () => {
+    const { dir, cleanup } = createTestTempDir("raw-index-pool-");
+    const path = join(dir, "fixture.db");
+    const db = new Database(path);
+    const coordination = new Int32Array(new SharedArrayBuffer(3 * Int32Array.BYTES_PER_ELEMENT));
+    Atomics.store(coordination, 0, 1);
+    let release: ReturnType<typeof setTimeout> | undefined;
+    try {
+        initializeDatabase(db);
+        const first = runSessionWarmupWorker(
+            {
+                path,
+                sessionId: "index-pool",
+                harness: "opencode",
+                indexMessage: {
+                    id: "first",
+                    ordinal: 1,
+                    role: "user",
+                    parts: [{ type: "text", text: "first source" }],
+                },
+                indexCoordination: coordination.buffer,
+                indexOwner: 1,
+            },
+            { acceptMs: 5_000, resultMs: 5_000 },
+        );
+        const queued = runSessionWarmupWorker(
+            {
+                path,
+                sessionId: "index-pool",
+                harness: "opencode",
+                indexMessage: {
+                    id: "queued",
+                    ordinal: 2,
+                    role: "user",
+                    parts: [{ type: "text", text: "queued source" }],
+                },
+                indexCoordination: coordination.buffer,
+                indexOwner: 2,
+            },
+            { acceptMs: 200, resultMs: 5_000 },
+        );
+        release = setTimeout(() => {
+            Atomics.store(coordination, 0, 0);
+            Atomics.notify(coordination, 0);
+        }, 500);
+        const [firstResult, queuedResult] = await Promise.all([first, queued]);
+        expect(firstResult).toMatchObject({ inserted: true });
+        expect(queuedResult).toMatchObject({ inserted: true });
+        expect(
+            db.prepare("SELECT message_id FROM message_history_fts ORDER BY message_id").all(),
+        ).toEqual([{ message_id: "first" }, { message_id: "queued" }]);
+    } finally {
+        clearTimeout(release);
+        Atomics.store(coordination, 0, 0);
+        resetRawSessionOrdinalWarmupsForTest();
+        await resetReusableWorkersForTests();
+        db.close();
+        cleanup();
+    }
+});
+
+test("different raw job kinds cannot queue behind an index writer", async () => {
+    const { dir, cleanup } = createTestTempDir("raw-index-kinds-");
+    const path = join(dir, "fixture.db");
+    const db = new Database(path);
+    const coordination = new Int32Array(new SharedArrayBuffer(3 * Int32Array.BYTES_PER_ELEMENT));
+    Atomics.store(coordination, 0, 1);
+    try {
+        initializeDatabase(db);
+        let writerSettled = false;
+        const writer = runSessionWarmupWorker(
+            {
+                path,
+                sessionId: "index-kinds",
+                harness: "opencode",
+                indexMessage: {
+                    id: "writer",
+                    ordinal: 1,
+                    role: "user",
+                    parts: [{ type: "text", text: "writer source" }],
+                },
+                indexCoordination: coordination.buffer,
+                indexOwner: 1,
+            },
+            { acceptMs: 5_000, resultMs: 1_000 },
+        ).then((result) => {
+            writerSettled = true;
+            return result;
+        });
+        const backfill = await runSessionWarmupWorker(
+            { path, sessionId: "", harness: "opencode", ftsRowidBackfill: true },
+            { acceptMs: 5_000, resultMs: 5_000 },
+        );
+        expect(backfill).toMatchObject({ ftsRowidBackfill: true });
+        expect(writerSettled).toBe(false);
+        Atomics.store(coordination, 0, 0);
+        Atomics.notify(coordination, 0);
+        expect(await writer).toMatchObject({ inserted: true });
+    } finally {
+        Atomics.store(coordination, 0, 0);
+        resetRawSessionOrdinalWarmupsForTest();
+        await resetReusableWorkersForTests();
+        db.close();
+        cleanup();
+    }
+});

packages/plugin/src/hooks/magic-context/raw-ordinal-warmup.ts

@@ -195,6 +195,7 @@ export function runSessionWarmupWorker(
             clearTimeout(deadline);
             options.signal?.removeEventListener("abort", abort);
             if (worker) {
+                worker.off("online", startDeadline);
                 activeWorkers.delete(worker);
                 const stopped = shutdown?.close() ?? worker.terminate();
                 worker.unref();
@@ -215,6 +216,9 @@ export function runSessionWarmupWorker(
             );
             deadline.unref?.();
         };
+        const startDeadline = () => {
+            if (!settled) arm("not-accepted", options.acceptMs ?? ACCEPT_TIMEOUT_MS);
+        };
         const abort = () => finish({ failure: "stopped", detail: "stopped" });
         if (options.signal?.aborted) {
             abort();
@@ -228,7 +232,11 @@ export function runSessionWarmupWorker(
                           workerData: { ...input, openCancellation: openCancellation.buffer },
                       })
                     : reuseWorker(defaultEntry(), {
-                          workerData: { ...input, openCancellation: openCancellation.buffer },
+                          workerData: {
+                              ...input,
+                              job: rawWorkerJobKind(input),
+                              openCancellation: openCancellation.buffer,
+                          },
                       }),
             );
             shutdown = createWorkerShutdown(worker, "raw-ordinal-warmup", WARMUP_CLOSE_DEADLINE_MS);
@@ -239,7 +247,10 @@ export function runSessionWarmupWorker(
         // Never keep the host process alive for a warm-up.
         worker.unref();
         activeWorkers.set(worker, (failure) => finish({ failure, detail: "stopped" }));
-        arm("not-accepted", options.acceptMs ?? ACCEPT_TIMEOUT_MS);
+        // Pool wait time is not thread startup. Arm the startup watchdog only
+        // when this job is dispatched to its thread, not while it is queued.
+        if (options.entry) startDeadline();
+        else worker.once("online", startDeadline);
         worker.on("message", (reply: RawOrdinalWarmupReply) => {
             if (reply.kind === "accepted") arm("no-result", options.resultMs ?? RESULT_TIMEOUT_MS);
             else if (reply.kind === "scan") finish({ scan: reply.scan });
@@ -263,6 +274,15 @@ export function runSessionWarmupWorker(
     });
 }
 
+function rawWorkerJobKind(input: RawOrdinalWarmupInput): string {
+    if (input.indexMessage) return "index-message";
+    if (input.ftsRowidBackfill) return "fts-rowid-backfill";
+    if (input.searchArchive) return "search-archive";
+    if (input.exactOrdinalIds) return "exact-ordinal-ids";
+    if (input.tagRatio !== undefined) return input.tagState ? "tag-delta" : "tag-summary";
+    return "ordinal-scan";
+}
+
 /**
  * Prepare exact message numbers, excluding finished compaction summaries, without
  * scanning the session on the host thread. The worker records prefix counts

packages/plugin/src/shared/reusable-worker.fixture.ts

@@ -7,10 +7,23 @@ serveWorkerJobs<{
     late?: boolean;
     exitCode?: boolean;
     value?: string;
+    shared?: SharedArrayBuffer;
 }>(async (input, port) => {
     if (input.crash) throw new Error("injected worker crash");
     if (input.hang) return await new Promise<void>(() => {});
     if (input.exitCode) process.exitCode = 1;
+    if (input.shared) {
+        const state = new Int32Array(input.shared);
+        port.postMessage({ threadId, shared: Atomics.load(state, 0) });
+        await new Promise<void>((resolve) =>
+            port.once("message", () => {
+                port.postMessage({ threadId, shared: Atomics.add(state, 0, 1) });
+                resolve();
+            }),
+        );
+        port.close();
+        return;
+    }
     if (input.late)
         setTimeout(() => {
             port.postMessage({ value: "stale" });

packages/plugin/src/shared/reusable-worker.test.ts

@@ -1,9 +1,10 @@
 import { describe, expect, test } from "bun:test";
 import { spawnSync } from "node:child_process";
 import { once } from "node:events";
-import { writeFileSync } from "node:fs";
+import { readFileSync, writeFileSync } from "node:fs";
 import { join } from "node:path";
-import type { Worker } from "node:worker_threads";
+import { fileURLToPath } from "node:url";
+import { MessageChannel, type Worker } from "node:worker_threads";
 import { reuseWorker } from "./reusable-worker";
 import { withTestTempDir } from "./test-temp-dir";
 
@@ -107,6 +108,21 @@ describe("reusable job workers", () => {
         expect((await result(job)).value).toBe("submitted");
     });
 
+    test("job input keeps shared coordination live in both directions", async () => {
+        const state = new Int32Array(new SharedArrayBuffer(4));
+        Atomics.store(state, 0, 1);
+        const worker = start(`shared-${++key}`, { shared: state.buffer });
+        const first = once(worker, "message");
+        const exit = once(worker, "exit");
+        expect((await first)[0].shared).toBe(1);
+        const changed = once(worker, "message");
+        Atomics.store(state, 0, 7);
+        worker.postMessage({ read: true });
+        expect((await changed)[0].shared).toBe(7);
+        await exit;
+        expect(Atomics.load(state, 0)).toBe(8);
+    });
+
     test("a nonzero job exit fails the lease and replaces the thread", async () => {
         const path = `exit-code-${++key}`;
         const first = await result(start(path));
@@ -139,43 +155,99 @@ describe("reusable job workers", () => {
 
     test("a file ending with queued jobs cannot interrupt the next parallel file", () => {
         withTestTempDir("reusable-worker-files-", (dir) => {
+            const outcomes = join(dir, "outcomes.jsonl");
+            writeFileSync(outcomes, "");
             writeFileSync(
                 join(dir, "01-owner.test.ts"),
                 `import { test } from "bun:test";
+import { appendFileSync } from "node:fs";
 import { reuseWorker } from ${JSON.stringify(new URL("./reusable-worker.ts", import.meta.url).href)};
 test("owner finishes with queued work", async () => {
     const entry = new URL(${JSON.stringify(entry.href)});
     const active = reuseWorker(entry, { workerData: { path: "file-exit", hang: true } });
-    const queued = reuseWorker(entry, { workerData: { path: "file-exit", value: "queued" } });
     active.on("error", () => {});
-    queued.on("error", () => console.log("queued failure reached client"));
-    queued.on("exit", (code) => console.log("queued exit:", code));
+    for (let id = 0; id < 3; id++) {
+        const queued = reuseWorker(entry, { workerData: { path: "file-exit", value: "queued" } });
+        let replied = false;
+        let failed = false;
+        queued.on("message", () => { replied = true; });
+        queued.on("error", () => { failed = true; });
+        queued.on("exit", code => appendFileSync(${JSON.stringify(outcomes)}, JSON.stringify({
+            id, outcome: code === 0 && replied && !failed ? "completed" : "failed",
+        }) + "\\n"));
+    }
     await Promise.resolve();
 });`,
             );
-            // Top-level await lets the old file's native Worker exit handler run
-            // while the next file is importing.
             writeFileSync(
                 join(dir, "02-next.test.ts"),
                 `import { describe, expect, test } from "bun:test";
+import { readFileSync } from "node:fs";
 await new Promise(resolve => setTimeout(resolve, 50));
 describe("next file", () => {
-    test("still registers", () => expect(1).toBe(1));
+    test("still registers after every queued client settles once", () => {
+        const outcomes = readFileSync(${JSON.stringify(outcomes)}, "utf8").trim().split("\\n").filter(Boolean).map(line => JSON.parse(line));
+        expect(outcomes).toHaveLength(3);
+        for (let id = 0; id < 3; id++) {
+            const client = outcomes.filter(outcome => outcome.id === id);
+            expect(client).toHaveLength(1);
+            expect(["completed", "failed"]).toContain(client[0].outcome);
+        }
+    });
 });`,
             );
-            const child = spawnSync(process.execPath, ["test", "--parallel=1", dir], {
-                cwd: dir,
-                encoding: "utf8",
-                timeout: 5_000,
-                windowsHide: true,
-            });
+            // The real preload owns file teardown. Bun may discard a bare file's
+            // callbacks without notifying its workers' clients during VM teardown.
+            const child = spawnSync(
+                process.execPath,
+                [
+                    "test",
+                    "--parallel=1",
+                    "--preload",
+                    fileURLToPath(new URL("../../test-preload.ts", import.meta.url)),
+                    dir,
+                ],
+                { cwd: dir, encoding: "utf8", timeout: 5_000, windowsHide: true },
+            );
             const output = child.stdout + child.stderr;
             expect(output).toContain("2 pass");
             expect(output).toContain("0 fail");
             expect(output).not.toContain("Unhandled error");
-            expect(output).toContain("queued failure reached client");
-            expect(output).toContain("queued exit: 1");
             expect(child.status).toBe(0);
+            const settled = readFileSync(outcomes, "utf8").trim().split("\n").filter(Boolean);
+            expect(settled).toHaveLength(3);
         });
     }, 10_000);
+
+    test("replacement construction failure settles every queued client exactly once", async () => {
+        const path = `replacement-failure-${++key}`;
+        const { port1, port2 } = new MessageChannel();
+        try {
+            const crashed = reuseWorker(entry, {
+                workerData: { path, crash: true },
+                transferList: [port1],
+            });
+            crashed.on("error", () => {});
+            const clients = [start(path), start(path)].map((worker) => {
+                const errors: unknown[] = [];
+                const exits: number[] = [];
+                worker.on("error", (error) => errors.push(error));
+                const done = new Promise<void>((resolve) =>
+                    worker.on("exit", (code) => {
+                        exits.push(code);
+                        resolve();
+                    }),
+                );
+                return { errors, exits, done };
+            });
+            await Promise.all(clients.map((client) => client.done));
+            await new Promise((resolve) => setImmediate(resolve));
+            for (const client of clients) {
+                expect(client.errors).toHaveLength(1);
+                expect(client.exits).toEqual([1]);
+            }
+        } finally {
+            port2.close();
+        }
+    });
 });