Skip to content

Commit faf86a1

Browse files
icecrasher321icecrasher321
committed
fix(execution): preserve pending subscriptions through reconnects
1 parent 836afae commit faf86a1

2 files changed

Lines changed: 145 additions & 43 deletions

File tree

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

Lines changed: 108 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -124,29 +124,111 @@ describe('ExecutionSignalHub', () => {
124124
expect(mockSubscribe).toHaveBeenCalledWith('execution:signal:execution-new', 'execution:cancel')
125125
})
126126

127-
it.each(['error', 'end'])(
128-
'rejects readiness waiters on %s and allows a fresh attempt',
129-
async (event) => {
130-
connection.status = 'connect'
131-
const hub = getExecutionSignalHub()
132-
const handler = vi.fn()
133-
const subscription = hub.subscribe('execution-1', handler)
134-
const rejected = expect(subscription).rejects.toThrow('Execution signal subscription failed:')
127+
it('keeps waiting through recoverable connection errors', async () => {
128+
connection.status = 'connecting'
129+
const hub = getExecutionSignalHub()
130+
const handler = vi.fn()
131+
const subscription = hub.subscribe('execution-1', handler)
132+
const settled = vi.fn()
133+
void subscription.then(settled, settled)
135134

136-
connection.status = 'end'
137-
connection.client?.emit(event, new Error('connection failed'))
138-
await rejected
139-
expect(mockSubscribe).not.toHaveBeenCalled()
140-
expect(connection.client?.listenerCount('ready')).toBe(1)
141-
expect(connection.client?.listenerCount('error')).toBe(1)
142-
expect(connection.client?.listenerCount('end')).toBe(0)
135+
connection.client?.emit('error', new Error('ECONNREFUSED'))
136+
connection.status = 'reconnecting'
137+
connection.client?.emit('close')
138+
await Promise.resolve()
139+
expect(settled).not.toHaveBeenCalled()
140+
expect(mockSubscribe).not.toHaveBeenCalled()
143141

144-
connection.status = 'ready'
145-
connection.client?.emit('ready')
146-
await hub.subscribe('execution-1', handler)
147-
expect(mockSubscribe).toHaveBeenCalledOnce()
148-
}
149-
)
142+
connection.status = 'ready'
143+
connection.client?.emit('ready')
144+
await subscription
145+
expect(mockSubscribe).toHaveBeenCalledOnce()
146+
connection.client?.emit('message', 'execution:signal:execution-1', 'cancelled')
147+
expect(handler).toHaveBeenCalledWith('cancelled')
148+
})
149+
150+
it('rejects readiness waiters when the subscriber stops reconnecting', async () => {
151+
connection.status = 'connect'
152+
const hub = getExecutionSignalHub()
153+
const subscription = hub.subscribe('execution-1', vi.fn())
154+
const rejected = expect(subscription).rejects.toThrow('Redis subscriber connection ended')
155+
156+
connection.status = 'end'
157+
connection.client?.emit('end')
158+
await rejected
159+
expect(mockSubscribe).not.toHaveBeenCalled()
160+
expect(connection.client?.listenerCount('ready')).toBe(1)
161+
expect(connection.client?.listenerCount('error')).toBe(1)
162+
expect(connection.client?.listenerCount('end')).toBe(0)
163+
})
164+
165+
it('keeps a new channel independent of an existing channel reconnect failure', async () => {
166+
const hub = getExecutionSignalHub()
167+
connection.client?.emit('ready')
168+
const existingHandler = vi.fn()
169+
await hub.subscribe('execution-existing', existingHandler)
170+
mockSubscribe.mockClear()
171+
connection.status = 'reconnecting'
172+
connection.client?.emit('close')
173+
const newHandler = vi.fn()
174+
const subscription = hub.subscribe('execution-new', newHandler)
175+
let rejectReconnect!: (error: Error) => void
176+
let acknowledgeNew!: (count: number) => void
177+
mockSubscribe.mockReturnValueOnce(
178+
new Promise<number>((_resolve, reject) => {
179+
rejectReconnect = reject
180+
})
181+
)
182+
mockSubscribe.mockReturnValueOnce(
183+
new Promise<number>((resolve) => {
184+
acknowledgeNew = resolve
185+
})
186+
)
187+
188+
connection.status = 'ready'
189+
connection.client?.emit('ready')
190+
await vi.waitFor(() => expect(mockSubscribe).toHaveBeenCalledTimes(2))
191+
expect(mockSubscribe).toHaveBeenNthCalledWith(
192+
1,
193+
'execution:signal:execution-existing',
194+
'execution:cancel'
195+
)
196+
expect(mockSubscribe).toHaveBeenNthCalledWith(
197+
2,
198+
'execution:signal:execution-new',
199+
'execution:cancel'
200+
)
201+
202+
rejectReconnect(new Error('Command timed out'))
203+
await vi.waitFor(() => expect(existingHandler).toHaveBeenCalledWith('unavailable'))
204+
expect(newHandler).not.toHaveBeenCalled()
205+
acknowledgeNew(3)
206+
await subscription
207+
connection.client?.emit('message', 'execution:signal:execution-new', 'cancelled')
208+
expect(newHandler).toHaveBeenCalledExactlyOnceWith('cancelled')
209+
})
210+
211+
it('preserves the pending acknowledgement when Redis reconnects before it arrives', async () => {
212+
const hub = getExecutionSignalHub()
213+
connection.client?.emit('ready')
214+
let acknowledge!: (count: number) => void
215+
mockSubscribe.mockReturnValueOnce(
216+
new Promise<number>((resolve) => {
217+
acknowledge = resolve
218+
})
219+
)
220+
const handler = vi.fn()
221+
const subscription = hub.subscribe('execution-new', handler)
222+
connection.status = 'reconnecting'
223+
connection.client?.emit('close')
224+
connection.status = 'ready'
225+
connection.client?.emit('ready')
226+
227+
expect(mockSubscribe).toHaveBeenCalledOnce()
228+
acknowledge(2)
229+
await subscription
230+
expect(handler).not.toHaveBeenCalled()
231+
})
150232

151233
it('bounds the readiness wait and removes failed handlers before a later ready event', async () => {
152234
vi.useFakeTimers()
@@ -159,7 +241,11 @@ describe('ExecutionSignalHub', () => {
159241
'Timed out waiting for Redis subscriber readiness'
160242
)
161243

162-
await Promise.all([rejected, vi.advanceTimersByTimeAsync(5000)])
244+
const timeout = vi.advanceTimersByTimeAsync(4000).then(() => {
245+
connection.client?.emit('error', new Error('ECONNREFUSED'))
246+
return vi.advanceTimersByTimeAsync(1000)
247+
})
248+
await Promise.all([rejected, timeout])
163249
expect(mockSubscribe).not.toHaveBeenCalled()
164250
expect(connection.client?.listenerCount('ready')).toBe(1)
165251
expect(connection.client?.listenerCount('error')).toBe(1)

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

Lines changed: 37 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@ export const LEGACY_EXECUTION_CANCEL_CHANNEL = 'execution:cancel'
1212
export type ExecutionSignalReason = 'event' | 'cancelled' | 'reconnected' | 'unavailable'
1313
export type ExecutionSignalHandler = (reason: ExecutionSignalReason) => void
1414

15+
interface ChannelSubscription {
16+
ready: Promise<void>
17+
acknowledged: boolean
18+
}
19+
1520
export interface ExecutionSignalHub {
1621
subscribe(executionId: string, handler: ExecutionSignalHandler): Promise<() => void>
1722
}
@@ -23,7 +28,7 @@ export function getExecutionSignalChannel(executionId: string): string {
2328
class RedisExecutionSignalHub implements ExecutionSignalHub {
2429
private readonly subscriber: Redis
2530
private readonly handlers = new Map<string, Set<ExecutionSignalHandler>>()
26-
private readonly subscriptionReady = new Map<string, Promise<void>>()
31+
private readonly subscriptions = new Map<string, ChannelSubscription>()
2732
private connectionReady: Promise<void> | undefined
2833
private connectedOnce = false
2934

@@ -70,16 +75,16 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
7075
}
7176
channelHandlers.add(handler)
7277

73-
let ready = this.subscriptionReady.get(channel)
74-
if (!ready) {
75-
ready = this.subscribeChannels(channel, LEGACY_EXECUTION_CANCEL_CHANNEL)
76-
this.subscriptionReady.set(channel, ready)
78+
let subscription = this.subscriptions.get(channel)
79+
if (!subscription) {
80+
subscription = this.createSubscription([channel])
81+
this.subscriptions.set(channel, subscription)
7782
}
7883
try {
79-
await ready
84+
await subscription.ready
8085
} catch (error) {
81-
if (this.subscriptionReady.get(channel) === ready) {
82-
this.subscriptionReady.delete(channel)
86+
if (this.subscriptions.get(channel) === subscription) {
87+
this.subscriptions.delete(channel)
8388
}
8489
channelHandlers.delete(handler)
8590
if (channelHandlers.size === 0) this.handlers.delete(channel)
@@ -94,7 +99,7 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
9499
current.delete(handler)
95100
if (current.size > 0) return
96101
this.handlers.delete(channel)
97-
this.subscriptionReady.delete(channel)
102+
this.subscriptions.delete(channel)
98103
void this.subscriber.unsubscribe(channel).catch((error) => {
99104
logger.warn('Execution signal unsubscribe failed', {
100105
channel,
@@ -104,6 +109,16 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
104109
}
105110
}
106111

112+
private createSubscription(channels: string[]): ChannelSubscription {
113+
const subscription: ChannelSubscription = {
114+
acknowledged: false,
115+
ready: this.subscribeChannels(...channels, LEGACY_EXECUTION_CANCEL_CHANNEL).then(() => {
116+
subscription.acknowledged = true
117+
}),
118+
}
119+
return subscription
120+
}
121+
107122
/**
108123
* ioredis can send SUBSCRIBE during its handshake because Redis permits it
109124
* while loading. Wait until the handshake's INFO completes before entering
@@ -126,24 +141,22 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
126141
const cleanup = () => {
127142
clearTimeout(timeout)
128143
this.subscriber.removeListener('ready', onReady)
129-
this.subscriber.removeListener('error', onError)
130144
this.subscriber.removeListener('end', onEnd)
131145
}
132146
const onReady = () => {
133147
cleanup()
134148
resolve()
135149
}
136-
const onError = (error: Error) => {
150+
const fail = (error: Error) => {
137151
cleanup()
138152
reject(error)
139153
}
140-
const onEnd = () => onError(new Error('Redis subscriber connection ended'))
154+
const onEnd = () => fail(new Error('Redis subscriber connection ended'))
141155
const timeout = setTimeout(
142-
() => onError(new Error('Timed out waiting for Redis subscriber readiness')),
156+
() => fail(new Error('Timed out waiting for Redis subscriber readiness')),
143157
SUBSCRIBER_TIMEOUT_MS
144158
)
145159
this.subscriber.once('ready', onReady)
146-
this.subscriber.once('error', onError)
147160
this.subscriber.once('end', onEnd)
148161
}).finally(() => {
149162
this.connectionReady = undefined
@@ -156,15 +169,18 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
156169
this.connectedOnce = true
157170
if (!reconnect || this.handlers.size === 0) return
158171

159-
const channels = [...this.handlers.keys()]
160-
const ready = this.subscribeChannels(...channels, LEGACY_EXECUTION_CANCEL_CHANNEL)
172+
const channels = [...this.handlers.keys()].filter(
173+
(channel) => this.subscriptions.get(channel)?.acknowledged
174+
)
175+
if (channels.length === 0) return
176+
const subscription = this.createSubscription(channels)
161177
for (const channel of channels) {
162-
if (this.handlers.has(channel)) this.subscriptionReady.set(channel, ready)
178+
if (this.handlers.has(channel)) this.subscriptions.set(channel, subscription)
163179
}
164180
try {
165-
await ready
181+
await subscription.ready
166182
for (const channel of channels) {
167-
if (this.handlers.has(channel) && this.subscriptionReady.get(channel) === ready) {
183+
if (this.handlers.has(channel) && this.subscriptions.get(channel) === subscription) {
168184
this.dispatch(channel, 'reconnected')
169185
} else if (!this.handlers.has(channel)) {
170186
void this.subscriber.unsubscribe(channel)
@@ -173,8 +189,8 @@ class RedisExecutionSignalHub implements ExecutionSignalHub {
173189
} catch (error) {
174190
logger.error('Execution signal resubscription failed', { error: toError(error).message })
175191
for (const channel of channels) {
176-
if (this.subscriptionReady.get(channel) !== ready) continue
177-
this.subscriptionReady.delete(channel)
192+
if (this.subscriptions.get(channel) !== subscription) continue
193+
this.subscriptions.delete(channel)
178194
this.dispatch(channel, 'unavailable')
179195
}
180196
}

0 commit comments

Comments
 (0)