feat: reconnect when websocket closed

packages/web/package.json

@@ -11,10 +11,10 @@
     "clsx": "^1.2.1",
     "copy-to-clipboard": "^3.3.3",
     "detect-port": "^1.5.1",
+    "eventemitter3": "^5.0.1",
     "fast-deep-equal": "^3.1.3",
     "immer": "^9.0.21",
     "mime": "^3.0.0",
-    "mitt": "^3.0.0",
     "nanoid": "^4.0.2",
     "normalize.css": "^8.0.1",
     "react-transition-group": "^4.4.5",

packages/web/src/model/default.ts

@@ -1,6 +1,7 @@
 import { AppState } from './types';
 
 export const initialState: AppState = {
+  online: false,
   userInfo: null,
   userInfoDict: {},
   users: [],

packages/web/src/model/effects.ts

@@ -2,25 +2,21 @@ import { nanoid } from 'nanoid';
 import { MAX_LONG_TEXT_LENGTH, MAX_TEXT_LENGTH } from '#/constants';
 import { addPreLongText, uploadFile } from '#/services/file';
 import { fetchUserInfo, fetchUsers, updateUserInfo } from '#/services/user';
-import { createWs, parseWebSocketMessageBody } from '#/utils/ws';
+import { WS } from '#/utils/ws';
 import { createDefineEffectFor } from '#/utils/zustand';
 import {
-  MailSendDetailed, MailType, User, WebSocketServerMessage,
+  MailSendDetailed, MailType, User,
 } from '../types';
 import type { UseAppStoreExtended } from '.';
 
 const defineEffect = createDefineEffectFor<UseAppStoreExtended>();
 
 export const connect = defineEffect((store) => {
-  const ws = createWs();
+  const ws = new WS();
 
   store.ws = ws;
 
-  ws.instance.addEventListener('message', (event) => {
-    console.log('message:', event);
-    const message = parseWebSocketMessageBody<WebSocketServerMessage>(event.data as string);
-    console.log('emit', message);
-    if (!message) return;
+  ws.on('message', (message) => {
     switch (message.type) {
       case 'users': {
         store.reducers.updateUsers(message.content);
@@ -36,19 +32,24 @@ export const connect = defineEffect((store) => {
     }
   });
 
-  ws.instance.addEventListener('open', () => {
+  ws.on('open', () => {
     void Promise.all([
       fetchUsers(),
       fetchUserInfo(),
     ]).then(([users, userInfo]) => {
       store.reducers.updateUsers(users);
       store.reducers.updateUserInfo(userInfo);
     });
+    store.reducers.turnOnline();
+  });
+
+  ws.on('close', () => {
+    store.reducers.turnOffline();
   });
 });
 
 export const disconnect = defineEffect((store) => {
-  store.ws?.instance.close();
+  store.ws?.close();
   store.ws = null;
 });
 

packages/web/src/model/index.ts

@@ -1,7 +1,7 @@
 import { create } from 'zustand';
 import { persist } from 'zustand/middleware';
 import { bindEffects, bindMutateReducers } from '#/utils/zustand';
-import { WS } from '#/types';
+import { WS } from '#/utils/ws';
 import { initialState } from './default';
 import { AppState } from './types';
 import {
@@ -15,6 +15,8 @@ import {
   exitChat,
   exitMyProfile,
   clearUnreadCount,
+  turnOnline,
+  turnOffline,
 } from './reducers';
 import {
   connect, disconnect, modifyUserInfo, sendMessage,
@@ -51,6 +53,8 @@ export const reducers = {
     exitChat,
     exitMyProfile,
     clearUnreadCount,
+    turnOnline,
+    turnOffline,
   }, useAppStoreBase),
 };
 

packages/web/src/model/reducers.ts

@@ -100,3 +100,11 @@ export const clearUnreadCount = defineMutateReducer((draft, userId: string) => {
     channel.unreadCount = 0;
   }
 });
+
+export const turnOnline = defineMutateReducer((draft) => {
+  draft.online = true;
+});
+
+export const turnOffline = defineMutateReducer((draft) => {
+  draft.online = false;
+});

packages/web/src/model/types.ts

@@ -15,6 +15,7 @@ export type UserInfoDict = ReadonlyRecord<string, User | undefined>;
 export type ChatChannelDict = ReadonlyRecord<string, ChatChannel | undefined>;
 
 export type AppState = Readonly<{
+  online: boolean;
   userInfo: User | null;
   /// a userId -> user map, including offline users
   userInfoDict: ReadonlyRecord<string, User | undefined>;

packages/web/src/types/ws.ts

@@ -27,11 +27,3 @@ export type WebSocketClientMessage = GetMessageType<WebSocketClientMessageMap>;
 export type ClientMessageType = keyof WebSocketClientMessageMap;
 
 export type ServerMessageType = keyof WebSocketServerMessageMap;
-
-export interface WS {
-  instance: WebSocket;
-  sendMessage: <T extends ClientMessageType>(
-    type: T,
-    content: WebSocketClientMessageMap[T]
-  ) => void;
-}

packages/web/src/utils/ws.ts

@@ -1,7 +1,7 @@
+import { EventEmitter } from 'eventemitter3';
 import {
   ClientMessageType,
   ServerMessageType,
-  WS,
   WebSocketClientMessage,
   WebSocketClientMessageMap,
   WebSocketServerMessage,
@@ -41,15 +41,92 @@ export function parseWebSocketMessageBody<T extends WebSocketServerMessage | Web
   return null;
 }
 
-export function createWs(): WS {
-  const ws = new WebSocket(`ws://${window.location.host}/ws`);
+interface Events {
+  message: WebSocketServerMessage;
+  open: void;
+  close: void;
+}
+
+export class WS extends EventEmitter<keyof Events> {
+  private status : 'lost' | 'connecting' | 'connected' | 'closed' = 'lost';
+
+  private instance?: WebSocket;
+
+  private sendingMessageQueue: string[] = [];
+
+  constructor() {
+    super();
+
+    this.connect();
+  }
+
+  on<T extends keyof Events>(type: T, listener: (payload: Events[T]) => void): this {
+    return super.on(type, listener);
+  }
+
+  emit<T extends keyof Events>(type: T, payload: Events[T]): boolean {
+    return super.emit(type, payload);
+  }
 
-  const sendMessage: WS['sendMessage'] = (type, content) => {
-    ws.send(createWebSocketMessageBody(type, content));
-  };
+  sendMessage<T extends ClientMessageType>(
+    type: T,
+    content: WebSocketClientMessageMap[T],
+  ) {
+    const message = createWebSocketMessageBody(type, content);
+    if (this.instance && this.status === 'connected') {
+      this.instance.send(message);
+    } else {
+      this.sendingMessageQueue.push(message);
+    }
+  }
+
+  private flush() {
+    this.sendingMessageQueue.forEach((message) => {
+      this.instance?.send(message);
+    });
+    this.sendingMessageQueue = [];
+  }
+
+  private connect() {
+    if (this.status === 'connecting') {
+      return;
+    }
+    const ws = new WebSocket(`ws://${window.location.host}/ws`);
+    this.instance = ws;
+    this.status = 'connecting';
+
+    ws.addEventListener('open', () => {
+      if (ws !== this.instance) return;
+      this.status = 'connected';
+      this.flush();
+      this.emit('open', undefined);
+    });
+
+    ws.addEventListener('close', () => {
+      if (ws !== this.instance) return;
+      this.status = 'lost';
+      this.emit('close', undefined);
 
-  return {
-    instance: ws,
-    sendMessage,
-  };
+      setTimeout(() => {
+        if (this.status !== 'closed') {
+          this.connect();
+        }
+      }, 2000);
+    });
+
+    ws.addEventListener('message', (event) => {
+      if (ws !== this.instance) return;
+      console.log('message:', event);
+      const message = parseWebSocketMessageBody<WebSocketServerMessage>(event.data as string);
+      console.log('emit', message);
+      if (!message) return;
+
+      this.emit('message', message);
+    });
+  }
+
+  close() {
+    this.status = 'closed';
+    this.instance?.close();
+  }
 }