前端WebSocket 封装

前言

实际工作中可能会遇到需要封装WebSocket的场景，以下基于ts做了WebSocket的封装，包括心跳机制和重连

1、封装逻辑如下，新建ws文件：

// websocket 封装
enum ConnectionState { // websocket 连接状态
  'CONNECTING',
  'CONNECTED',
  'DISCONNECTED',
  'RECONNECTING',
  'RECOVERY'
}

export interface ConnectCallback {
  (): void;
}

export interface MessageCallback {
  (response: MessageEvent<any>): void;
}

export interface ErrorCallback {
  (error: Error): void;
}

export interface RequestCallback {
  (error: Error): void;
}

export interface SuccessResponseCallback {
  (response: string): void;
}

export interface ErrorResponseCallback {
  (response: string): void;
}

export interface StateChangeCallback {
  (prevState: ConnectionState, state: ConnectionState, retryCount: number): void;
}

export interface ReconnectFailedCallback {
  (): void;
}

export interface HeartCheck {
  timeout: number;
  serverTimeout: number;
  timer: number;
  serverTimer: number;
  reset: () => {};
  start: () => void;
}

export interface Options {
  wsUrl: string;
  extendInfo: Record<string, any>;
  onConnect: ConnectCallback; // websocket连接成功回调
  onMessage: MessageCallback; // 客户接收服务端数据回调
  one rror: ErrorCallback; // 通信发送错误回调
  onWsStateChange?: StateChangeCallback;// websocket连接状态改变的回调
  onWsReconnectFailed?: ReconnectFailedCallback; // 重连失败回调
}

export class WS{
  private wsUrl: string;
  private extendInfo: Record<string, any>;
  private options: Options;
  // websocket
  private wsSocket: WebSocket | null;
  public state: ConnectionState; // websocket当前连接状态
  public prevState: ConnectionState; // websocket上次连接状态
  private currentId: number; // 用于和请求消息匹配
  private times: number; // 超时时间 默认 60s
  private timer: any; // 延迟器
  public autoReconnected: boolean; // 自动重连标志
  private lockReconnect: boolean; // 避免重复连接
  private heartCheck: HeartCheck; // 心跳检测机制
  private retryTimerId: number;// 重连定时器对象
  private retryCount: number; // 当前重连次数
  private maxRetryCount: number; // 最大重连次数
  private pendingRequests: Array<string>; // websocket 未准备好时发送消息，缓存消息的队列
  private successCallbacks: Map<string, SuccessResponseCallback>; // 发送消息，请求成功回调
  private errorCallbacks: Map<string, ErrorResponseCallback>; // 发送消息，请求失败回调

  public constructor (options: Options) {
    console.log('options', options);
    this.wsUrl = options.wsUrl;
    this.options = options;
    this.wsSocket = null;
    this.state = ConnectionState.DISCONNECTED;
    this.prevState = ConnectionState.DISCONNECTED;
    this.times = 60 * 1000;
    this.timer = null;
    this.autoReconnected = true;
    this.lockReconnect = false;
    this.retryTimerId = null;
    this.retryCount = 0;
    this.maxRetryCount = 30;
    this.pendingRequests = new Array<string>();
    this.successCallbacks = new Map<string, SuccessResponseCallback>();
    this.errorCallbacks = new Map<string, ErrorResponseCallback>();
    this.heartCheck = this.initHeartCheck();
  }

  // 建立websocket连接
  public connect (): boolean {
    if (!this.wsUrl && !this.wsUrl.length) {
      console.error('Websocket url is empty!')
      return false
    }
    this.retryTimerId && clearTimeout(this.retryTimerId);
    if (!this.wsSocket) {
      this.updateWsStateChange(ConnectionState.CONNECTING);
      this.wsSocket = new WebSocket(this.wsUrl);
      if (this.wsUrl) {
        this.wsSocket.onopen = this.onConnect.bind(this);
        this.wsSocket.onmessage = this.onWsMessage.bind(this);
        this.wsSocket.onclose = this.onWsClose.bind(this);
        this.wsSocket.onerror = this.onWsError.bind(this);
      }
    }
    return !!this.wsSocket;
  }

  // 发送消息
  public sendMessage (request: any, successCb?: SuccessResponseCallback, errorCb?: ErrorResponseCallback): boolean {
    if (!request) {
      return false;
    }
    const requestId = `${Date.now()}_${this.currentId++}`;
    request.id = requestId; // 需要后端配合，将id带给后端,后端再将id返回
    if (this.wsSocket) {
      if (this.wsSocket.readyState < 1) {
        this.pendingRequests.push(request);
      } else {
        console.log(`send message:${request}`);
        this.send(request)
      }
      if (successCb && typeof successCb === 'function') {
        this.successCallbacks.set(requestId, successCb)
      }
      if (errorCb && typeof errorCb === 'function') {
        this.errorCallbacks.set(requestId, errorCb)
      }
      return true;
    }
  }

  /**
  * 发送消息 + 超时机制
  */
  private send (msg: any) {
    this.wsSocket.send(msg);
    if (!this.timer) {
      this.timer = setTimeout(() => {
        console.error('server timeout')
        if (typeof this.options.onError === 'function') {
          this.options.onError(new Error('server timeout'));
        }
      }, this.times);
    }
  }

  // 关闭websocket连接
  public close (): void {
    if (this.wsSocket) {
      this.autoReconnected = false;
      this.pendingRequests = [];
      this.wsSocket.close();
      this.resetWs();
      this.timer && clearTimeout(this.timer);
      this.retryTimerId && clearTimeout(this.retryTimerId)
      this.heartCheck.reset();
      this.heartCheck = null;
      console.log('close websocket')
    }
  }

  // 清空websocket的绑定事件
  private resetWs () {
    if (this.wsSocket) {
      this.wsSocket.onmessage = null;
      this.wsSocket.onclose = null;
      this.wsSocket.onerror = null;
      this.wsSocket.onopen = null;
      this.wsSocket = null;
    }
  }

  // 连接建立成功时触发
  private onConnect () {
    this.heartCheck.start();
    this.prevState = this.state;
    if (this.retryTimerId) {
      clearTimeout(this.retryTimerId);
      this.retryTimerId = null;
      this.updateWsStateChange(ConnectionState.RECOVERY);
    } else {
      this.updateWsStateChange(ConnectionState.CONNECTED);
    }
    if (this.pendingRequests.length) {
      let req: string | undefined;
      while ((req = this.pendingRequests.shift())) {
        this.send(req)
      }
    }
  }

  // 收到服务端消息时触发
  private onWsMessage (event: MessageEvent<any>) {
    this.heartCheck && this.heartCheck.start();
    this.timer && clearTimeout(this.timer);
    this.timer = null;
    const message = JSON.parse(event.data);
    if (message === 'pong') return
    if (typeof this.options.onMessage === 'function') {
      this.options.onMessage(event)
    }
    // 根据消息id判断是否要执行请求回调
    const id = message.id;
    if (id) {
      const successCb = this.successCallbacks.get(id);
      const errorCb = this.errorCallbacks.get(id);
      // 成功
      successCb(message);
      // 失败
      errorCb(message);
      this.successCallbacks.delete(id);
      this.errorCallbacks.delete(id);
    }
  }

  // 连接关闭时触发
  private onWsClose (event) {
    console.log('onWsClose', event);
    this.heartCheck && this.heartCheck.reset();
    this.pendingRequests = [];
    this.updateWsStateChange(ConnectionState.DISCONNECTED);
  }

  // 连接报错时触发
  private onWsError (event) {
    console.log('onWsError', event);
    this.heartCheck && this.heartCheck.reset();
    this.pendingRequests = [];
    this.updateWsStateChange(ConnectionState.DISCONNECTED);
  }

  /**
  * 更新websocket连接状态
  */
  private updateWsStateChange (state: ConnectionState) {
    this.prevState = this.state;
    this.state = state;
    if (this.options.onWsStateChange && typeof this.options.onWsStateChange === 'function') {
      this.prevState !== this.state && this.options.onWsStateChange(this.prevState, this.state, this.retryCount);
    }
  }

  /**
   * websocket 重连
   */
  public reconnect () {
    if (!this.autoReconnected && this.lockReconnect) return
    this.lockReconnect = true;
    this.retryTimerId && clearTimeout(this.retryTimerId)
    if (this.retryCount < this.maxRetryCount) {
      this.retryTimerId = window.setTimeout(() => {
        this.retryCount++;
        console.log(`${new Date().toLocaleString()} Try to reconnect, count: ${this.retryCount}`);
        this.updateWsStateChange(ConnectionState.RECONNECTING);
        this.resetWs();
        this.connect();
        this.lockReconnect = false;
      }, this.getReconnectDelay(this.retryCount));
    } else {
      console.warn(`SDK has tried reconnect signal channel for ${this.maxRetryCount} times, but all failed. please check your network`);
      if (this.options.onWsReconnectFailed) {
        this.options.onWsReconnectFailed()
      }
    }
  }

  /**
   * 心跳机制
   */
  private initHeartCheck (): HeartCheck {
    // eslint-disable-next-line @typescript-eslint/no-this-alias
    const that = this;
    return {
      timeout: 2 * 1000, // 每2s向服务端发送一次消息
      serverTimeout: 10 * 1000, // 10s收不到服务端消息算超时
      timer: null,
      serverTimer: null,
      reset () { // 心跳检测重置
        clearTimeout(this.timer);
        clearTimeout(this.serverTimer);
        this.timer = null;
        this.serverTimer = null;
        return this;
      },
      start () { // 心跳检测启动
        this.reset();
        this.timer = setTimeout(() => {
          that.wsSocket.send('ping'); // 定时向服务端发送消息
          if (!this.serverTimer) {
            this.serverTimer = setTimeout(() => {
              console.log(new Date().toLocaleString(), 'not received pong, close the websocket');
              that.reconnect(); // 重连
            }, this.serverTimeout);
          }
        }, this.timeout);
      }
    }
  }

  /**
    * 获取重连的定时时间
    * @param count 重连次数
    */
  private getReconnectDelay (count: number) {
    const t = Math.round(count / 2) + 1;
    return t > 6 ? 13 * 1000 : 3 * 1000;
  }
}

 2、调用

import {WS} from './ws.ts'

const ws = new WS({
    wsUrl:'xxxxx',
    onConnect:()=>{
      // websocket的连接成功业务逻辑
    }，
   onMessage:(data)=>{
     // 接收到客户端数据逻辑
   }，
   one rror:(err)=>{
    // 通信发送错误逻辑
   }
}) 

// 建立连接
ws.connect();

// 发送消息
ws.sendMessage(msg,(data)=>{
  //成功回调
},(err)=>{
  // 失败回调
})