From c854300d24cbc15288248ec875a7bb8d63516d6d Mon Sep 17 00:00:00 2001
From: Shigma <1700011071@pku.edu.cn>
Date: Fri, 6 Mar 2020 15:54:38 +0800
Subject: [PATCH] feat(core): support reconnection at runtime

---
 packages/koishi-core/src/server.ts | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/packages/koishi-core/src/server.ts b/packages/koishi-core/src/server.ts
index b0f89b1883..b72bbcf1e1 100644
--- a/packages/koishi-core/src/server.ts
+++ b/packages/koishi-core/src/server.ts
@@ -299,23 +299,33 @@ export class WsClient extends Server {
       if (token) headers.Authorization = `Bearer ${token}`
       this.socket = new WebSocket(server, { headers })
 
-      this.socket.once('error', (error: Error) => {
-        if (!retryTimeout) reject(error)
-        if (this._retryCount >= retryTimes) reject(error)
+      this.socket.on('error', error => this.debug(error))
+
+      this.socket.once('close', (code: number) => {
+        if (code === 1000) {
+          return reject(new Error('connection was closed manually'))
+        }
+
+        const message = `failed to connect to ${server}`
+        if (!retryTimeout || this._retryCount >= retryTimes) {
+          return reject(new Error(message))
+        }
+
         this._retryCount++
-        this.debug(error)
-        this.app?.logger('koishi').warn(`failed to connect to ${server}, will retry in ${ms(this.app.options.retryTimeout)}...`)
-        setTimeout(() => connect(resolve, reject), this.app.options.retryTimeout)
+        this.app?.logger('koishi').warn(`${message}, will retry in ${ms(retryTimeout)}...`)
+        setTimeout(() => connect(resolve, reject), retryTimeout)
       })
 
       this.socket.once('open', () => {
+        this._retryCount = 0
+
         this.socket.send(JSON.stringify({
           action: 'get_version_info',
           echo: -1,
         }), (error) => {
           if (error) reject(error)
         })
-  
+
         let resolved = false
         this.socket.on('message', (data) => {
           data = data.toString()
@@ -326,11 +336,13 @@ export class WsClient extends Server {
           } catch (error) {
             return reject(new Error(data))
           }
+
           if (!resolved) {
             resolved = true
             this.debug('connect to ws server:', this.app.options.server)
             resolve()
           }
+
           if ('post_type' in parsed) {
             const meta = this.prepareMeta(parsed)
             if (meta) this.dispatchMeta(meta)
@@ -347,7 +359,7 @@ export class WsClient extends Server {
   }
 
   _close () {
-    this.socket.close()
+    this.socket.close(1000)
     this.debug('websocket client closed')
   }
 }