From 7eb78a9fbc7084c2a518d784c2e573593226fc21 Mon Sep 17 00:00:00 2001
From: Vladimir Sheremet <sleuths.slews0s@icloud.com>
Date: Mon, 8 Jan 2024 17:59:32 +0100
Subject: [PATCH] chore: add a separate class for messages

---
 packages/vite/src/client/client.ts |  2 +-
 packages/vite/src/shared/hmr.ts    | 44 ++++++++++++++++++------------
 2 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/packages/vite/src/client/client.ts b/packages/vite/src/client/client.ts
index 31db051ae2e868..c198c6efef9d79 100644
--- a/packages/vite/src/client/client.ts
+++ b/packages/vite/src/client/client.ts
@@ -160,7 +160,7 @@ async function handleMessage(payload: HMRPayload) {
   switch (payload.type) {
     case 'connected':
       console.debug(`[vite] connected.`)
-      hmrClient.flushMessageQueue()
+      hmrClient.messenger.flush()
       // proxy(nginx, docker) hmr ws maybe caused timeout,
       // so send ping package let ws keep alive.
       setInterval(() => {
diff --git a/packages/vite/src/shared/hmr.ts b/packages/vite/src/shared/hmr.ts
index 4748441f6852d9..9617ddb3a949a4 100644
--- a/packages/vite/src/shared/hmr.ts
+++ b/packages/vite/src/shared/hmr.ts
@@ -15,7 +15,7 @@ interface HotCallback {
   fn: (modules: Array<ModuleNamespace | undefined>) => void
 }
 
-interface HMRConnection {
+export interface HMRConnection {
   /**
    * Checked before sending messages to the client.
    */
@@ -146,7 +146,9 @@ export class HMRContext implements ViteHotContext {
   }
 
   send<T extends string>(event: T, data?: InferCustomEventPayload<T>): void {
-    this.hmrClient.sendMessage(JSON.stringify({ type: 'custom', event, data }))
+    this.hmrClient.messenger.send(
+      JSON.stringify({ type: 'custom', event, data }),
+    )
   }
 
   private acceptDeps(
@@ -165,6 +167,24 @@ export class HMRContext implements ViteHotContext {
   }
 }
 
+class HMRMessenger {
+  constructor(private connection: HMRConnection) {}
+
+  private queue: string[] = []
+
+  public send(message: string): void {
+    this.queue.push(message)
+    this.flush()
+  }
+
+  public flush(): void {
+    if (this.connection.isReady()) {
+      this.queue.forEach((msg) => this.connection.send(msg))
+      this.queue = []
+    }
+  }
+}
+
 export class HMRClient {
   public hotModulesMap = new Map<string, HotModule>()
   public disposeMap = new Map<string, (data: any) => void | Promise<void>>()
@@ -173,14 +193,16 @@ export class HMRClient {
   public customListenersMap: CustomListenersMap = new Map()
   public ctxToListenersMap = new Map<string, CustomListenersMap>()
 
-  private messageQueue: string[] = []
+  public messenger: HMRMessenger
 
   constructor(
     public logger: Console,
-    private connection: HMRConnection,
+    connection: HMRConnection,
     // This allows implementing reloading via different methods depending on the environment
     private importUpdatedModule: (update: Update) => Promise<ModuleNamespace>,
-  ) {}
+  ) {
+    this.messenger = new HMRMessenger(connection)
+  }
 
   public async notifyListeners<T extends string>(
     event: T,
@@ -206,18 +228,6 @@ export class HMRClient {
     })
   }
 
-  public sendMessage(message: string): void {
-    this.messageQueue.push(message)
-    this.flushMessageQueue()
-  }
-
-  public flushMessageQueue(): void {
-    if (this.connection.isReady()) {
-      this.messageQueue.forEach((msg) => this.connection.send(msg))
-      this.messageQueue = []
-    }
-  }
-
   protected warnFailedUpdate(err: Error, path: string | string[]): void {
     if (!err.message.includes('fetch')) {
       this.logger.error(err)