UHD_Media_Tool/apps/server/test/worker-lanes.test.ts

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<string, string> {
  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<string, string>();
  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<boolean>, 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<DB>;
  let booted: Bootstrapped;
  let worker: RunningWorker;
  const bus = new EventBus(50);
  const started: string[] = [];
  const done: string[] = [];
  const gates = new Map<string, () => 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<void>((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<DB>({
      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);
});