UHD_Media_Tool/apps/server/src/http/sse-pg.ts

import type { EventEnvelope } from "@umt/contracts";
import { type DB, type Kysely, sql } from "@umt/db";
import pg from "pg";
import type { Logger } from "pino";
import type { BusListener, EventBusPort, SinceResult } from "./sse.ts";

/** 默认通知通道 */
export const EVENTS_CHANNEL = "umt_events";
/** 发布时串行化 seq 分配与提交的咨询锁键('umte') */
const PUBLISH_LOCK_KEY = 0x756d7465;

export interface PgEventBusOptions {
  db: Kysely<DB>;
  /** LISTEN 专用连接(不走连接池) */
  connectionString: string;
  logger: Logger;
  channel?: string;
  /** sse_event 保留最近多少条供 ?after= 补读,默认 10000 */
  retain?: number;
  /** 每发布多少条裁剪一次,默认 200 */
  pruneEvery?: number;
  /** LISTEN 断线重连退避上限,默认 10s */
  maxReconnectMs?: number;
  /** 仅测试:LISTEN 连接的 search_path 等 pg 启动参数 */
  pgOptions?: string;
}

type Row = { seq: number; type: string; at: Date; payload: unknown };
const toEnvelope = (r: Row): EventEnvelope => ({
  seq: Number(r.seq),
  type: r.type,
  at: r.at.toISOString(),
  payload: r.payload,
});

/**
 * 多进程 SSE 总线:发布 = 在 sse_event 追加一行(持咨询锁,保证 seq 顺序即提交顺序)+ pg_notify(channel, seq);
 * 每个进程一条 LISTEN 连接,收到通知后按 seq 从表里补读并分发给本进程订阅者(包括自己发布的)。
 * LISTEN 断线期间的事件在重连后按 lastDispatched 补读,不丢。
 */
export class PgEventBus implements EventBusPort {
  private readonly listeners = new Set<BusListener>();
  private readonly channel: string;
  private readonly retain: number;
  private readonly pruneEvery: number;
  private client: pg.Client | null = null;
  private lastDispatched = 0;
  private draining: Promise<void> | null = null;
  private drainAgain = false;
  private published = 0;
  private stopped = false;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;

  private constructor(private readonly opts: PgEventBusOptions) {
    this.channel = opts.channel ?? EVENTS_CHANNEL;
    if (!/^[a-z_][a-z0-9_]*$/.test(this.channel)) throw new Error(`非法通道名 ${this.channel}`);
    this.retain = opts.retain ?? 10_000;
    this.pruneEvery = opts.pruneEvery ?? 200;
  }

  /** 建总线并开始 LISTEN;从当前最大 seq 之后开始分发 */
  static async start(opts: PgEventBusOptions): Promise<PgEventBus> {
    const bus = new PgEventBus(opts);
    bus.lastDispatched = await bus.maxSeq();
    await bus.connect();
    return bus;
  }

  private get log() {
    return this.opts.logger;
  }

  private async maxSeq(): Promise<number> {
    const r = await sql<{ m: number | null }>`select max(seq) as m from sse_event`.execute(this.opts.db);
    return Number(r.rows[0]?.m ?? 0);
  }

  private async connect(): Promise<void> {
    const client = new pg.Client({
      connectionString: this.opts.connectionString,
      ...(this.opts.pgOptions ? { options: this.opts.pgOptions } : {}),
    });
    client.on("notification", (msg) => {
      if (msg.channel === this.channel) this.scheduleDrain();
    });
    const onLost = (err?: Error) => {
      if (this.client !== client) return;
      this.client = null;
      client.removeAllListeners("notification");
      client.end().catch(() => {});
      if (!this.stopped) {
        this.log.warn({ err }, "SSE 总线 LISTEN 连接断开,准备重连");
        this.reconnect(500);
      }
    };
    client.on("error", onLost);
    client.on("end", () => onLost());
    await client.connect();
    await client.query(`listen ${this.channel}`);
    this.client = client;
    // 断线期间可能漏掉通知:连上后补读一次
    this.scheduleDrain();
  }

  private reconnect(delayMs: number) {
    if (this.stopped || this.reconnectTimer) return;
    this.reconnectTimer = setTimeout(async () => {
      this.reconnectTimer = null;
      try {
        await this.connect();
        this.log.info("SSE 总线 LISTEN 已重连");
      } catch (err) {
        this.log.warn({ err }, "SSE 总线重连失败");
        this.reconnect(Math.min(delayMs * 2, this.opts.maxReconnectMs ?? 10_000));
      }
    }, delayMs);
  }

  private scheduleDrain() {
    if (this.draining) {
      this.drainAgain = true;
      return;
    }
    this.draining = (async () => {
      try {
        do {
          this.drainAgain = false;
          await this.drainOnce();
        } while (this.drainAgain && !this.stopped);
      } catch (err) {
        this.log.warn({ err }, "SSE 总线补读失败");
      } finally {
        this.draining = null;
      }
    })();
  }

  private async drainOnce() {
    for (;;) {
      const rows = await sql<Row>`select seq, type, at, payload from sse_event
        where seq > ${this.lastDispatched} order by seq limit 500`.execute(this.opts.db);
      for (const r of rows.rows) {
        const e = toEnvelope(r);
        this.lastDispatched = e.seq;
        for (const l of this.listeners) {
          try {
            l(e);
          } catch (err) {
            this.log.warn({ err, type: e.type }, "SSE 订阅者异常");
          }
        }
      }
      if (rows.rows.length < 500) return;
    }
  }

  async publish<P>(type: string, payload: P): Promise<EventEnvelope<P> | null> {
    try {
      const row = await this.opts.db.transaction().execute(async (trx) => {
        await sql`select pg_advisory_xact_lock(${PUBLISH_LOCK_KEY})`.execute(trx);
        const r = await sql<Row>`insert into sse_event (type, payload)
          values (${type}, ${JSON.stringify(payload ?? null)}::jsonb) returning seq, type, at, payload`.execute(trx);
        const inserted = r.rows[0]!;
        // NOTIFY 在提交时才投递
        await sql`select pg_notify(${this.channel}, ${String(inserted.seq)})`.execute(trx);
        return inserted;
      });
      this.published += 1;
      if (this.published % this.pruneEvery === 0) void this.prune();
      return toEnvelope(row) as EventEnvelope<P>;
    } catch (err) {
      this.log.error({ err, type }, "SSE 事件发布失败");
      return null;
    }
  }

  /** 只保留最近 retain 条 */
  async prune(): Promise<void> {
    try {
      await sql`delete from sse_event where seq <= (select max(seq) from sse_event) - ${this.retain}`.execute(
        this.opts.db,
      );
    } catch (err) {
      this.log.warn({ err }, "sse_event 裁剪失败");
    }
  }

  async since(after: number, limit = 1000): Promise<SinceResult> {
    const [rows, bounds] = await Promise.all([
      sql<Row>`select seq, type, at, payload from sse_event where seq > ${after} order by seq limit ${limit}`.execute(
        this.opts.db,
      ),
      sql<{ lo: number | null; hi: number | null }>`select min(seq) as lo, max(seq) as hi from sse_event`.execute(
        this.opts.db,
      ),
    ]);
    const lo = bounds.rows[0]?.lo;
    const hi = Number(bounds.rows[0]?.hi ?? 0);
    // 裁剪过:请求起点早于最老保留条目
    const gap = after > 0 && lo != null && after < Number(lo) - 1;
    return { events: rows.rows.map(toEnvelope), gap, latestSeq: Math.max(hi, after) };
  }

  subscribe(l: BusListener): () => void {
    this.listeners.add(l);
    return () => this.listeners.delete(l);
  }

  async stop(): Promise<void> {
    this.stopped = true;
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    const c = this.client;
    this.client = null;
    await c?.end().catch(() => {});
    await this.draining;
  }
}