UHD_Media_Tool/apps/server/src/modules/schedules/scheduler.ts

import { SCHEDULE_CHANGED_EVENT } from "@umt/contracts";
import { type DB, type Kysely, type ScheduledJobRow, sql } from "@umt/db";
import type { Logger } from "pino";
import type { EventBusPort } from "../../http/sse.ts";
import type { ServerContext } from "../types.ts";
import { runAgentCheck } from "./agent-checks.ts";
import type { JobRunFn, JobRunResult, ScheduleRegistry } from "./registry.ts";
import { computeNext, LOCK_NS, parseSpec, slotOf, syncDefinitions, trimRuns } from "./store.ts";

// 调度器:只在 io worker(或 all)运行;多个实例用 pg_try_advisory_lock 选出唯一执行者(锁在一条专用连接上,
// 连接断开即释放,其它实例下一轮接手)。执行者每 tickMs 查一次到期(enabled 且 next_run_at 已到)或被请求
// 立即执行(run_requested_at)的任务,逐个异步执行(不同任务并行,同一任务不重入),结果写回表与运行历史。
// interval 任务的下次时间 = 本次结束 + 周期(与原来的自调度循环一致);cron 任务取下一个匹配时间。
// 立即执行不改下次运行时间。

export interface SchedulerOptions {
  /** 查到期任务的间隔(缺省 1s) */
  tickMs?: number;
  /** 抢锁失败后的重试间隔(缺省 5s) */
  lockRetryMs?: number;
  /** 回填巡检 Agent 结果的间隔(缺省 15s) */
  harvestMs?: number;
  /** 停止时等待进行中任务的上限(缺省 10s) */
  stopWaitMs?: number;
  now?: () => Date;
}

const TERMINAL_AGENT = ["succeeded", "failed", "cancelled", "budget_exceeded"];

const errText = (err: unknown) => (err instanceof Error ? err.message : String(err)).slice(0, 2000);

export class Scheduler {
  private stopped = false;
  private wake: (() => void) | null = null;
  private loop: Promise<void> | null = null;
  private inflight = new Map<string, Promise<void>>();
  private syncedVersion = -1;
  private lastHarvest = 0;
  leader = false;
  private readonly db: Kysely<DB>;
  private readonly log: Logger;
  private readonly events: EventBusPort;
  private readonly now: () => Date;

  constructor(
    private readonly ctx: ServerContext,
    private readonly registry: ScheduleRegistry,
    private readonly opts: SchedulerOptions = {},
  ) {
    this.db = ctx.db;
    this.log = ctx.logger.child({ component: "scheduler" });
    this.events = ctx.events;
    this.now = opts.now ?? (() => new Date());
  }

  start(): void {
    if (this.loop) return;
    this.stopped = false;
    this.loop = this.leaderLoop();
  }

  async stop(): Promise<void> {
    this.stopped = true;
    this.wake?.();
    await this.loop;
    this.loop = null;
    const pending = [...this.inflight.values()];
    if (pending.length) {
      await Promise.race([Promise.allSettled(pending), this.sleep(this.opts.stopWaitMs ?? 10_000, true)]);
    }
  }

  /** 可被 stop 打断的等待 */
  private sleep(ms: number, ignoreStop = false): Promise<void> {
    return new Promise((resolve) => {
      const t = setTimeout(done, ms);
      function done() {
        clearTimeout(t);
        resolve();
      }
      if (!ignoreStop) this.wake = done;
    });
  }

  private async leaderLoop(): Promise<void> {
    const retry = this.opts.lockRetryMs ?? 5000;
    while (!this.stopped) {
      try {
        await this.db.connection().execute(async (conn) => {
          while (!this.stopped) {
            const r = await sql<{ ok: boolean }>`
              select pg_try_advisory_lock(${LOCK_NS}::int4, hashtext(current_schema())) as ok`.execute(conn);
            if (r.rows[0]?.ok) break;
            await this.sleep(retry);
          }
          if (this.stopped) return;
          this.leader = true;
          this.log.info("调度器取得执行权");
          try {
            await this.recoverInterrupted();
            while (!this.stopped) {
              await this.tick();
              await this.sleep(this.opts.tickMs ?? 1000);
              // 探活:专用连接断了锁也就没了,抛错后重新抢锁
              await sql`select 1`.execute(conn);
            }
          } finally {
            this.leader = false;
            await sql`select pg_advisory_unlock(${LOCK_NS}::int4, hashtext(current_schema()))`
              .execute(conn)
              .catch(() => undefined);
          }
        });
      } catch (err) {
        this.leader = false;
        this.log.warn({ err: { message: errText(err) } }, "调度器循环出错，稍后重试");
        if (!this.stopped) await this.sleep(retry);
      }
    }
  }

  /** 上一个执行者中断时留下的 running 状态:标为失败 */
  private async recoverInterrupted(): Promise<void> {
    const msg = "调度器重启，上次运行中断";
    await this.db
      .updateTable("scheduled_job_run")
      .set({ status: "failed", error: msg, finished_at: this.now() })
      .where("status", "=", "running")
      .execute();
    await this.db
      .updateTable("scheduled_job")
      .set({ running_since: null, last_status: "failed", last_error: msg })
      .where("running_since", "is not", null)
      .execute();
  }

  /** 一轮:同步定义 → 找到期任务 → 异步执行;定期回填巡检结果。测试可直接调用 */
  async tick(): Promise<string[]> {
    const now = this.now();
    if (this.registry.version !== this.syncedVersion) {
      const v = this.registry.version;
      await syncDefinitions(this.db, this.registry.list(), now);
      this.syncedVersion = v;
    }
    const busy = [...this.inflight.keys()];
    const due = await this.db
      .selectFrom("scheduled_job")
      .selectAll()
      .where("available", "=", true)
      .where((eb) =>
        eb.or([
          eb("run_requested_at", "is not", null),
          eb.and([eb("enabled", "=", true), eb("next_run_at", "<=", now)]),
        ]),
      )
      .$if(busy.length > 0, (q) => q.where("id", "not in", busy))
      .orderBy("next_run_at")
      .limit(50)
      .execute();
    const started: string[] = [];
    for (const job of due) {
      const fn = this.runFor(job);
      if (!fn) continue;
      const p = this.execute(job, fn).finally(() => this.inflight.delete(job.id));
      this.inflight.set(job.id, p);
      started.push(job.id);
    }
    if (Date.now() - this.lastHarvest >= (this.opts.harvestMs ?? 15_000)) {
      this.lastHarvest = Date.now();
      await this.harvest().catch((err) => this.log.warn({ err: { message: errText(err) } }, "巡检结果回填失败"));
    }
    return started;
  }

  /** 等进行中的任务全部结束(测试用) */
  async idle(): Promise<void> {
    while (this.inflight.size) await Promise.allSettled([...this.inflight.values()]);
  }

  private runFor(job: ScheduledJobRow): JobRunFn | null {
    const def = this.registry.get(job.id);
    if (def) return def.run;
    // 自定义巡检不在代码里登记,共用巡检执行体
    if (job.kind === "agent_check") return (c) => runAgentCheck(this.ctx, c);
    return null;
  }

  private publish(jobId: string, status: string | null) {
    void this.events.publish(SCHEDULE_CHANGED_EVENT, { jobId, status });
  }

  private async execute(job: ScheduledJobRow, fn: JobRunFn): Promise<void> {
    const manual = job.run_requested_at !== null;
    const spec = parseSpec(job.schedule);
    const startedAt = this.now();
    const t0 = performance.now();
    // 认领:清掉请求标志、标运行中(条件更新,别的执行者抢先认领过就放弃)
    const claimed = await this.db
      .updateTable("scheduled_job")
      .set({ running_since: startedAt, last_status: "running", run_requested_at: null })
      .where("id", "=", job.id)
      .where("running_since", "is", null)
      .returning("id")
      .executeTakeFirst();
    if (!claimed) return;
    const run = await this.db
      .insertInto("scheduled_job_run")
      .values({ job_id: job.id, trigger: manual ? "manual" : "schedule", status: "running", started_at: startedAt })
      .returning("id")
      .executeTakeFirstOrThrow();
    if (manual || job.last_status !== "ok") this.publish(job.id, "running");

    let status: "ok" | "skipped" | "failed" = "ok";
    let error: string | null = null;
    let result: JobRunResult | undefined;
    try {
      result =
        (await fn({
          jobId: job.id,
          runId: run.id,
          trigger: manual ? "manual" : "schedule",
          now: startedAt,
          config: (job.config ?? {}) as Record<string, unknown>,
          slot: manual ? `manual:${run.id}` : slotOf(spec, startedAt, job.next_run_at),
        })) ?? undefined;
      status = result?.status ?? "ok";
    } catch (err) {
      status = "failed";
      error = errText(err);
      this.log.error({ err: { message: error }, job: job.id }, "定时任务执行失败");
    }
    const finishedAt = this.now();
    const durationMs = Math.max(0, Math.round(performance.now() - t0));
    // 计划触发(或从未排过)才推进下次时间;立即执行保持原计划
    const advance = !manual || job.next_run_at === null || job.next_run_at <= startedAt;
    try {
      await this.db.transaction().execute(async (trx) => {
        const cur = await trx
          .selectFrom("scheduled_job")
          .select(["schedule"])
          .where("id", "=", job.id)
          .executeTakeFirst();
        if (!cur) return; // 运行期间被删除
        const liveSpec = parseSpec(cur.schedule);
        await trx
          .updateTable("scheduled_job")
          .set({
            running_since: null,
            last_run_at: startedAt,
            last_status: status,
            last_duration_ms: durationMs,
            last_error: error,
            run_count: sql`run_count + 1`,
            ...(status === "failed" ? { fail_count: sql`fail_count + 1` } : {}),
            ...(advance
              ? { next_run_at: computeNext(liveSpec, liveSpec.type === "interval" ? finishedAt : startedAt) }
              : {}),
          })
          .where("id", "=", job.id)
          .execute();
        await trx
          .updateTable("scheduled_job_run")
          .set({
            status,
            finished_at: finishedAt,
            duration_ms: durationMs,
            error,
            summary: result?.summary ?? null,
            agent_run_id: result?.agentRunId ?? null,
          })
          .where("id", "=", run.id)
          .execute();
        await trimRuns(trx, job.id);
      });
    } catch (err) {
      this.log.error({ err: { message: errText(err) }, job: job.id }, "定时任务结果写回失败");
    }
    if (manual || status !== job.last_status || status === "failed") this.publish(job.id, status);
  }

  /** 巡检入队的 agent_run 结束后:回填终态与摘要;配置了 reportAsDecision 的写一条 info 待处理 */
  async harvest(): Promise<number> {
    const rows = await sql<{
      id: string;
      job_id: string;
      job_name: string;
      config: Record<string, unknown>;
      agent_run_id: string;
      status: string;
      summary: string | null;
    }>`
      select r.id, r.job_id, j.name as job_name, j.config, r.agent_run_id, a.status, a.summary
        from scheduled_job_run r
        join scheduled_job j on j.id = r.job_id
        join agent_run a on a.id = r.agent_run_id
       where r.agent_status is null and a.status = any(${TERMINAL_AGENT})
       limit 50`.execute(this.db);
    for (const r of rows.rows) {
      const summary = (r.summary ?? "").trim() || `智能体运行结束：${r.status}`;
      await this.db
        .updateTable("scheduled_job_run")
        .set({ agent_status: r.status, summary: summary.slice(0, 4000) })
        .where("id", "=", r.id)
        .execute();
      if (r.config?.reportAsDecision === true && r.status === "succeeded") {
        const ins = await sql<{ id: string }>`
          insert into decision (kind, severity, title, summary, subject_type, subject_id, dedupe_key, evidence, agent_run_id)
          values ('other', 'info', ${`巡检报告：${r.job_name}`}, ${summary.slice(0, 4000)}, 'agent_run', ${r.agent_run_id},
                  ${`schedule-report:${r.id}`}, ${JSON.stringify({ scheduledJobId: r.job_id, scheduledRunId: r.id })}::jsonb,
                  ${r.agent_run_id})
          on conflict do nothing
          returning id`.execute(this.db);
        const d = ins.rows[0];
        if (d)
          await this.ctx.outbox.emit(
            this.db,
            "decision.created",
            { decisionId: d.id, kind: "other", severity: "info", title: `巡检报告：${r.job_name}` },
            { idempotencyKey: `decision.created:${d.id}` },
          );
      }
      this.publish(r.job_id, null);
    }
    return rows.rows.length;
  }
}