WebSocket

封装工具类

class WebSocketManager {
  constructor(options = {}) {
    // 默认配置
    this.options = {
      url: options.url || "",
      reconnectInterval: options.reconnectInterval || 5000, // 重连间隔时间（毫秒）
      heartbeatInterval: options.heartbeatInterval || 30000, // 心跳间隔时间（毫秒）
      maxReconnectAttempts: options.maxReconnectAttempts || 5, // 最大重连次数
      heartbeatMessage:
        options.heartbeatMessage || JSON.stringify({ type: "ping" }), // 心跳消息
      onOpen: options.onOpen || (() => {}),
      onClose: options.onClose || (() => {}),
      onError: options.onError || (() => {}),
      onMessage: options.onMessage || (() => {}),
      onReconnect: options.onReconnect || (() => {}),
      ...options,
    };

    this.ws = null;
    this.reconnectAttempts = 0;
    this.isClosing = false;
    this.isManuallyClosed = false;
    this.heartbeatTimer = null;
    this.reconnectTimer = null;
    this.id = options.id || this.generateId();
  }

  // 生成唯一ID
  generateId() {
    return `ws_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;
  }

  // 连接WebSocket
  connect() {
    if (this.isClosing || this.isManuallyClosed) return;

    try {
      this.ws = new WebSocket(this.options.url);

      this.ws.onopen = (event) => {
        console.log(`WebSocket连接已建立: ${this.id}`);
        this.reconnectAttempts = 0; // 连接成功，重置重连次数
        this.startHeartbeat(); // 开始心跳
        this.options.onOpen(event, this.id);
      };

      this.ws.onmessage = (event) => {
        // 如果收到心跳响应，不执行用户回调
        if (event.data === this.options.heartbeatMessage) return;
        this.options.onMessage(event, this.id);
      };

      this.ws.onclose = (event) => {
        console.log(`WebSocket连接已关闭: ${this.id}`);
        this.stopHeartbeat();
        this.handleReconnect(event);
        this.options.onClose(event, this.id);
      };

      this.ws.onerror = (event) => {
        console.error(`WebSocket发生错误: ${this.id}`, event);
        this.stopHeartbeat();
        this.options.onError(event, this.id);
      };
    } catch (error) {
      console.error(`WebSocket连接失败: ${this.id}`, error);
      this.handleReconnect();
    }
  }

  // 开始心跳
  startHeartbeat() {
    this.stopHeartbeat(); // 先清除可能存在的定时器
    this.heartbeatTimer = setInterval(() => {
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(this.options.heartbeatMessage);
      }
    }, this.options.heartbeatInterval);
  }

  // 停止心跳
  stopHeartbeat() {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
      this.heartbeatTimer = null;
    }
  }

  // 处理重连逻辑
  handleReconnect(event) {
    if (this.isClosing || this.isManuallyClosed) return;

    if (this.reconnectAttempts < this.options.maxReconnectAttempts) {
      this.reconnectAttempts++;
      console.log(
        `WebSocket尝试重连 (${this.reconnectAttempts}/${this.options.maxReconnectAttempts}): ${this.id}`
      );
      this.options.onReconnect(
        this.reconnectAttempts,
        this.options.maxReconnectAttempts,
        this.id
      );

      this.reconnectTimer = setTimeout(() => {
        this.connect();
      }, this.options.reconnectInterval);
    } else {
      console.warn(`WebSocket重连次数已达上限: ${this.id}`);
    }
  }

  // 发送消息
  send(data) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(data);
      return true;
    } else {
      console.warn(`WebSocket未连接，无法发送消息: ${this.id}`);
      return false;
    }
  }

  // 主动关闭连接（单个）
  close(code = 1000, reason = "") {
    this.isManuallyClosed = true;
    this.isClosing = true;
    this.stopHeartbeat();

    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }

    if (this.ws) {
      this.ws.close(code, reason);
      this.ws = null;
    }

    console.log(`WebSocket连接已手动关闭: ${this.id}`);
  }

  // 获取当前连接状态
  getReadyState() {
    if (!this.ws) return WebSocket.CLOSED;
    return this.ws.readyState;
  }

  // 获取连接状态文本
  getReadyStateText() {
    const states = {
      [WebSocket.CONNECTING]: "CONNECTING",
      [WebSocket.OPEN]: "OPEN",
      [WebSocket.CLOSING]: "CLOSING",
      [WebSocket.CLOSED]: "CLOSED",
    };
    return states[this.getReadyState()];
  }
}

// WebSocket连接池管理类
class WebSocketPool {
  constructor() {
    this.connections = new Map(); // 存储WebSocket连接
  }

  // 添加连接到池中
  addConnection(options) {
    const wsManager = new WebSocketManager(options);
    this.connections.set(wsManager.id, wsManager);
    wsManager.connect();
    return wsManager.id;
  }

  // 根据ID获取连接
  getConnection(id) {
    return this.connections.get(id);
  }

  // 发送消息到指定连接
  sendTo(id, data) {
    const connection = this.getConnection(id);
    if (connection) {
      return connection.send(data);
    }
    console.warn(`WebSocket连接不存在: ${id}`);
    return false;
  }

  // 关闭指定连接
  closeConnection(id, code, reason) {
    const connection = this.connections.get(id);
    if (connection) {
      connection.close(code, reason);
      this.connections.delete(id);
      console.log(`连接 ${id} 已从连接池中移除`);
    }
  }

  // 批量关闭所有连接
  closeAllConnections(code = 1000, reason = "") {
    for (const [id, connection] of this.connections) {
      connection.close(code, reason);
    }
    this.connections.clear();
    console.log("所有WebSocket连接已关闭");
  }

  // 获取所有连接信息
  getAllConnectionsInfo() {
    const info = [];
    for (const [id, connection] of this.connections) {
      info.push({
        id,
        url: connection.options.url,
        readyState: connection.getReadyStateText(),
      });
    }
    return info;
  }

  // 获取连接数量
  getConnectionCount() {
    return this.connections.size;
  }

  // 移除连接
  removeConnection(id) {
    const connection = this.connections.get(id);
    if (connection) {
      connection.close();
      this.connections.delete(id);
    }
  }
}

// 创建全局连接池实例
const wsPool = new WebSocketPool();

// 导出WebSocket管理类和连接池
export { WebSocketManager, WebSocketPool, wsPool };

// 提供便捷的API
export const createWebSocket = (options) => {
  return wsPool.addConnection(options);
};

export const closeWebSocket = (id, code, reason) => {
  wsPool.closeConnection(id, code, reason);
};

export const closeAllWebSockets = (code, reason) => {
  wsPool.closeAllConnections(code, reason);
};

export const sendWebSocketMessage = (id, data) => {
  return wsPool.sendTo(id, data);
};