fix(execution): wait for Redis readiness before subscribing

Author: icecrasher321

apps/sim/lib/execution/execution-signal.test.ts

@@ -1,20 +1,28 @@
 /**
  * @vitest-environment node
  */
+import { EventEmitter } from 'node:events'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { listeners, mockRedisUrl, mockSubscribe, mockUnsubscribe } = vi.hoisted(() => ({
-  listeners: new Map<string, (...args: unknown[]) => void>(),
+const { connection, mockRedisUrl, mockSubscribe, mockUnsubscribe } = vi.hoisted(() => ({
+  connection: {
+    status: 'ready',
+    client: undefined as EventEmitter | undefined,
+  },
   mockRedisUrl: { value: 'redis://localhost:6379' as string | undefined },
   mockSubscribe: vi.fn(),
   mockUnsubscribe: vi.fn(),
 }))
 
 vi.mock('ioredis', () => ({
-  default: class {
-    on(event: string, handler: (...args: unknown[]) => void) {
-      listeners.set(event, handler)
-      return this
+  default: class extends EventEmitter {
+    constructor() {
+      super()
+      connection.client = this
+    }
+
+    get status() {
+      return connection.status
     }
 
     subscribe = mockSubscribe
@@ -35,14 +43,140 @@ import {
 describe('ExecutionSignalHub', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    listeners.clear()
+    connection.status = 'ready'
+    connection.client = undefined
     mockSubscribe.mockResolvedValue(1)
     mockUnsubscribe.mockResolvedValue(0)
     mockRedisUrl.value = 'redis://localhost:6379'
     const signalGlobal = globalThis as typeof globalThis & { _executionSignalHub?: unknown }
     signalGlobal._executionSignalHub = undefined
   })
 
+  it.each(['connecting', 'connect', 'reconnecting'])(
+    'waits for Redis readiness while %s before subscribing concurrent execution channels',
+    async (status) => {
+      connection.status = status
+      const hub = getExecutionSignalHub()
+      const subscriptions = Array.from({ length: 12 }, (_, index) =>
+        hub.subscribe(`execution-${index}`, vi.fn())
+      )
+
+      await Promise.resolve()
+      expect(mockSubscribe).not.toHaveBeenCalled()
+      expect(connection.client?.listenerCount('ready')).toBeLessThanOrEqual(2)
+
+      connection.status = 'ready'
+      connection.client?.emit('ready')
+      await Promise.all(subscriptions)
+
+      expect(mockSubscribe).toHaveBeenCalledTimes(12)
+      expect(connection.client?.listenerCount('ready')).toBe(1)
+      expect(connection.client?.listenerCount('error')).toBe(1)
+      expect(connection.client?.listenerCount('end')).toBe(0)
+    }
+  )
+
+  it('rechecks readiness when the connection closes before waiting subscriptions resume', async () => {
+    connection.status = 'connect'
+    const hub = getExecutionSignalHub()
+    const subscription = hub.subscribe('execution-1', vi.fn())
+
+    connection.status = 'ready'
+    connection.client?.emit('ready')
+    connection.status = 'connect'
+    connection.client?.emit('close')
+    await vi.waitFor(() => expect(connection.client?.listenerCount('ready')).toBe(2))
+    expect(mockSubscribe).not.toHaveBeenCalled()
+
+    connection.status = 'ready'
+    connection.client?.emit('ready')
+    await subscription
+    expect(mockSubscribe).toHaveBeenCalled()
+  })
+
+  it('rejects immediately when the subscriber has already ended', async () => {
+    connection.status = 'end'
+
+    await expect(getExecutionSignalHub().subscribe('execution-1', vi.fn())).rejects.toThrow(
+      'Redis subscriber connection ended'
+    )
+    expect(mockSubscribe).not.toHaveBeenCalled()
+    expect(connection.client?.listenerCount('ready')).toBe(1)
+  })
+
+  it('does not subscribe new channels while Redis is reconnecting', async () => {
+    const hub = getExecutionSignalHub()
+    connection.client?.emit('ready')
+    const handler = vi.fn()
+    await hub.subscribe('execution-existing', handler)
+    mockSubscribe.mockClear()
+    connection.status = 'connect'
+    connection.client?.emit('close')
+    const subscription = hub.subscribe('execution-new', vi.fn())
+
+    await Promise.resolve()
+    expect(mockSubscribe).not.toHaveBeenCalled()
+
+    connection.status = 'ready'
+    connection.client?.emit('ready')
+    await subscription
+    await vi.waitFor(() => expect(handler).toHaveBeenCalledWith('reconnected'))
+    expect(mockSubscribe).toHaveBeenCalledWith('execution:signal:execution-new', 'execution:cancel')
+  })
+
+  it.each(['error', 'end'])(
+    'rejects readiness waiters on %s and allows a fresh attempt',
+    async (event) => {
+      connection.status = 'connect'
+      const hub = getExecutionSignalHub()
+      const handler = vi.fn()
+      const subscription = hub.subscribe('execution-1', handler)
+      const rejected = expect(subscription).rejects.toThrow('Execution signal subscription failed:')
+
+      connection.status = 'end'
+      connection.client?.emit(event, new Error('connection failed'))
+      await rejected
+      expect(mockSubscribe).not.toHaveBeenCalled()
+      expect(connection.client?.listenerCount('ready')).toBe(1)
+      expect(connection.client?.listenerCount('error')).toBe(1)
+      expect(connection.client?.listenerCount('end')).toBe(0)
+
+      connection.status = 'ready'
+      connection.client?.emit('ready')
+      await hub.subscribe('execution-1', handler)
+      expect(mockSubscribe).toHaveBeenCalledOnce()
+    }
+  )
+
+  it('bounds the readiness wait and removes failed handlers before a later ready event', async () => {
+    vi.useFakeTimers()
+    try {
+      connection.status = 'connect'
+      const hub = getExecutionSignalHub()
+      const handler = vi.fn()
+      const subscription = hub.subscribe('execution-1', handler)
+      const rejected = expect(subscription).rejects.toThrow(
+        'Timed out waiting for Redis subscriber readiness'
+      )
+
+      await Promise.all([rejected, vi.advanceTimersByTimeAsync(5000)])
+      expect(mockSubscribe).not.toHaveBeenCalled()
+      expect(connection.client?.listenerCount('ready')).toBe(1)
+      expect(connection.client?.listenerCount('error')).toBe(1)
+      expect(connection.client?.listenerCount('end')).toBe(0)
+      expect(vi.getTimerCount()).toBe(0)
+
+      connection.status = 'ready'
+      connection.client?.emit('ready')
+      connection.client?.emit('message', 'execution:signal:execution-1', 'cancelled')
+      expect(handler).not.toHaveBeenCalled()
+      await hub.subscribe('execution-1', handler)
+      expect(mockSubscribe).toHaveBeenCalledOnce()
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('waits for one shared subscription acknowledgement before resolving concurrent subscribers', async () => {
     let acknowledge: (() => void) | undefined
     mockSubscribe.mockReturnValueOnce(
@@ -69,12 +203,12 @@ describe('ExecutionSignalHub', () => {
 
   it('marks every affected subscription unavailable when reconnect acknowledgement fails', async () => {
     const hub = getExecutionSignalHub()
-    listeners.get('ready')?.()
+    connection.client?.emit('ready')
     const handler = vi.fn()
     await hub.subscribe('execution-1', handler)
     mockSubscribe.mockRejectedValueOnce(new Error('Redis unavailable'))
 
-    listeners.get('ready')?.()
+    connection.client?.emit('ready')
 
     await vi.waitFor(() => expect(handler).toHaveBeenCalledWith('unavailable'))
   })
@@ -87,7 +221,11 @@ describe('ExecutionSignalHub', () => {
     await hub.subscribe('execution-2', secondHandler)
 
     expect(mockSubscribe).toHaveBeenCalledWith('execution:signal:execution-1', 'execution:cancel')
-    listeners.get('message')?.('execution:cancel', JSON.stringify({ executionId: 'execution-1' }))
+    connection.client?.emit(
+      'message',
+      'execution:cancel',
+      JSON.stringify({ executionId: 'execution-1' })
+    )
 
     expect(firstHandler).toHaveBeenCalledWith('cancelled')
     expect(secondHandler).not.toHaveBeenCalled()
@@ -98,19 +236,20 @@ describe('ExecutionSignalHub', () => {
     const handler = vi.fn()
     await hub.subscribe('execution-1', handler)
 
-    listeners.get('message')?.(
+    connection.client?.emit(
+      'message',
       'execution:cancel',
       JSON.stringify({ executionId: 'execution-1', executionSignalPublished: true })
     )
-    listeners.get('message')?.('execution:signal:execution-1', 'cancelled')
+    connection.client?.emit('message', 'execution:signal:execution-1', 'cancelled')
 
     expect(handler).toHaveBeenCalledOnce()
     expect(handler).toHaveBeenCalledWith('cancelled')
   })
 
   it('does not deliver a stale reconnect failure to a replacement subscriber', async () => {
     const hub = getExecutionSignalHub()
-    listeners.get('ready')?.()
+    connection.client?.emit('ready')
     const oldHandler = vi.fn()
     const unsubscribeOld = await hub.subscribe('execution-1', oldHandler)
     let rejectOldReconnect!: (error: Error) => void
@@ -120,7 +259,7 @@ describe('ExecutionSignalHub', () => {
       })
     )
 
-    listeners.get('ready')?.()
+    connection.client?.emit('ready')
     unsubscribeOld()
     mockSubscribe.mockResolvedValueOnce(1)
     const replacement = vi.fn()

apps/sim/lib/execution/execution-signal.ts

@@ -6,6 +6,7 @@ import { getConfiguredRedisUrl, getRedisConnectionDefaults } from '@/lib/core/co
 
 const logger = createLogger('ExecutionSignalHub')
 const EXECUTION_SIGNAL_PREFIX = 'execution:signal:'
+const SUBSCRIBER_TIMEOUT_MS = 5000
 export const LEGACY_EXECUTION_CANCEL_CHANNEL = 'execution:cancel'
 
 export type ExecutionSignalReason = 'event' | 'cancelled' | 'reconnected' | 'unavailable'
@@ -23,12 +24,13 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
   private readonly subscriber: Redis
   private readonly handlers = new Map<string, Set<ExecutionSignalHandler>>()
   private readonly subscriptionReady = new Map<string, Promise<void>>()
+  private connectionReady: Promise<void> | undefined
   private connectedOnce = false
 
   constructor(redisUrl: string) {
     const options = {
       ...getRedisConnectionDefaults(redisUrl),
-      commandTimeout: 5000,
+      commandTimeout: SUBSCRIBER_TIMEOUT_MS,
       connectionName: 'execution-signal-hub',
       maxRetriesPerRequest: null,
       retryStrategy: (attempt: number) => Math.min(attempt * 500, 5000),
@@ -70,9 +72,7 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
 
     let ready = this.subscriptionReady.get(channel)
     if (!ready) {
-      ready = this.subscriber
-        .subscribe(channel, LEGACY_EXECUTION_CANCEL_CHANNEL)
-        .then(() => undefined)
+      ready = this.subscribeChannels(channel, LEGACY_EXECUTION_CANCEL_CHANNEL)
       this.subscriptionReady.set(channel, ready)
     }
     try {
@@ -104,15 +104,60 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
     }
   }
 
+  /**
+   * ioredis can send SUBSCRIBE during its handshake because Redis permits it
+   * while loading. Wait until the handshake's INFO completes before entering
+   * subscriber mode, including when new executions arrive during reconnect.
+   */
+  private async subscribeChannels(...channels: string[]): Promise<void> {
+    while (this.subscriber.status !== 'ready') {
+      await this.waitForConnectionReady()
+    }
+    await this.subscriber.subscribe(...channels)
+  }
+
+  private waitForConnectionReady(): Promise<void> {
+    if (this.connectionReady) return this.connectionReady
+    if (this.subscriber.status === 'end') {
+      return Promise.reject(new Error('Redis subscriber connection ended'))
+    }
+
+    this.connectionReady = new Promise<void>((resolve, reject) => {
+      const cleanup = () => {
+        clearTimeout(timeout)
+        this.subscriber.removeListener('ready', onReady)
+        this.subscriber.removeListener('error', onError)
+        this.subscriber.removeListener('end', onEnd)
+      }
+      const onReady = () => {
+        cleanup()
+        resolve()
+      }
+      const onError = (error: Error) => {
+        cleanup()
+        reject(error)
+      }
+      const onEnd = () => onError(new Error('Redis subscriber connection ended'))
+      const timeout = setTimeout(
+        () => onError(new Error('Timed out waiting for Redis subscriber readiness')),
+        SUBSCRIBER_TIMEOUT_MS
+      )
+      this.subscriber.once('ready', onReady)
+      this.subscriber.once('error', onError)
+      this.subscriber.once('end', onEnd)
+    }).finally(() => {
+      this.connectionReady = undefined
+    })
+    return this.connectionReady
+  }
+
   private async handleReady(): Promise<void> {
     const reconnect = this.connectedOnce
     this.connectedOnce = true
     if (!reconnect || this.handlers.size === 0) return
 
     const channels = [...this.handlers.keys()]
-    const ready = this.subscriber
-      .subscribe(...channels, LEGACY_EXECUTION_CANCEL_CHANNEL)
-      .then(() => undefined)
+    const ready = this.subscribeChannels(...channels, LEGACY_EXECUTION_CANCEL_CHANNEL)
     for (const channel of channels) {
       if (this.handlers.has(channel)) this.subscriptionReady.set(channel, ready)
     }