chanjs 2.7.2 → 2.7.4

package/realtime/sse.js

@@ -0,0 +1,424 @@
+/**
+ * SSE (Server-Sent Events) 管理器
+ *
+ * ============================================================
+ * 特性
+ * ============================================================
+ * 1. 零依赖：基于 Node 原生 http/express 的 res 对象实现
+ * 2. 房间（Room）机制：可分组广播，也可对单个连接推送
+ * 3. 心跳保活：默认 30s 发送 :heartbeat 注释，防止代理超时断开
+ * 4. 鉴权钩子：可选的 authenticate 回调，按业务校验 req
+ * 5. 优雅停机：shutdown() 关闭所有连接，清理定时器
+ *
+ * ============================================================
+ * 使用方法
+ * ============================================================
+ *
+ * // 1. 注册 SSE 路由（业务侧只需一行）
+ * import { sse } from 'chanjs';
+ * app.get('/sse', sse.middleware(), (req, res) => {
+ *   sse.join(res, req.query.room || 'default');
+ * });
+ *
+ * // 2. 单推
+ * sse.send(clientId, { type: 'message', data: 'hello' });
+ *
+ * // 3. 广播到房间
+ * sse.broadcast('room-1', { type: 'notice', data: 'hello all' });
+ *
+ * // 4. 全局广播
+ * sse.broadcastAll({ type: 'ping', data: Date.now() });
+ *
+ * // 5. 带鉴权
+ * app.get('/sse-auth', sse.middleware({
+ *   authenticate: (req) => req.headers.token === 'xxx',
+ * }), (req, res) => sse.join(res, 'user-' + req.user.id));
+ *
+ * // 6. 状态查询
+ * sse.getStatus();  // { connections, rooms, uptime }
+ *
+ * ============================================================
+ * 协议说明
+ * ============================================================
+ * SSE 响应头：
+ *   Content-Type: text/event-stream
+ *   Cache-Control: no-cache
+ *   Connection: keep-alive
+ * 消息格式：
+ *   data: <json>\n\n
+ *   event: <event-type>\ndata: <json>\n\n
+ * 注释（心跳）：
+ *   : heartbeat\n\n
+ */
+
+/**
+ * SSE 默认心跳间隔（毫秒）
+ * 30s 是代理超时的安全值（Nginx proxy_read_timeout 默认 60s）
+ */
+const DEFAULT_HEARTBEAT_INTERVAL = 30000;
+
+/**
+ * SSE 管理器
+ * @class SSEManager
+ */
+class SSEManager {
+  constructor() {
+    /** @type {Map<string, {res: Object, rooms: Set<string>, joinedAt: number}>} 连接表 id -> client */
+    this._clients = new Map();
+    /** @type {Map<string, Set<string>>} 房间表 room -> Set<clientId> */
+    this._rooms = new Map();
+    /** 心跳定时器句柄 */
+    this._heartbeatTimer = null;
+    /** 心跳间隔 */
+    this._heartbeatInterval = DEFAULT_HEARTBEAT_INTERVAL;
+    /** 自增连接 id */
+    this._idSeq = 0;
+    /** 启动时间 */
+    this._startedAt = Date.now();
+    /** 是否已 shutdown */
+    this._shuttingDown = false;
+    /** 最大并发连接数上限，超出直接返回 503，防止连接洪泛耗尽内存 */
+    this._maxConnections = 10000;
+  }
+
+  /**
+   * Express 中间件：建立 SSE 连接
+   * @param {Object} [options]
+   * @param {Function} [options.authenticate] - 鉴权钩子 (req) => boolean | Promise<boolean>
+   * @param {number} [options.heartbeatInterval] - 心跳间隔，默认 30000
+   * @returns {Function} Express middleware
+   */
+  middleware(options = {}) {
+    const { authenticate, heartbeatInterval } = options;
+
+    return async (req, res, next) => {
+      if (this._shuttingDown) {
+        res.status(503).json({ success: false, code: 5002, msg: '服务不可用' });
+        return;
+      }
+
+      // 连接数上限校验，防止连接洪泛导致内存耗尽
+      if (this._clients.size >= this._maxConnections) {
+        res.status(503).json({ success: false, code: 5002, msg: 'SSE 连接数已达上限，请稍后再试' });
+        return;
+      }
+
+      // 鉴权钩子
+      if (typeof authenticate === 'function') {
+        let ok;
+        try {
+          ok = await authenticate(req);
+        } catch (e) {
+          ok = false;
+        }
+        if (!ok) {
+          res.status(401).json({ success: false, code: 1001, msg: '认证失败' });
+          return;
+        }
+      }
+
+      // 设置 SSE 响应头
+      res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
+      res.setHeader('Cache-Control', 'no-cache, no-transform');
+      res.setHeader('Connection', 'keep-alive');
+      // 禁用代理缓冲（Nginx X-Accel-Buffering）
+      res.setHeader('X-Accel-Buffering', 'no');
+      res.flushHeaders?.();
+
+      // 生成连接 id 并注册
+      const clientId = `sse_${++this._idSeq}_${Date.now()}`;
+      this._clients.set(clientId, {
+        res,
+        rooms: new Set(),
+        joinedAt: Date.now(),
+      });
+
+      // 启动心跳（首次连接时启动，全局共享）
+      this._ensureHeartbeat(heartbeatInterval);
+
+      // 注入到 req，便于业务侧拿到 clientId
+      req.sseClientId = clientId;
+
+      // 连接关闭时清理（res 是 Express 响应对象，标准 Node EventEmitter）
+      const cleanup = () => {
+        this._removeClient(clientId);
+      };
+      req.on?.('close', cleanup);
+      req.on?.('aborted', cleanup);
+      res.on?.('close', cleanup);
+      res.on?.('error', cleanup);
+
+      next?.();
+    };
+  }
+
+  /**
+   * 将连接加入房间
+   * @param {Object} res - Express res（建立 SSE 时的那个 res）
+   * @param {string} room - 房间名
+   * @returns {string|null} clientId（成功时返回，便于业务侧记录）
+   */
+  join(res, room) {
+    const clientId = this._findClientIdByRes(res);
+    if (!clientId) return null;
+
+    const client = this._clients.get(clientId);
+    if (!client) return null;
+
+    client.rooms.add(room);
+    if (!this._rooms.has(room)) this._rooms.set(room, new Set());
+    this._rooms.get(room).add(clientId);
+    return clientId;
+  }
+
+  /**
+   * 将连接离开房间
+   * @param {Object|string} target - res 或 clientId
+   * @param {string} room - 房间名
+   */
+  leave(target, room) {
+    const clientId = typeof target === 'string' ? target : this._findClientIdByRes(target);
+    if (!clientId) return false;
+    const client = this._clients.get(clientId);
+    if (!client) return false;
+
+    client.rooms.delete(room);
+    const members = this._rooms.get(room);
+    if (members) {
+      members.delete(clientId);
+      if (members.size === 0) this._rooms.delete(room);
+    }
+    return true;
+  }
+
+  /**
+   * 向单个连接发送消息
+   * @param {string} clientId - 连接 id
+   * @param {*} data - 任意可序列化数据
+   * @param {Object} [options]
+   * @param {string} [options.event] - SSE event 字段（可选）
+   * @returns {boolean} 是否发送成功
+   */
+  send(clientId, data, options = {}) {
+    const client = this._clients.get(clientId);
+    if (!client) return false;
+    return this._write(client.res, data, options.event);
+  }
+
+  /**
+   * 广播到指定房间
+   * @param {string} room - 房间名
+   * @param {*} data - 任意可序列化数据
+   * @param {Object} [options]
+   * @param {string} [options.event] - SSE event 字段
+   * @returns {number} 实际送达的连接数
+   */
+  broadcast(room, data, options = {}) {
+    const members = this._rooms.get(room);
+    if (!members || members.size === 0) return 0;
+
+    let sent = 0;
+    for (const clientId of members) {
+      if (this.send(clientId, data, options)) sent++;
+    }
+    return sent;
+  }
+
+  /**
+   * 广播到所有连接
+   * @param {*} data - 任意可序列化数据
+   * @param {Object} [options]
+   * @param {string} [options.event] - SSE event 字段
+   * @returns {number} 实际送达的连接数
+   */
+  broadcastAll(data, options = {}) {
+    let sent = 0;
+    for (const clientId of this._clients.keys()) {
+      if (this.send(clientId, data, options)) sent++;
+    }
+    return sent;
+  }
+
+  /**
+   * 获取当前连接数
+   * @returns {number}
+   */
+  size() {
+    return this._clients.size;
+  }
+
+  /**
+   * 获取指定房间连接数
+   * @param {string} room
+   * @returns {number}
+   */
+  roomSize(room) {
+    return this._rooms.get(room)?.size || 0;
+  }
+
+  /**
+   * 获取运行状态
+   * @returns {Object} { connections, rooms, uptime }
+   */
+  getStatus() {
+    return {
+      connections: this._clients.size,
+      rooms: this._rooms.size,
+      roomList: Array.from(this._rooms.keys()),
+      uptime: Date.now() - this._startedAt,
+    };
+  }
+
+  /**
+   * 优雅停机：关闭所有连接，清理心跳
+   */
+  shutdown() {
+    this._shuttingDown = true;
+
+    // 停止心跳
+    if (this._heartbeatTimer) {
+      clearInterval(this._heartbeatTimer);
+      this._heartbeatTimer = null;
+    }
+
+    // 主动关闭所有连接
+    for (const [, client] of this._clients) {
+      try {
+        client.res.end();
+      } catch (e) {
+        // 忽略已关闭的连接
+      }
+    }
+    this._clients.clear();
+    this._rooms.clear();
+    console.log(`[SSE] 已关闭所有 SSE 连接`);
+  }
+
+  /**
+   * 写入 SSE 数据（底层方法）
+   * @private
+   * @param {Object} res - Express res
+   * @param {*} data - 数据
+   * @param {string} [event] - 事件类型
+   * @returns {boolean} 是否写入成功
+   */
+  _write(res, data, event) {
+    if (res.writableEnded || res.destroyed) return false;
+
+    const payload = typeof data === 'string' ? data : JSON.stringify(data);
+    let chunk = '';
+    if (event) chunk += `event: ${event}\n`;
+    // 多行 data 拼接：每行前缀 data:
+    chunk += payload
+      .split('\n')
+      .map((line) => `data: ${line}`)
+      .join('\n');
+    chunk += '\n\n';
+
+    try {
+      res.write(chunk);
+      return true;
+    } catch (e) {
+      return false;
+    }
+  }
+
+  /**
+   * 启动心跳（懒加载，首次连接时启动）
+   * @private
+   */
+  _ensureHeartbeat(interval) {
+    if (this._heartbeatTimer) return;
+    if (typeof interval === 'number' && interval > 0) {
+      this._heartbeatInterval = interval;
+    }
+
+    this._heartbeatTimer = setInterval(() => {
+      if (this._clients.size === 0) return;
+      for (const [, client] of this._clients) {
+        // SSE 注释格式：: heartbeat\n\n（不会触发前端 message 事件）
+        this._writeComment(client.res, 'heartbeat');
+      }
+    }, this._heartbeatInterval);
+
+    // 不阻止进程退出
+    if (this._heartbeatTimer.unref) this._heartbeatTimer.unref();
+  }
+
+  /**
+   * 写入 SSE 注释（心跳用，不触发前端 message 事件）
+   * @private
+   */
+  _writeComment(res, text) {
+    if (res.writableEnded || res.destroyed) return false;
+    try {
+      res.write(`: ${text}\n\n`);
+      return true;
+    } catch (e) {
+      return false;
+    }
+  }
+
+  /**
+   * 移除客户端（清理房间索引）
+   * @private
+   */
+  _removeClient(clientId) {
+    const client = this._clients.get(clientId);
+    if (!client) return;
+
+    // 从所有房间移除
+    for (const room of client.rooms) {
+      const members = this._rooms.get(room);
+      if (members) {
+        members.delete(clientId);
+        if (members.size === 0) this._rooms.delete(room);
+      }
+    }
+
+    this._clients.delete(clientId);
+  }
+
+  /**
+   * 根据 res 反查 clientId
+   * @private
+   */
+  _findClientIdByRes(res) {
+    for (const [id, client] of this._clients) {
+      if (client.res === res) return id;
+    }
+    return null;
+  }
+}
+
+// 全局单例（懒加载：首次访问属性时才实例化，避免 import 即实例化浪费内存）
+let _sseInstance = null;
+const getSSEInstance = () => {
+  if (!_sseInstance) _sseInstance = new SSEManager();
+  return _sseInstance;
+};
+
+// 用 Proxy 保持 `import { sse } from 'chanjs'` API 兼容，访问任意属性时才实例化
+const sse = new Proxy(
+  {},
+  {
+    get(_target, prop) {
+      const inst = getSSEInstance();
+      const val = inst[prop];
+      // 方法绑定 this，避免解构调用时丢失上下文
+      return typeof val === 'function' ? val.bind(inst) : val;
+    },
+    set(_target, prop, value) {
+      getSSEInstance()[prop] = value;
+      return true;
+    },
+    has(_target, prop) {
+      return prop in getSSEInstance();
+    },
+    getPrototypeOf() {
+      return SSEManager.prototype;
+    },
+  }
+);
+
+export { sse as default, sse, SSEManager, DEFAULT_HEARTBEAT_INTERVAL };