TS websockets with heartbeats

285 阅读1分钟
const DURATION = 20000;
const HEART_CHECK_MSG = { ping: 'HeartBeat' };

export interface WebsocketParams {
  url: string;
  authTokenUrl?: string;
  onOpen?: (e?: Event) => void;
  onMessage?: (e?: MessageEvent) => void;
  onError?: (e?: CloseEvent) => void;
  onClose?: (e?: Event) => void;
}

class WS {
  private websocket: WebSocket;
  private heartCheck: HeartCheck;
  private params: WebsocketParams;

  constructor(params: WebsocketParams) {
    this.params = params;
    this.init();
  }
  private async init() {
    const url = await this.getUrl();
    this.websocket = new WebSocket(url);
    this.heartCheck = new HeartCheck(this.websocket);
    this.websocket.onopen = () => {
      this.heartCheck.start();
    };
    this.websocket.onmessage = () => {
      this.heartCheck.reset().start();
    };
    this.websocket.onclose = () => {
      this.reconnect();
    };
    this.websocket.onerror = () => {
      this.reconnect();
    };
  }
  private reconnect() {
    this.heartCheck.reset();
    this.init();
  }
  private async getUrl() {
    const authToken = await this.tryGetWSToken();
    return encodeURI(
      authToken ? `${this.params.url}?token=${authToken}` : this.params.url
    );
  }

  private async tryGetWSToken() {
    if (this.params.authTokenUrl) {
      const resp = await fetch(this.params.authTokenUrl);
      const respJson = await resp.json();
      if (respJson) {
        const token = respJson['token'];
        if (token) {
          return token as string;
        }
      }
    }
  }
}

class HeartCheck {
  private duration: number;
  private websocket: WebSocket;
  private timeout;
  private serverTimeout;
  constructor(websocket: WebSocket, duration = DURATION) {
    this.duration = duration;
    this.websocket = websocket;
  }
  start() {
    this.timeout = setTimeout(() => {
      this.websocket.send(JSON.stringify(HEART_CHECK_MSG));
      this.serverTimeout = setTimeout(() => {
        this.websocket.close(); // will trigger onclose
      }, this.duration);
    }, this.duration);
  }
  reset() {
    this.serverTimeout && clearTimeout(this.serverTimeout);
    this.timeout && clearTimeout(this.timeout);
    return this;
  }
}

// Example
new WS({ url: 'url'})