refactor(electron): reduce the number of listeners for ipcmain

toeverything · Aug 5, 2024 · c4e431f · c4e431f
1 parent 9f57ed5
commit c4e431f
Show file tree

Hide file tree

Showing 5 changed files with 104 additions and 58 deletions.
diff --git a/packages/frontend/electron/src/main/events.ts b/packages/frontend/electron/src/main/events.ts
@@ -1,5 +1,6 @@
 import { app, BrowserWindow, WebContentsView } from 'electron';
 
+import { AFFINE_EVENT_CHANNEL_NAME } from '../shared/type';
 import { applicationMenuEvents } from './application-menu';
 import { logger } from './logger';
 import { sharedStorageEvents } from './shared-storage';
@@ -39,14 +40,14 @@ export function registerEvents() {
             return;
           }
           // .webContents could be undefined if the window is destroyed
-          win.webContents?.send(chan, ...args);
+          win.webContents?.send(AFFINE_EVENT_CHANNEL_NAME, chan, ...args);
           win.contentView.children.forEach(child => {
             if (
               child instanceof WebContentsView &&
               child.webContents &&
               !child.webContents.isDestroyed()
             ) {
-              child.webContents?.send(chan, ...args);
+              child.webContents?.send(AFFINE_EVENT_CHANNEL_NAME, chan, ...args);
             }
           });
         });

diff --git a/packages/frontend/electron/src/main/handlers.ts b/packages/frontend/electron/src/main/handlers.ts
@@ -1,5 +1,6 @@
 import { ipcMain } from 'electron';
 
+import { AFFINE_API_CHANNEL_NAME } from '../shared/type';
 import { clipboardHandlers } from './clipboard';
 import { configStorageHandlers } from './config-storage';
 import { exportHandlers } from './export';
@@ -31,45 +32,60 @@ export const allHandlers = {
 };
 
 export const registerHandlers = () => {
-  // TODO(@Peng): listen to namespace instead of individual event types
-  ipcMain.setMaxListeners(100);
-  for (const [namespace, namespaceHandlers] of Object.entries(allHandlers)) {
-    for (const [key, handler] of Object.entries(namespaceHandlers)) {
-      const chan = `${namespace}:${key}`;
-      const wrapper = async (
-        e: Electron.IpcMainInvokeEvent,
-        ...args: any[]
-      ) => {
-        const start = performance.now();
-        try {
-          const result = await handler(e, ...args);
-          logger.debug(
-            '[ipc-api]',
-            chan,
-            args.filter(
-              arg => typeof arg !== 'function' && typeof arg !== 'object'
-            ),
-            '-',
-            (performance.now() - start).toFixed(2),
-            'ms'
-          );
-          return result;
-        } catch (error) {
-          logger.error('[ipc]', chan, error);
-        }
-      };
-      // for ipcRenderer.invoke
-      ipcMain.handle(chan, wrapper);
-      // for ipcRenderer.sendSync
-      ipcMain.on(chan, (e, ...args) => {
-        wrapper(e, ...args)
-          .then(ret => {
-            e.returnValue = ret;
-          })
-          .catch(() => {
-            // never throw
-          });
-      });
+  const handleIpcMessage = async (
+    e: Electron.IpcMainInvokeEvent,
+    ...args: any[]
+  ) => {
+    // args[0] is the `{namespace:key}`
+    if (typeof args[0] !== 'string') {
+      logger.error('invalid ipc message', args);
+      return;
+    }
+    const channel = args[0] as string;
+    const [namespace, key] = channel.split(':');
+
+    if (!namespace || !key) {
+      logger.error('invalid ipc message', args);
+      return;
+    }
+
+    // @ts-expect-error - ignore here
+    const handler = allHandlers[namespace]?.[key];
+
+    if (!handler) {
+      logger.error('handler not found for ', args[0]);
+      return;
     }
-  }
+
+    const start = Date.now();
+    const realArgs = args.slice(1);
+    const result = await handler(e, ...realArgs);
+
+    logger.debug(
+      '[ipc-api]',
+      channel,
+      realArgs.filter(
+        arg => typeof arg !== 'function' && typeof arg !== 'object'
+      ),
+      '-',
+      Date.now() - start,
+      'ms'
+    );
+
+    return result;
+  };
+
+  ipcMain.handle(AFFINE_API_CHANNEL_NAME, async (e, ...args: any[]) => {
+    return handleIpcMessage(e, ...args);
+  });
+
+  ipcMain.on(AFFINE_API_CHANNEL_NAME, (e, ...args: any[]) => {
+    handleIpcMessage(e, ...args)
+      .then(ret => {
+        e.returnValue = ret;
+      })
+      .catch(() => {
+        // never throw
+      });
+  });
 };
diff --git a/packages/frontend/electron/src/preload/electron-api.ts b/packages/frontend/electron/src/preload/electron-api.ts
@@ -5,10 +5,12 @@ import { ipcRenderer } from 'electron';
 import { Subject } from 'rxjs';
 import { z } from 'zod';
 
-import type {
-  ExposedMeta,
-  HelperToRenderer,
-  RendererToHelper,
+import {
+  AFFINE_API_CHANNEL_NAME,
+  AFFINE_EVENT_CHANNEL_NAME,
+  type ExposedMeta,
+  type HelperToRenderer,
+  type RendererToHelper,
 } from '../shared/type';
 
 export function getElectronAPIs() {
@@ -65,7 +67,11 @@ function getMainAPIs() {
         return [
           name,
           (...args: any[]) => {
-            return ipcRenderer.invoke(channel, ...args);
+            return ipcRenderer.invoke(
+              AFFINE_API_CHANNEL_NAME,
+              channel,
+              ...args
+            );
           },
         ];
       });
@@ -79,24 +85,40 @@ function getMainAPIs() {
   const events: any = (() => {
     const { events: eventsMeta } = meta;
 
-    // NOTE: ui may try to listen to a lot of the same events, so we increase the limit...
-    ipcRenderer.setMaxListeners(100);
+    // channel -> callback[]
+    const listenersMap = new Map<string, ((...args: any[]) => void)[]>();
+
+    ipcRenderer.on(AFFINE_EVENT_CHANNEL_NAME, (_event, channel, ...args) => {
+      if (typeof channel !== 'string') {
+        console.error('invalid ipc event', channel);
+        return;
+      }
+      const [namespace, name] = channel.split(':');
+      if (!namespace || !name) {
+        console.error('invalid ipc event', channel);
+        return;
+      }
+      const listeners = listenersMap.get(channel) ?? [];
+      listeners.forEach(listener => listener(...args));
+    });
 
     const all = eventsMeta.map(([namespace, eventNames]) => {
       const namespaceEvents = eventNames.map(name => {
         const channel = `${namespace}:${name}`;
         return [
           name,
           (callback: (...args: any[]) => void) => {
-            const fn: (
-              event: Electron.IpcRendererEvent,
-              ...args: any[]
-            ) => void = (_, ...args) => {
-              callback(...args);
-            };
-            ipcRenderer.on(channel, fn);
+            listenersMap.set(channel, [
+              ...(listenersMap.get(channel) ?? []),
+              callback,
+            ]);
+
             return () => {
-              ipcRenderer.off(channel, fn);
+              const listeners = listenersMap.get(channel) ?? [];
+              const index = listeners.indexOf(callback);
+              if (index !== -1) {
+                listeners.splice(index, 1);
+              }
             };
           },
         ];

diff --git a/packages/frontend/electron/src/preload/shared-storage.ts b/packages/frontend/electron/src/preload/shared-storage.ts
@@ -1,15 +1,19 @@
 import { MemoryMemento } from '@toeverything/infra';
 import { ipcRenderer } from 'electron';
 
+import { AFFINE_API_CHANNEL_NAME } from '../shared/type';
+
 const initialGlobalState = ipcRenderer.sendSync(
+  AFFINE_API_CHANNEL_NAME,
   'sharedStorage:getAllGlobalState'
 );
 const initialGlobalCache = ipcRenderer.sendSync(
+  AFFINE_API_CHANNEL_NAME,
   'sharedStorage:getAllGlobalCache'
 );
 
 function invokeWithCatch(key: string, ...args: any[]) {
-  ipcRenderer.invoke(key, ...args).catch(err => {
+  ipcRenderer.invoke(AFFINE_API_CHANNEL_NAME, key, ...args).catch(err => {
     console.error(`Failed to invoke ${key}`, err);
   });
 }

diff --git a/packages/frontend/electron/src/shared/type.ts b/packages/frontend/electron/src/shared/type.ts
@@ -27,3 +27,6 @@ export type MainToHelper = Pick<
   | 'showItemInFolder'
   | 'getPath'
 >;
+
+export const AFFINE_API_CHANNEL_NAME = 'affine-ipc-api';
+export const AFFINE_EVENT_CHANNEL_NAME = 'affine-ipc-event';