fix: clean up daemon health socket listeners (#3771)

src/main/daemon/daemon-health-socket-cleanup.test.ts

@@ -0,0 +1,82 @@
+import { EventEmitter } from 'events'
+import { mkdtempSync, rmSync, writeFileSync } from 'fs'
+import { tmpdir } from 'os'
+import { join } from 'path'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { netConnectMock } = vi.hoisted(() => ({
+  netConnectMock: vi.fn()
+}))
+
+vi.mock('net', () => ({ connect: netConnectMock }))
+
+import { healthCheckDaemon, killStaleDaemon } from './daemon-health'
+
+class FakeSocket extends EventEmitter {
+  destroy = vi.fn()
+  write = vi.fn()
+}
+
+describe('daemon health socket listener cleanup', () => {
+  let dir: string
+  let socketPath: string
+  let tokenPath: string
+
+  beforeEach(() => {
+    dir = mkdtempSync(join(tmpdir(), 'daemon-health-socket-cleanup-'))
+    socketPath = join(dir, 'daemon.sock')
+    tokenPath = join(dir, 'daemon.token')
+    writeFileSync(socketPath, '')
+    writeFileSync(tokenPath, 'token')
+    netConnectMock.mockReset()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+    rmSync(dir, { recursive: true, force: true })
+  })
+
+  it('removes health-check socket listeners after a daemon response', async () => {
+    const socket = new FakeSocket()
+    netConnectMock.mockReturnValueOnce(socket)
+
+    const result = healthCheckDaemon(socketPath, tokenPath)
+    socket.emit('connect')
+    socket.emit('data', Buffer.from('{"type":"hello","ok":true}\n{"id":"health-1","ok":true}\n'))
+
+    await expect(result).resolves.toBe(true)
+    expect(socket.listenerCount('connect')).toBe(0)
+    expect(socket.listenerCount('error')).toBe(0)
+    expect(socket.listenerCount('data')).toBe(0)
+    expect(socket.destroy).toHaveBeenCalledTimes(1)
+  })
+
+  it('removes health-check socket listeners after a timeout', async () => {
+    vi.useFakeTimers()
+    const socket = new FakeSocket()
+    netConnectMock.mockReturnValueOnce(socket)
+
+    const result = healthCheckDaemon(socketPath, tokenPath)
+    await vi.advanceTimersByTimeAsync(3_000)
+
+    await expect(result).resolves.toBe(false)
+    expect(socket.listenerCount('connect')).toBe(0)
+    expect(socket.listenerCount('error')).toBe(0)
+    expect(socket.listenerCount('data')).toBe(0)
+    expect(socket.destroy).toHaveBeenCalledTimes(1)
+  })
+
+  it('removes stale-socket probe listeners after a timeout', async () => {
+    vi.useFakeTimers()
+    const socket = new FakeSocket()
+    netConnectMock.mockReturnValueOnce(socket)
+
+    const result = killStaleDaemon(dir, socketPath, tokenPath)
+    await vi.advanceTimersByTimeAsync(500)
+
+    await expect(result).resolves.toBe(false)
+    expect(socket.listenerCount('connect')).toBe(0)
+    expect(socket.listenerCount('error')).toBe(0)
+    expect(socket.destroy).toHaveBeenCalledTimes(1)
+  })
+})

src/main/daemon/daemon-health.ts

@@ -32,19 +32,33 @@ function canConnectSocket(socketPath: string): Promise<boolean> {
       return
     }
     const sock = connect({ path: socketPath })
+    let settled = false
+    const cleanup = (): void => {
+      clearTimeout(timer)
+      sock.off('connect', onConnect)
+      sock.off('error', onError)
+    }
+    const settle = (result: boolean): void => {
+      if (settled) {
+        return
+      }
+      settled = true
+      cleanup()
+      resolve(result)
+    }
+    const onConnect = (): void => {
+      settle(true)
+      sock.destroy()
+    }
+    const onError = (): void => {
+      settle(false)
+    }
     const timer = setTimeout(() => {
+      settle(false)
       sock.destroy()
-      resolve(false)
     }, 500)
-    sock.on('connect', () => {
-      clearTimeout(timer)
-      sock.destroy()
-      resolve(true)
-    })
-    sock.on('error', () => {
-      clearTimeout(timer)
-      resolve(false)
-    })
+    sock.on('connect', onConnect)
+    sock.on('error', onError)
   })
 }
 
@@ -71,14 +85,17 @@ export function healthCheckDaemon(socketPath: string, tokenPath: string): Promis
       }
       settled = true
       clearTimeout(timer)
+      removeSocketListeners()
       sock?.destroy()
       resolve(result)
     }
-    const timer = setTimeout(() => settle(false), HEALTH_CHECK_TIMEOUT_MS)
-
-    sock = connect({ path: socketPath })
-    sock.on('error', () => settle(false))
-    sock.on('connect', () => {
+    const removeSocketListeners = (): void => {
+      sock?.off('error', onError)
+      sock?.off('connect', onConnect)
+      sock?.off('data', onData)
+    }
+    const onError = (): void => settle(false)
+    const onConnect = (): void => {
       const hello: HelloMessage = {
         type: 'hello',
         version: PROTOCOL_VERSION,
@@ -87,10 +104,8 @@ export function healthCheckDaemon(socketPath: string, tokenPath: string): Promis
         role: 'control'
       }
       sock?.write(encodeNdjson(hello))
-    })
-
-    let buffer = ''
-    sock.on('data', (chunk: Buffer) => {
+    }
+    const onData = (chunk: Buffer): void => {
       if (settled) {
         return
       }
@@ -128,7 +143,15 @@ export function healthCheckDaemon(socketPath: string, tokenPath: string): Promis
           return
         }
       }
-    })
+    }
+    const timer = setTimeout(() => settle(false), HEALTH_CHECK_TIMEOUT_MS)
+
+    sock = connect({ path: socketPath })
+    sock.on('error', onError)
+    sock.on('connect', onConnect)
+
+    let buffer = ''
+    sock.on('data', onData)
   })
 }
 
@@ -167,14 +190,17 @@ export function getMacDaemonSystemResolverHealth(
       }
       settled = true
       clearTimeout(timer)
+      removeSocketListeners()
       sock?.destroy()
       resolve(result)
     }
-    const timer = setTimeout(() => settle('unknown'), RESOLVER_HEALTH_CHECK_TIMEOUT_MS)
-
-    sock = connect({ path: socketPath })
-    sock.on('error', () => settle('unknown'))
-    sock.on('connect', () => {
+    const removeSocketListeners = (): void => {
+      sock?.off('error', onError)
+      sock?.off('connect', onConnect)
+      sock?.off('data', onData)
+    }
+    const onError = (): void => settle('unknown')
+    const onConnect = (): void => {
       const hello: HelloMessage = {
         type: 'hello',
         version: protocolVersion,
@@ -183,10 +209,8 @@ export function getMacDaemonSystemResolverHealth(
         role: 'control'
       }
       sock?.write(encodeNdjson(hello))
-    })
-
-    let buffer = ''
-    sock.on('data', (chunk: Buffer) => {
+    }
+    const onData = (chunk: Buffer): void => {
       if (settled) {
         return
       }
@@ -231,7 +255,15 @@ export function getMacDaemonSystemResolverHealth(
           return
         }
       }
-    })
+    }
+    const timer = setTimeout(() => settle('unknown'), RESOLVER_HEALTH_CHECK_TIMEOUT_MS)
+
+    sock = connect({ path: socketPath })
+    sock.on('error', onError)
+    sock.on('connect', onConnect)
+
+    let buffer = ''
+    sock.on('data', onData)
   })
 }