refs #125 Implement ping/pong to keep connection in misskey

src/misskey/web_socket.ts

@@ -1,4 +1,5 @@
 import WS from 'ws'
+import moment, { Moment } from 'moment'
 import { v4 as uuid } from 'uuid'
 import { EventEmitter } from 'events'
 import { WebSocketInterface } from '../megalodon'
@@ -16,6 +17,9 @@ export default class WebSocket extends EventEmitter implements WebSocketInterfac
   private _connectionClosed: boolean
   private _client: WS | null = null
   private _channelID: string
+  private _pongReceivedTimestamp: Moment
+  private _heartbeatInterval: number = 60000
+  private _pongWaiting: boolean = false
 
   constructor(
     url: string,
@@ -36,6 +40,7 @@ export default class WebSocket extends EventEmitter implements WebSocketInterfac
     this._reconnectCurrentAttempts = 0
     this._connectionClosed = false
     this._channelID = uuid()
+    this._pongReceivedTimestamp = moment()
   }
 
   public start() {
@@ -188,14 +193,19 @@ export default class WebSocket extends EventEmitter implements WebSocketInterfac
       }
     })
     client.on('pong', () => {
+      this._pongWaiting = false
       this.emit('pong', {})
+      this._pongReceivedTimestamp = moment()
+      // It is required to anonymous function since get this scope in checkAlive.
+      setTimeout(() => this._checkAlive(this._pongReceivedTimestamp), this._heartbeatInterval)
     })
     client.on('open', () => {
       this.emit('connect', {})
       this._channel()
+      // Call first ping event.
       setTimeout(() => {
-        client.pong('')
-      }, 1000)
+        client.ping('')
+      }, 10000)
     })
     client.on('message', (data: WS.Data) => {
       this.parser.parse(data, this._channelID)
@@ -219,6 +229,29 @@ export default class WebSocket extends EventEmitter implements WebSocketInterfac
       this.emit('parser-error', err)
     })
   }
+
+  /**
+   * Call ping and wait to pong.
+   */
+  private _checkAlive(timestamp: Moment) {
+    const now: Moment = moment()
+    // Block multiple calling, if multiple pong event occur.
+    // It the duration is less than interval, through ping.
+    if (now.diff(timestamp) > this._heartbeatInterval - 1000 && !this._connectionClosed) {
+      // Skip ping when client is connecting.
+      // https://github.com/websockets/ws/blob/7.2.1/lib/websocket.js#L289
+      if (this._client && this._client.readyState !== WS.CONNECTING) {
+        this._pongWaiting = true
+        this._client.ping('')
+        setTimeout(() => {
+          if (this._pongWaiting) {
+            this._pongWaiting = false
+            this._reconnect()
+          }
+        }, 10000)
+      }
+    }
+  }
 }
 
 /**