Skip to content

1、Ws连接基础类

ts
type Timeout = NodeJS.Timeout;

export type Options = {
  reconnectLimit: number; // 最大重新连接次数
  reconnectInterval: number; // 重新连接的时间间隔
  onOpen?: () => void; // 连接打开时的回调
  onMessage?: (msg: string) => void; // 收到消息时的回调
  onClose?: () => void; // 连接关闭时的回调
  onError?: () => void; // 连接出错时的回调
  onReconnect?: (times: number) => void; // 重新连接的回调
  onReadStateChange?: (readState: WsReadyState) => void; // 连接状态变化时的回调
};

// 连接状态
export enum WsReadyState {
  CONNECTING = 0, // 等待连接
  OPEN = 1, // 连接中
  CLOSING = 2, // 关闭中
  CLOSED = 3, // 已关闭
  ERROR = 4, // 异常
}

export class Ws {
  private options: Options; // 参数
  private wsUrl: string; // 连接地址
  private client: WebSocket | null = null; // 连接实例
  private readyState: WsReadyState = WsReadyState.CONNECTING; // 连接状态
  private message: string | undefined; // 最后一次收到的消息
  private reConnectNum: number = 0; // 重新连接数
  private timer: Timeout | null = null; // 重新请求的timer

  /**
   * 构造方法
   * @param {string} url websocket 地址
   * @param {number} maxConnectNumber  最大重连次数
   */
  constructor(
    url: string,
    options: Options = { reconnectLimit: 10, reconnectInterval: 1000 }
  ) {
    this.wsUrl = url;
    this.options = options;
  }

  /**
   * 连接状态变化时的回调
   */
  public onReadStateChange() {
    this.options?.onReadStateChange &&
      this.options.onReadStateChange(this.readyState);
  }

  /**
   * 清除定时器
   */
  private clearTimer() {
    if (this.timer) {
      clearTimeout(this.timer);
      this.timer = null;
    }
  }

  public connect() {
    if (this.client == null) {
      // 手动打开时重置
      if (this.reConnectNum > this.options.reconnectLimit) {
        this.reConnectNum = 0;
      }
      this.readyState = WsReadyState.CONNECTING;

      const ws = new WebSocket(this.wsUrl);
      ws.onopen = this.onOpen.bind(this);
      ws.onerror = this.onError.bind(this);
      ws.onclose = this.onClose.bind(this);
      ws.onmessage = this.onMessage.bind(this);
      this.onReadStateChange();
      this.client = ws;
    }
  }

  private onOpen() {
    this.readyState = WsReadyState.OPEN;
    this.onReadStateChange();
    // 连接成功重置重连次数
    this.reConnectNum = 0;
    this.options?.onOpen && this.options.onOpen();
  }

  private async onError() {
    this.readyState = WsReadyState.ERROR;
    this.client = null;
    this.onReadStateChange();
    this.options?.onError && this.options.onError();
    // 重新尝试连接
    await this.reconnect();
  }

  private async onClose() {
    this.readyState = WsReadyState.CLOSING;
    this.readyState = WsReadyState.CLOSED;
    this.onReadStateChange();
    this.options?.onClose && this.options.onClose();
    this.client = null;
  }

  private onMessage(res: { data: string }) {
    const msg = res.data;
    this.message = msg;
    this.options?.onMessage && this.options.onMessage(msg);
  }

  private async reconnect() {
    this.clearTimer();
    this.timer = setTimeout(() => {
      if (this.readyState === WsReadyState.OPEN || this.client) {
        this.clearTimer();
        return;
      }
      this.reConnectNum++;
      if (this.reConnectNum <= this.options.reconnectLimit) {
        console.log(
          "第",
          this.reConnectNum,
          "次连接。最大连接次数:",
          this.options.reconnectLimit
        );
        this.connect();
        this.options?.onReconnect &&
          this.options.onReconnect(this.reConnectNum);
      }
    }, this.options.reconnectInterval);
  }

  /**
   * @description: 获取连接实例
   * @return {WebSocket | null}
   */
  public getClient(): WebSocket | null {
    return this.client;
  }

  /**
   * @description: 获取连接状态
   * @return {WsReadyState}
   */
  public getReadyState(): WsReadyState {
    return this.readyState;
  }

  /**
   * 关闭连接
   */
  public disconnect() {
    this.reConnectNum = this.options.reconnectLimit + 1;
    this.clearTimer();
    this.client?.close();
  }

  /**
   * @description: 发送websocket消息
   * @param {string} message
   * @return {*}
   */
  public sendMessage(message: string): boolean {
    if (this.client == null) return false;
    try {
      this.client.send(message);
      return true;
    } catch {
      return false;
    }
  }

  /**
   * 获取收到的最新消息
   * @returns {string | undefined}
   */
  public getMessage(): string | undefined {
    return this.message;
  }
}

2、组合式函数实现连接websocket

ts
import { ref, type Ref, watchEffect, onBeforeUnmount } from 'vue';
import { WsReadyState, Ws } from './websocket.core';
import type { Options as WsOptions } from './websocket.core';
import { notification } from 'ant-design-vue';
import { defaults } from 'lodash';

export type Options = WsOptions & {
  manual?: boolean; // 是否手动开启连接
  ready?: Ref<boolean>; // ready不为false时开启连接
  onOpen?: () => void;
};

const defaultOptions: Options = {
  reconnectLimit: 10,
  reconnectInterval: 1000,
  manual: false,
};

export default function useWebsocket(url: string, options: Partial<Options> = {}) {
  const targetOptions: Options = defaults(options, defaultOptions);
  let ws: Ws | undefined;
  const isConnected = ref(false);
  const message = ref<string>();
  const readState = ref<WsReadyState>(WsReadyState.CONNECTING);

  const init = () => {
    if (ws) return;
    console.log('useWebsocket:', url);
    ws = new Ws(url, {
      ...targetOptions,
      onMessage: (value: string) => {
        message.value = value;
        targetOptions.onMessage && targetOptions.onMessage(value);
      },
      onReadStateChange: (value: WsReadyState) => {
        readState.value = value;
        targetOptions?.onReadStateChange && targetOptions.onReadStateChange(value);
      },
      onOpen() {
        isConnected.value = true;
        options.onOpen && options.onOpen();
      },
    });
  };

  init();

  const connect = () => {
    ws?.connect();
  };

  const disconnect = () => {
    if (readState.value === WsReadyState.OPEN) {
      ws?.disconnect();
      isConnected.value = false;
    }
  };

  const sendMessage = (msg: any) => {
    ws?.sendMessage(msg);
  };

  if (targetOptions.ready) {
    watchEffect(() => {
      if (targetOptions.ready?.value) {
        if (ws) ws.disconnect();
        readState.value = WsReadyState.CONNECTING;
        connect();
      } else {
        disconnect();
      }
    });
  }

  if (!(targetOptions.manual === true)) {
    connect();
  }

  // 检查连接情况
  const checkConnect = (msg: string = '打印组件未连接,请确保打印组件已打开!') => {
    return new Promise((resolve, reject) => {
      if (readState.value === WsReadyState.OPEN) {
        resolve(true);
        return;
      } else {
        isConnected.value = false;
      }
      ws?.connect();
      setTimeout(() => {
        if (isConnected.value) {
          resolve(true);
        } else {
          notification.error({
            message: '连接提醒!',
            description: msg,
          });
          reject(false);
        }
      }, 100);
    });
  };

  onBeforeUnmount(disconnect);

  return {
    readState,
    message,
    connect,
    disconnect,
    isConnected,
    sendMessage,
    checkConnect,
  };
}