import { readFileSync } from "node:fs"; import { fileURLToPath } from "node:url"; import { type DB, migrateToLatest } from "@umt/db"; import { defineTask, type TaskLane } from "@umt/tasks"; import { Kysely, PostgresDialect } from "kysely"; import pg from "pg"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { z } from "zod"; import { ConfigSchema } from "../src/config.ts"; import { type Bootstrapped, bootstrap, buildTaskRegistry } from "../src/context.ts"; import { EventBus } from "../src/http/sse.ts"; import { SETTINGS_CHANGED_EVENT } from "../src/modules/auth-settings/provided.ts"; import { SERVER_MODULES } from "../src/modules/index.ts"; import type { ServerContext, ServerModule } from "../src/modules/types.ts"; import { type RunningWorker, startWorker } from "../src/worker/index.ts"; import { groupTasksByLane, LaneController, laneOverridesFromEnv } from "../src/worker/lanes.ts"; import { withSchemaLock } from "./auth-settings/harness.ts"; import { memoryLogger } from "./helpers.ts"; // WK-1 IO Worker 分道:归属表完整性(注册表 ↔ runbook)、控制器规则(并发热更新、让路、env 映射)、 // 真实 PG 上三条道互不阻塞与维护道让路。 const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); const fakeCtx = { role: "all" } as unknown as ServerContext; /** docs/runbook/dev.md「任务归属表」:| `name` | lane | 说明 | */ function runbookLaneTable(): Map { const md = readFileSync(fileURLToPath(new URL("../../../docs/runbook/dev.md", import.meta.url)), "utf8"); const section = md.slice(md.indexOf("#### 任务归属表")); const end = section.indexOf("\n\n", section.indexOf("| 任务 |")); const rows = section.slice(0, end).split("\n"); const out = new Map(); for (const r of rows) { const m = /^\| `([A-Za-z0-9_.-]+)` \| (\w+) \|/.exec(r); if (m) out.set(m[1]!, m[2]!); } return out; } describe("任务归属表", () => { const reg = buildTaskRegistry(SERVER_MODULES, fakeCtx); it("注册表里每个任务都显式声明了 lane", () => { const { missing } = groupTasksByLane(reg.all()); expect(missing, "这些任务没在 defineTask 里声明 lane").toEqual([]); }); it("runbook 的归属表与注册表逐项一致(新增任务要同时写进 docs/runbook/dev.md)", () => { const table = runbookLaneTable(); const actual = new Map(reg.all().map((d) => [d.name, d.lane as string])); expect(Object.fromEntries(table)).toEqual(Object.fromEntries(actual)); }); it("没声明 lane:runAgent/resolveDecision 兜底归 agent,其余归网络道并列入 missing", () => { const g = groupTasksByLane([{ name: "x" }, { name: "runAgent" }, { name: "f", lane: "file" }]); expect(g.byLane.network).toEqual(["x"]); expect(g.byLane.agent).toEqual(["runAgent"]); expect(g.missing).toEqual(["runAgent", "x"]); }); }); describe("LaneController 规则", () => { const { logger } = memoryLogger(); const lanes = { file: ["f"], network: ["n"], maintenance: ["m"] }; it("旧 WORKER_CONCURRENCY 映射到网络道;新变量优先", () => { expect(laneOverridesFromEnv({ WORKER_CONCURRENCY: 6 })).toEqual({ network: 6 }); expect( laneOverridesFromEnv({ WORKER_CONCURRENCY: 6, WORKER_NETWORK_CONCURRENCY: 3, WORKER_FILE_CONCURRENCY: 1 }), ).toEqual({ network: 3, file: 1 }); expect(laneOverridesFromEnv({})).toEqual({}); }); it("并发读设置、热更新;环境变量覆盖设置;槽位按上限开", async () => { let settings = { file: 2, network: 4, maintenance: 1 }; const c = new LaneController({ engine: {}, logger, lanes, readSettings: async () => settings, overrides: { maintenance: 20 }, }); await c.refresh(); expect([c.limit("file"), c.limit("network"), c.limit("maintenance")]).toEqual([2, 4, 20]); expect([c.slots("file"), c.slots("maintenance")]).toEqual([16, 20]); expect(c.source("maintenance")).toBe("env"); settings = { file: 5, network: 1, maintenance: 3 }; await c.refresh(); expect([c.limit("file"), c.limit("network"), c.limit("maintenance")]).toEqual([5, 1, 20]); }); it("维护道:其它道有活让路,空闲放行;连续让路超时放行一个;立即执行不让路;内联占并发", async () => { let busy = { queued: 1, running: 0 }; let now = 0; const c = new LaneController({ engine: { queueCounts: async (groups) => groups.map(() => busy) }, logger, lanes, readSettings: async () => ({ file: 2, network: 4, maintenance: 1 }), busyCacheMs: 0, maxYieldMs: 1000, now: () => now, }); await c.refresh(); expect(await c.enterMaintenance()).toBeNull(); expect(c.yielding()).toBe(true); now = 500; expect(await c.maintenanceMayStart()).toBe(false); now = 1000; expect(await c.maintenanceMayStart()).toBe(true); // 让路满 1s 放行一个 now = 1200; expect(await c.maintenanceMayStart()).toBe(false); // 重新计时 const forced = await c.enterMaintenance(true); expect(forced).not.toBeNull(); expect(c.active("maintenance")).toBe(1); expect(c.maxActive("maintenance")).toBe(0); // 引擎维护 worker 让出被内联占用的并发 forced?.(); busy = { queued: 0, running: 0 }; const r1 = await c.enterMaintenance(); expect(r1).not.toBeNull(); expect(c.yielding()).toBe(false); expect(await c.enterMaintenance()).toBeNull(); // 并发 1 已被占 r1?.(); const snap = await c.snapshot(); expect(snap.map((s) => [s.lane, s.concurrency, s.running, s.queued])).toEqual([ ["file", 2, 0, 0], ["network", 4, 0, 0], ["maintenance", 1, 0, 0], ]); }); }); // ── 真实 PG ── const url = process.env.DATABASE_URL ?? "postgres://umt:umt@localhost:55433/umt"; const uniq = `${process.pid.toString(36)}${Math.random().toString(36).slice(2, 7)}`; const biz = `server_lanes_${uniq}`; const ow = `${biz}_ow`; async function pgAvailable() { const c = new pg.Client({ connectionString: url, connectionTimeoutMillis: 2000 }); try { await c.connect(); await c.end(); return true; } catch { return false; } } const hasPg = process.env.UMT_SKIP_DB_TESTS !== "1" && (await pgAvailable()); async function waitUntil(fn: () => boolean | Promise, ms = 15_000) { const until = Date.now() + ms; while (Date.now() < until) { if (await fn()) return; await sleep(50); } throw new Error("等待超时"); } describe.skipIf(!hasPg)("io worker 分道(真实 PG)", () => { let db: Kysely; let booted: Bootstrapped; let worker: RunningWorker; const bus = new EventBus(50); const started: string[] = []; const done: string[] = []; const gates = new Map void>(); const task = (name: string, lane: TaskLane, block: boolean) => defineTask({ name, version: 1, lane, inputSchema: z.object({ id: z.string() }), retryPolicy: { maximumAttempts: 1 }, handler: async (t) => { started.push(t.input.id); if (block) await new Promise((resolve) => gates.set(t.input.id, resolve)); done.push(t.input.id); return null; }, }); const probe: ServerModule = { name: "probe", registerRoutes() {}, registerTasks(registry) { registry.register( task("laneFile", "file", true), task("laneNet", "network", false), task("laneNetBlock", "network", true), task("laneMaint", "maintenance", false), ); }, }; const enqueue = (name: string, id: string) => booted.ctx.engine.enqueue(name, { id }, { idempotencyKey: `${name}:${id}`, traceId: id }); beforeAll(async () => { const admin = new pg.Client({ connectionString: url }); await admin.connect(); await admin.query(`create schema ${biz}`); await admin.end(); db = new Kysely({ dialect: new PostgresDialect({ pool: new pg.Pool({ connectionString: url, max: 8, options: `-c search_path=${biz},public` }), }), }); await withSchemaLock(() => migrateToLatest(db, { migrationTableSchema: biz })); await db .insertInto("settings") .values({ key: "automation", value: JSON.stringify({ enabled: true, readOnlyReconcile: false, workerLanes: { file: 1, network: 2, maintenance: 1 }, }), }) .execute(); const config = ConfigSchema.parse({ DATABASE_URL: url, SECRET_MASTER_KEY: "test-master-key-0123456789" }); const { logger } = memoryLogger(); booted = await bootstrap({ role: "worker", config, logger, db, events: bus, modules: [probe], engineSchema: ow, engineMigrations: true, }); worker = await startWorker({ logger, heartbeatMs: 60_000, runtime: { ctx: booted.ctx, modules: [probe], setRelay: booted.outbox.setRelay, outboxIntervalMs: 100, group: "io", tasks: booted.tasks.all(), lanes: { refreshMs: 100, busyCacheMs: 0 }, }, }); }, 60_000); afterAll(async () => { for (const g of gates.values()) g(); await worker?.stop(); await booted?.engine.stop(); await db?.destroy(); await withSchemaLock(async () => { const admin = new pg.Client({ connectionString: url }); await admin.connect(); await admin.query(`drop schema if exists ${biz} cascade`); await admin.query(`drop schema if exists ${ow} cascade`); await admin.end(); }); }, 60_000); it("文件道占满(并发 1)时网络道照常领取;维护道在文件道有活时让路,空了再跑;改设置热生效", async () => { const lanes = worker.lanes!; expect(lanes.limit("file")).toBe(1); await enqueue("laneFile", "f1"); await enqueue("laneFile", "f2"); await waitUntil(() => started.includes("f1") || started.includes("f2")); const first = started.includes("f1") ? "f1" : "f2"; const second = first === "f1" ? "f2" : "f1"; // 文件道占满:网络任务仍被领取并完成;第二个文件任务排队 await enqueue("laneNet", "n1"); await waitUntil(() => done.includes("n1")); expect(started).not.toContain(second); // 维护道让路:文件道有运行中 + 排队 await enqueue("laneMaint", "m1"); await sleep(1500); expect(started).not.toContain("m1"); expect(lanes.yielding()).toBe(true); const snap = await lanes.snapshot(); expect(snap.find((s) => s.lane === "file")).toMatchObject({ concurrency: 1, running: 1, queued: 1 }); expect(snap.find((s) => s.lane === "maintenance")).toMatchObject({ queued: 1, yielding: true }); // 热生效:文件道并发改 2,排队的第二个文件任务被领取 await db .updateTable("settings") .set({ value: JSON.stringify({ enabled: true, readOnlyReconcile: false, workerLanes: { file: 2, network: 2, maintenance: 1 }, }), }) .where("key", "=", "automation") .execute(); await bus.publish(SETTINGS_CHANGED_EVENT, { keys: ["automation"] }); await waitUntil(() => started.includes(second)); expect(lanes.limit("file")).toBe(2); expect(started).not.toContain("m1"); // 文件道清空 → 维护道领取 gates.get(first)?.(); gates.get(second)?.(); await waitUntil(() => done.includes("m1")); expect(lanes.yielding()).toBe(false); }, 60_000); it("网络道占满时文件道照常领取", async () => { await enqueue("laneNetBlock", "nb1"); await enqueue("laneNetBlock", "nb2"); await waitUntil(() => started.includes("nb1") && started.includes("nb2")); await enqueue("laneFile", "f3"); await waitUntil(() => started.includes("f3")); for (const id of ["nb1", "nb2", "f3"]) gates.get(id)?.(); await waitUntil(() => ["nb1", "nb2", "f3"].every((id) => done.includes(id))); }, 60_000); });