From f95f794dd5ef9639da85d19518bd332917c06f8e Mon Sep 17 00:00:00 2001
From: Shigma <1700011071@pku.edu.cn>
Date: Thu, 5 Mar 2020 18:12:12 +0800
Subject: [PATCH] feat(core): support ws reconnect

---
 packages/koishi-cli/src/worker.ts  |  3 +--
 packages/koishi-core/src/app.ts    |  3 +++
 packages/koishi-core/src/server.ts | 22 ++++++++++++++++------
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/packages/koishi-cli/src/worker.ts b/packages/koishi-cli/src/worker.ts
index 6347e45b25..af2c13d960 100644
--- a/packages/koishi-cli/src/worker.ts
+++ b/packages/koishi-cli/src/worker.ts
@@ -16,8 +16,7 @@ if (process.env.KOISHI_LOG_LEVEL !== undefined) {
 }
 
 function handleException (error: any) {
-  const message = types.isNativeError(error) ? error.stack : String(error)
-  logger.error(message, baseLogLevel)
+  logger.error(error, baseLogLevel)
   process.exit(1)
 }
 
diff --git a/packages/koishi-core/src/app.ts b/packages/koishi-core/src/app.ts
index 703f2f70a6..5fa1caf78d 100644
--- a/packages/koishi-core/src/app.ts
+++ b/packages/koishi-core/src/app.ts
@@ -20,6 +20,8 @@ export interface AppOptions {
   type?: ServerType
   database?: DatabaseConfig
   nickname?: string | string[]
+  retryTimes?: number
+  retryTimeout?: number
   maxMiddlewares?: number
   commandPrefix?: string | string[]
   quickOperationTimeout?: number
@@ -75,6 +77,7 @@ function createLeadingRE (patterns: string[], prefix = '', suffix = '') {
 
 const defaultOptions: AppOptions = {
   maxMiddlewares: 64,
+  retryTimeout: 15000,
 }
 
 export enum Status { closed, opening, open, closing }
diff --git a/packages/koishi-core/src/server.ts b/packages/koishi-core/src/server.ts
index 799664a673..b0f89b1883 100644
--- a/packages/koishi-core/src/server.ts
+++ b/packages/koishi-core/src/server.ts
@@ -7,6 +7,7 @@ import { Meta, VersionInfo, ContextType } from './meta'
 import { App } from './app'
 import { CQResponse } from './sender'
 import { format } from 'util'
+import ms from 'ms'
 
 export abstract class Server {
   public appList: App[] = []
@@ -29,7 +30,7 @@ export abstract class Server {
   }
 
   protected debug (format: any, ...params: any[]) {
-    this.app?.logger('koishi:sender').debug(format, ...params)
+    this.app?.logger('koishi:server').debug(format, ...params)
   }
 
   protected prepareMeta (data: any) {
@@ -277,6 +278,7 @@ let counter = 0
 
 export class WsClient extends Server {
   public socket: WebSocket
+  private _retryCount = 0
   private _listeners: Record<number, (response: CQResponse) => void> = {}
 
   send (data: any): Promise<CQResponse> {
@@ -290,14 +292,21 @@ export class WsClient extends Server {
   }
 
   _listen (): Promise<void> {
-    return new Promise((resolve, reject) => {
+    const connect = (resolve: () => void, reject: (reason: Error) => void) => {
       this.debug('websocket client opening')
       const headers: Record<string, string> = {}
-      const { token, server } = this.app.options
+      const { token, server, retryTimeout, retryTimes } = this.app.options
       if (token) headers.Authorization = `Bearer ${token}`
       this.socket = new WebSocket(server, { headers })
 
-      this.socket.once('error', reject)
+      this.socket.once('error', (error: Error) => {
+        if (!retryTimeout) reject(error)
+        if (this._retryCount >= retryTimes) reject(error)
+        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.socket.once('open', () => {
         this.socket.send(JSON.stringify({
@@ -306,7 +315,7 @@ export class WsClient extends Server {
         }), (error) => {
           if (error) reject(error)
         })
-
+  
         let resolved = false
         this.socket.on('message', (data) => {
           data = data.toString()
@@ -333,7 +342,8 @@ export class WsClient extends Server {
           }
         })
       })
-    })
+    }
+    return new Promise(connect)
   }
 
   _close () {