feat(typed-ipc): main process IpcListener cleanup functions

package.json

@@ -13,7 +13,8 @@
     "build:preload": "pnpm run -C packages/preload build",
     "build:ipc": "pnpm run -C packages/typed-ipc build",
     "build:utils": "pnpm run -C packages/utils build",
-    "dev": "pnpm run -C packages/playground start"
+    "dev": "pnpm run -C packages/playground start",
+    "test": "vitest"
   },
   "simple-git-hooks": {
     "pre-commit": "npx lint-staged",
@@ -49,6 +50,7 @@
     "simple-git-hooks": "^2.11.1",
     "tslib": "^2.6.2",
     "typescript": "^5.4.3",
-    "unbuild": "^2.0.0"
+    "unbuild": "^2.0.0",
+    "vitest": "^2.1.8"
   }
 }

packages/typed-ipc/src/main.test.ts

@@ -0,0 +1,122 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
+import { IpcEventMap, IpcListener } from './main'
+import { ipcMain } from 'electron'
+
+vi.mock('electron', () => ({
+  ipcMain: {
+    on: vi.fn(),
+    removeListener: vi.fn(),
+    handle: vi.fn(),
+    removeHandler: vi.fn()
+  }
+}))
+
+describe('IpcListener', () => {
+  let ipcListener: IpcListener<IpcEventMap>
+
+  beforeEach(() => {
+    ipcListener = new IpcListener()
+  })
+
+  afterEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('should register a listener and return a cleanup function', () => {
+    const channel = 'test-channel'
+    const listener1 = vi.fn()
+    const listener2 = vi.fn()
+
+    const cleanup1 = ipcListener.on(channel, listener1)
+    ipcListener.on(channel, listener2)
+
+    cleanup1()
+
+    expect(ipcMain.removeListener).toHaveBeenCalledWith(channel, listener1)
+    expect(ipcListener['listeners'].get(channel)?.has(listener1)).toBe(false)
+    expect(ipcListener['listeners'].get(channel)?.has(listener2)).toBe(true)
+  })
+
+  it('should delete the channel from the listeners map when all listeners are removed', () => {
+    const channel = 'test-channel'
+    const listener = vi.fn()
+
+    const cleanup = ipcListener.on(channel, listener)
+
+    cleanup()
+
+    expect(ipcListener['listeners'].has(channel)).toBe(false)
+  })
+
+  it('should register a handler and return a cleanup function', () => {
+    const channel1 = 'test-channel-1'
+    const channel2 = 'test-channel-2'
+    const handler1 = vi.fn()
+    const handler2 = vi.fn()
+
+    const cleanup1 = ipcListener.handle(channel1, handler1)
+    const cleanup2 = ipcListener.handle(channel2, handler2)
+
+    cleanup1()
+
+    expect(ipcMain.removeHandler).toHaveBeenCalledWith(channel1)
+    expect(ipcListener['handlers'].has(channel1)).toBe(false)
+    expect(ipcListener['handlers'].has(channel2)).toBe(true)
+
+    cleanup2()
+
+    expect(ipcMain.removeHandler).toHaveBeenCalledWith(channel2)
+    expect(ipcListener['handlers'].has(channel2)).toBe(false)
+  })
+
+  it('should dispose all listeners and handlers', () => {
+    const channel1 = 'test-channel-1'
+    const channel2 = 'test-channel-2'
+    const listener1 = vi.fn()
+    const listener2 = vi.fn()
+    const handler = vi.fn()
+
+    ipcListener.on(channel1, listener1)
+    ipcListener.on(channel2, listener2)
+    ipcListener.handle(channel1, handler)
+
+    ipcListener.dispose()
+
+    expect(ipcMain.removeListener).toHaveBeenCalledWith(channel1, listener1)
+    expect(ipcMain.removeListener).toHaveBeenCalledWith(channel2, listener2)
+    expect(ipcMain.removeHandler).toHaveBeenCalledWith(channel1)
+    expect(ipcListener['listeners'].size).toBe(0)
+    expect(ipcListener['handlers'].size).toBe(0)
+  })
+
+  it('should only dispose listeners and handlers registered by the instance', () => {
+    const channel1 = 'test-channel-1'
+    const channel2 = 'test-channel-2'
+    const listener1 = vi.fn()
+    const listener2 = vi.fn()
+    const handler = vi.fn()
+
+    ipcListener.on(channel1, listener1)
+    ipcListener.on(channel2, listener2)
+    ipcListener.handle(channel1, handler)
+
+    const otherIpcListener = new IpcListener()
+
+    const otherChannel = 'other-channel'
+    const otherListener = vi.fn()
+    const otherHandler = vi.fn()
+
+    otherIpcListener.on(otherChannel, otherListener)
+    otherIpcListener.handle(otherChannel, otherHandler)
+
+    ipcListener.dispose()
+
+    expect(ipcMain.removeListener).toHaveBeenCalledWith(channel1, listener1)
+    expect(ipcMain.removeListener).toHaveBeenCalledWith(channel2, listener2)
+    expect(ipcMain.removeHandler).toHaveBeenCalledWith(channel1)
+    expect(ipcMain.removeListener).not.toHaveBeenCalledWith(otherChannel, otherListener)
+    expect(ipcMain.removeHandler).not.toHaveBeenCalledWith(otherChannel)
+    expect(ipcListener['listeners'].size).toBe(0)
+    expect(ipcListener['handlers'].size).toBe(0)
+  })
+})

packages/typed-ipc/src/main.ts

@@ -8,42 +8,62 @@ export * from './types'
  * Typed listener for Electron `ipcMain`.
  */
 export class IpcListener<T extends IpcEventMap> {
-  private listeners: string[] = []
-  private handlers: string[] = []
+  private listeners: Map<string, Set<(...args: any[]) => void>> = new Map()
+  private handlers: Set<string> = new Set()
 
   /**
    * Listen to `channel`.
+   * @returns A function to remove the listener.
    */
   on<E extends keyof ExtractArgs<T>>(
     channel: Extract<E, string>,
     listener: (e: Electron.IpcMainEvent, ...args: ExtractArgs<T>[E]) => void | Promise<void>
-  ): void {
-    this.listeners.push(channel)
+  ): () => void {
+    let channelListeners = this.listeners.get(channel)
+    if (!channelListeners) {
+      channelListeners = new Set()
+      this.listeners.set(channel, channelListeners)
+    }
+    channelListeners.add(listener)
     ipcMain.on(channel, listener as (e: Electron.IpcMainEvent, ...args: any[]) => void)
+    return () => {
+      ipcMain.removeListener(channel, listener)
+      channelListeners.delete(listener)
+      if (channelListeners.size === 0) {
+        this.listeners.delete(channel)
+      }
+    }
   }
 
   /**
    * Handle a renderer invoke request.
+   * @returns A function to remove the handler.
    */
   handle<E extends keyof ExtractHandler<T>>(
     channel: Extract<E, string>,
     listener: (
       e: Electron.IpcMainInvokeEvent,
       ...args: Parameters<ExtractHandler<T>[E]>
     ) => ReturnType<ExtractHandler<T>[E]> | Promise<ReturnType<ExtractHandler<T>[E]>>
-  ): void {
-    this.handlers.push(channel)
+  ): () => void {
+    this.handlers.add(channel)
     ipcMain.handle(channel, listener as (event: Electron.IpcMainInvokeEvent, ...args: any[]) => any)
+    return () => {
+      ipcMain.removeHandler(channel)
+      this.handlers.delete(channel)
+    }
   }
 
   /**
-   * Dispose all listeners and handlers.
+   * Dispose all registered listeners and handlers.
    */
   dispose(): void {
-    this.listeners.forEach(c => ipcMain.removeAllListeners(c))
-    this.listeners = []
-    this.handlers.forEach(c => ipcMain.removeHandler(c))
-    this.handlers = []
+    this.listeners.forEach((listeners, channel) => {
+      listeners.forEach(listener => ipcMain.removeListener(channel, listener))
+    })
+    this.listeners.clear()
+    this.handlers.forEach(channel => ipcMain.removeHandler(channel))
+    this.handlers.clear()
   }
 }