feat: isWssReused, clean eventHandlers

src/ws/adapter.ts

@@ -60,12 +60,17 @@ export class WsAdapter {
     this.server = server;
     this.manager.onConnected(async (ws) => {
       logger.info('conn', '', 'connected');
-      ws.on('close', () => {
+      const handler = () => {
         logger.info('conn', '', pc.yellow('disconnected'));
-      });
+      };
+      ws.on('close', handler);
       await this.getDts();
       await this.handleHmrMessage();
+      return () => {
+        ws.off('close', handler);
+      };
     });
+    this.manager.checkIfWssReused();
   }
   async getDts() {
     const filename = this.server.config.viteburner.dts;

src/ws/allocator.ts

@@ -1,5 +1,6 @@
 import { WebSocketServer } from 'ws';
 
+let reused = false;
 let wss: WebSocketServer | null = null;
 let wssPort = -1;
 
@@ -11,6 +12,13 @@ export function getWss(port: number) {
     }
     wssPort = port;
     wss = new WebSocketServer({ port });
+    reused = false;
+  } else {
+    reused = true;
   }
   return wss;
 }
+
+export function isWssReused() {
+  return reused;
+}

src/ws/manager.ts

@@ -1,4 +1,5 @@
 import { RawData, WebSocket, WebSocketServer } from 'ws';
+import { z } from 'zod';
 import {
   wsResponseSchema,
   PushFileParams,
@@ -15,9 +16,8 @@ import {
   calculateRamResponseSchema,
   getDefinitionFileResponseSchema,
 } from './messages';
-import { z } from 'zod';
-import { logger } from '..';
-import { getWss } from './allocator';
+import { logger } from '@/console';
+import { getWss, isWssReused } from './allocator';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export interface PromiseHolder<T = any> {
@@ -39,14 +39,15 @@ export interface WsManagerOptions {
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type Fn = (...args: any[]) => any;
+type Promisable<T> = T | Promise<T>;
 
 export class WsManager {
   options: Required<WsManagerOptions>;
   ws: WebSocket | undefined;
   wss: WebSocketServer;
   trackers: PromiseHolder[];
   nextId: number;
-  handlers: [string, Fn][];
+  unregisters: Fn[];
   constructor(options: WsManagerOptions) {
     this.options = {
       timeout: 10000,
@@ -56,7 +57,7 @@ export class WsManager {
     this.nextId = 0;
     this.ws = undefined;
     this.wss = getWss(this.options.port);
-    this.handlers = [];
+    this.unregisters = [];
     this._registerHandler();
   }
   get connected() {
@@ -78,21 +79,30 @@ export class WsManager {
     };
     this.wss.on('connection', _onConnected);
     this.wss.on('error', _onError);
-    this.handlers.push(['connection', _onConnected]);
-    this.handlers.push(['error', _onError]);
+    this.unregisters.push(() => {
+      this.wss.off('connection', _onConnected);
+      this.wss.off('error', _onError);
+    });
   }
-  onConnected(cb: (ws: WebSocket) => void) {
-    const handler = (ws: WebSocket) => {
+  onConnected(cb: (ws: WebSocket) => Promisable<Fn | void>) {
+    const handler = async (ws: WebSocket) => {
       // ensure ws is saved before any sendMessage calls
       this.ws = ws;
-      cb(ws);
+      const unregister = await cb(ws);
+      unregister && this.unregisters.push(unregister);
     };
     this.wss.on('connection', handler);
-    this.handlers.push(['connection', handler]);
+    this.unregisters.push(() => {
+      this.wss.off('connection', handler);
+    });
   }
-  onDisconneted(cb: () => void) {
-    this.wss.on('close', cb);
-    this.handlers.push(['close', cb]);
+  checkIfWssReused() {
+    if (isWssReused() && this.wss.clients.size > 0) {
+      for (const client of this.wss.clients) {
+        this.ws = client;
+      }
+      this.wss.emit('connection', this.ws);
+    }
   }
   handleMessage(response: RawData) {
     const parsed = wsResponseSchema.parse(JSON.parse(response.toString()));
@@ -209,8 +219,6 @@ export class WsManager {
   }
   close() {
     // remove all handlers
-    this.handlers.forEach(([event, handler]) => {
-      this.wss.off(event, handler);
-    });
+    this.unregisters.forEach((unregister) => unregister());
   }
 }