From c15bbca848046403b85f7c225773be438afbe45d Mon Sep 17 00:00:00 2001 From: Lorent Lempereur Date: Fri, 31 Jul 2026 16:02:24 +0200 Subject: [PATCH] fix(core): Improve task broker resilience when a runner dies (#35298) --- .../@n8n/config/src/configs/runners.config.ts | 9 + packages/@n8n/config/test/config.test.ts | 1 + .../errors/task-request-timeout.error.ts | 13 +- .../__tests__/task-broker-ws-server.test.ts | 283 +++++-- .../__tests__/task-broker.service.test.ts | 708 +++++++++++++++++- .../task-requester-accept-timeout.error.ts | 7 + .../errors/task-runner-unreachable.error.ts | 7 + .../task-broker/task-broker-ws-server.ts | 145 ++-- .../task-broker/task-broker.service.ts | 266 +++++-- .../task-managers/task-requester.ts | 11 +- 10 files changed, 1251 insertions(+), 199 deletions(-) create mode 100644 packages/cli/src/task-runners/task-broker/errors/task-requester-accept-timeout.error.ts create mode 100644 packages/cli/src/task-runners/task-broker/errors/task-runner-unreachable.error.ts diff --git a/packages/@n8n/config/src/configs/runners.config.ts b/packages/@n8n/config/src/configs/runners.config.ts index 48b75c0c5d5..acd6677b7d8 100644 --- a/packages/@n8n/config/src/configs/runners.config.ts +++ b/packages/@n8n/config/src/configs/runners.config.ts @@ -62,6 +62,15 @@ export class TaskRunnersConfig { @Env('N8N_RUNNERS_TASK_REQUEST_TIMEOUT') taskRequestTimeout: number = 60; + /** + * How long (in seconds) the broker waits for a runner or requester to + * acknowledge a matched task before abandoning the match. + * Must be greater than 0. + * Increase on infrastructure where runners are slow to respond. + */ + @Env('N8N_RUNNERS_TASK_ACCEPT_TIMEOUT') + taskAcceptTimeout: number = 2; + /** Interval in seconds between heartbeats from runner to broker; missing heartbeats abort the task (and restart the runner in internal mode). Must be > 0. */ @Env('N8N_RUNNERS_HEARTBEAT_INTERVAL') heartbeatInterval: number = 30; diff --git a/packages/@n8n/config/test/config.test.ts b/packages/@n8n/config/test/config.test.ts index b3985fdf7c1..ce660a9009c 100644 --- a/packages/@n8n/config/test/config.test.ts +++ b/packages/@n8n/config/test/config.test.ts @@ -408,6 +408,7 @@ describe('GlobalConfig', () => { maxConcurrency: 10, taskTimeout: 300, taskRequestTimeout: 60, + taskAcceptTimeout: 2, heartbeatInterval: 30, grantTokenTtl: 30, insecureMode: false, diff --git a/packages/cli/src/task-runners/errors/task-request-timeout.error.ts b/packages/cli/src/task-runners/errors/task-request-timeout.error.ts index a7e7c4ba9d2..08a34a1e039 100644 --- a/packages/cli/src/task-runners/errors/task-request-timeout.error.ts +++ b/packages/cli/src/task-runners/errors/task-request-timeout.error.ts @@ -3,11 +3,18 @@ import { OperationalError } from 'n8n-workflow'; export class TaskRequestTimeoutError extends OperationalError { description: string; - constructor({ timeout, isSelfHosted }: { timeout: number; isSelfHosted: boolean }) { - super(`Task request timed out after ${timeout} ${timeout === 1 ? 'second' : 'seconds'}`); + constructor({ + elapsedSeconds, + isSelfHosted, + }: { + elapsedSeconds: number; + isSelfHosted: boolean; + }) { + // Keep the message static so Sentry groups all occurrences as one issue. + super('Task request timed out', { extra: { elapsedSeconds } }); const description = [ - 'Your Code node task was not matched to a runner within the timeout period. This indicates that the task runner is currently down, or not ready, or at capacity, so it cannot service your task.', + `Your Code node task was not matched to a runner within the timeout period (waited ${elapsedSeconds} ${elapsedSeconds === 1 ? 'second' : 'seconds'}). This indicates that the task runner is currently down, or not ready, or at capacity, so it cannot service your task.`, 'If you are repeatedly executing Code nodes with long-running tasks across your instance, please space them apart to give the runner time to catch up. If this does not describe your use case, please open a GitHub issue or reach out to support.', ]; diff --git a/packages/cli/src/task-runners/task-broker/__tests__/task-broker-ws-server.test.ts b/packages/cli/src/task-runners/task-broker/__tests__/task-broker-ws-server.test.ts index 5a584036139..b334821bf0a 100644 --- a/packages/cli/src/task-runners/task-broker/__tests__/task-broker-ws-server.test.ts +++ b/packages/cli/src/task-runners/task-broker/__tests__/task-broker-ws-server.test.ts @@ -4,35 +4,241 @@ import { mock } from 'vitest-mock-extended'; import type WebSocket from 'ws'; import { WsStatusCodes } from '@/constants'; +import type { DefaultTaskRunnerDisconnectAnalyzer } from '@/task-runners/default-task-runner-disconnect-analyzer'; import { TaskBrokerWsServer } from '@/task-runners/task-broker/task-broker-ws-server'; +import type { TaskBroker } from '@/task-runners/task-broker/task-broker.service'; const globalConfig = mock({ generic: { gracefulShutdownTimeout: 30 } }); +const WS_OPEN = 1; +const WS_CLOSING = 2; + +const setReadyState = (ws: WebSocket, readyState: number) => { + (ws as unknown as { readyState: number }).readyState = readyState; +}; + +const mockWs = (readyState = WS_OPEN, isAlive = true) => { + const ws = mock({ OPEN: WS_OPEN }); + setReadyState(ws, readyState); + ws.isAlive = isAlive; + return ws; +}; + +const createServer = ({ + taskBroker = mock(), + disconnectAnalyzer = mock(), + heartbeatInterval = 30, +}: { + taskBroker?: TaskBroker; + disconnectAnalyzer?: DefaultTaskRunnerDisconnectAnalyzer; + heartbeatInterval?: number; +} = {}) => + new TaskBrokerWsServer( + mock(), + taskBroker, + disconnectAnalyzer, + mock({ path: '/runners', heartbeatInterval }), + mock(), + globalConfig, + ); + +const wsMessage = (message: unknown) => Buffer.from(JSON.stringify(message)); + +const messageHandlerOf = (ws: WebSocket) => { + // `on` is overloaded, so its mock calls are unreachable through the ws types + const onCalls = (ws.on as unknown as { mock: { calls: Array<[string, unknown]> } }).mock.calls; + return onCalls.find(([event]) => event === 'message')?.[1] as ( + data: WebSocket.RawData, + ) => Promise; +}; + +const registerOverWs = async (server: TaskBrokerWsServer, id: string, ws: WebSocket) => { + server.add(id, ws); + + const onMessage = messageHandlerOf(ws); + await onMessage( + wsMessage({ type: 'runner:info', name: 'JS Task Runner', types: ['javascript'] }), + ); + + return onMessage; +}; + +const pendingAnalysis = () => { + let complete: (error: Error) => void = () => {}; + const disconnectAnalyzer = mock(); + disconnectAnalyzer.toDisconnectError.mockReturnValue( + new Promise((resolve) => { + complete = resolve; + }), + ); + return { disconnectAnalyzer, complete: (error: Error) => complete(error) }; +}; + describe('TaskBrokerWsServer', () => { describe('removeConnection', () => { it('should close with 1000 status code by default', async () => { - const server = new TaskBrokerWsServer(mock(), mock(), mock(), mock(), mock(), globalConfig); - const ws = mock(); + const server = createServer(); + const ws = mockWs(); server.runnerConnections.set('test-runner', ws); await server.removeConnection('test-runner'); expect(ws.close).toHaveBeenCalledWith(WsStatusCodes.CloseNormal); }); + + it('should stop routing to the runner before the disconnect analysis completes', async () => { + const { disconnectAnalyzer, complete } = pendingAnalysis(); + const taskBroker = mock(); + const server = createServer({ taskBroker, disconnectAnalyzer }); + const ws = mockWs(); + server.runnerConnections.set('test-runner', ws); + + const removal = server.removeConnection('test-runner'); + await Promise.resolve(); + + expect(server.runnerConnections.has('test-runner')).toBe(false); + expect(ws.close).toHaveBeenCalled(); + expect(taskBroker.deregisterRunner).not.toHaveBeenCalled(); + + complete(new Error('disconnected')); + await removal; + + expect(taskBroker.deregisterRunner).toHaveBeenCalled(); + }); + + it('should not deregister a runner that reconnects before disconnect analysis completes', async () => { + const { disconnectAnalyzer, complete } = pendingAnalysis(); + const taskBroker = mock(); + const server = createServer({ taskBroker, disconnectAnalyzer }); + const newWs = mockWs(); + server.runnerConnections.set('test-runner', mockWs()); + + const removal = server.removeConnection('test-runner'); + await Promise.resolve(); + + server.runnerConnections.set('test-runner', newWs); + complete(new Error('disconnected')); + await removal; + + expect(taskBroker.deregisterRunner).not.toHaveBeenCalled(); + expect(server.runnerConnections.get('test-runner')).toBe(newWs); + }); + + it('should fail the tasks in flight on the replaced connection when the runner reconnects', async () => { + const { disconnectAnalyzer, complete } = pendingAnalysis(); + const taskBroker = mock(); + taskBroker.getInFlightTaskIds.mockReturnValue(['task1', 'task2']); + const server = createServer({ taskBroker, disconnectAnalyzer }); + server.runnerConnections.set('test-runner', mockWs()); + + const removal = server.removeConnection('test-runner'); + await Promise.resolve(); + + server.runnerConnections.set('test-runner', mockWs()); + const disconnectError = new Error('disconnected'); + complete(disconnectError); + await removal; + + expect(taskBroker.deregisterRunner).not.toHaveBeenCalled(); + expect(taskBroker.failTasks).toHaveBeenCalledWith(['task1', 'task2'], disconnectError); + }); + + it('should ignore stale close events from replaced connections', async () => { + const taskBroker = mock(); + const server = createServer({ taskBroker }); + const staleWs = mockWs(); + const currentWs = mockWs(); + server.runnerConnections.set('test-runner', currentWs); + + await server.removeConnection('test-runner', { expectedConnection: staleWs }); + + expect(currentWs.close).not.toHaveBeenCalled(); + expect(taskBroker.deregisterRunner).not.toHaveBeenCalled(); + expect(server.runnerConnections.get('test-runner')).toBe(currentWs); + }); + }); + + describe('runner reachability', () => { + const registerAndGetReachability = async (ws: WebSocket) => { + const taskBroker = mock(); + const server = createServer({ taskBroker }); + + await registerOverWs(server, 'test-runner', ws); + + return { server, isRunnerReachable: taskBroker.registerRunner.mock.calls[0][2]! }; + }; + + it('should report the runner as unreachable once its connection stops being open', async () => { + const ws = mockWs(); + const { isRunnerReachable } = await registerAndGetReachability(ws); + + expect(isRunnerReachable()).toBe(true); + + setReadyState(ws, WS_CLOSING); + + expect(isRunnerReachable()).toBe(false); + }); + + it('should report the runner as unreachable once its connection is replaced', async () => { + const { server, isRunnerReachable } = await registerAndGetReachability(mockWs()); + + server.runnerConnections.set('test-runner', mockWs()); + + expect(isRunnerReachable()).toBe(false); + }); + }); + + describe('incoming runner messages', () => { + const offer = { + type: 'runner:taskoffer', + offerId: 'offer1', + taskType: 'javascript', + validFor: 5000, + }; + + it('should forward a message from the connection the runner is registered with', async () => { + const taskBroker = mock(); + const server = createServer({ taskBroker }); + + const onMessage = await registerOverWs(server, 'test-runner', mockWs()); + await onMessage(wsMessage(offer)); + + expect(taskBroker.onRunnerMessage).toHaveBeenCalledWith('test-runner', offer); + }); + + it('should drop a message buffered on a connection that was replaced', async () => { + const taskBroker = mock(); + const server = createServer({ taskBroker }); + + const onStaleMessage = await registerOverWs(server, 'test-runner', mockWs()); + await registerOverWs(server, 'test-runner', mockWs()); + + await onStaleMessage(wsMessage(offer)); + + expect(taskBroker.onRunnerMessage).not.toHaveBeenCalled(); + }); }); describe('heartbeat timer', () => { + const DEAD = false; + + const runHeartbeatCheck = async (...connections: WebSocket[]) => { + vi.useFakeTimers(); + const server = createServer(); + + connections.forEach((ws, i) => server.runnerConnections.set(`runner-${i}`, ws)); + + server.start(); + vi.advanceTimersByTime(30 * Time.seconds.toMilliseconds); + await Promise.resolve(); + + // Restoring real timers discards the heartbeat interval, so the server needs no stop. + vi.useRealTimers(); + }; + it('should set up heartbeat timer on server start', async () => { const setIntervalSpy = vi.spyOn(global, 'setInterval'); - - const server = new TaskBrokerWsServer( - mock(), - mock(), - mock(), - mock({ path: '/runners', heartbeatInterval: 30 }), - mock(), - globalConfig, - ); + const server = createServer(); server.start(); @@ -47,15 +253,7 @@ describe('TaskBrokerWsServer', () => { it('should clear heartbeat timer on server stop', async () => { vi.spyOn(global, 'setInterval'); const clearIntervalSpy = vi.spyOn(global, 'clearInterval'); - - const server = new TaskBrokerWsServer( - mock(), - mock(), - mock(), - mock({ path: '/runners', heartbeatInterval: 30 }), - mock(), - globalConfig, - ); + const server = createServer(); server.start(); await server.stop(); @@ -64,37 +262,38 @@ describe('TaskBrokerWsServer', () => { }); it('should close connection with protocol error code when heartbeat check fails', async () => { - vi.useFakeTimers(); - const server = new TaskBrokerWsServer( - mock(), - mock(), - mock(), - mock({ path: '/runners', heartbeatInterval: 30 }), - mock(), - globalConfig, - ); + const ws = mockWs(WS_OPEN, DEAD); - const ws = mock(); - ws.isAlive = false; - server.runnerConnections.set('test-runner', ws); - - server.start(); - - vi.advanceTimersByTime(30 * Time.seconds.toMilliseconds); - - await Promise.resolve(); + await runHeartbeatCheck(ws); expect(ws.close).toHaveBeenCalledWith(WsStatusCodes.CloseProtocolError); + }); - await server.stop(); - vi.useRealTimers(); + it('should close every dead connection in a single check', async () => { + const deadWs1 = mockWs(WS_OPEN, DEAD); + const deadWs2 = mockWs(WS_OPEN, DEAD); + + await runHeartbeatCheck(deadWs1, deadWs2); + + expect(deadWs1.close).toHaveBeenCalledWith(WsStatusCodes.CloseProtocolError); + expect(deadWs2.close).toHaveBeenCalledWith(WsStatusCodes.CloseProtocolError); + }); + + it('should keep pinging live connections listed after a dead one', async () => { + const liveWs = mockWs(); + + await runHeartbeatCheck(mockWs(WS_OPEN, DEAD), liveWs); + + expect(liveWs.ping).toHaveBeenCalled(); + expect(liveWs.isAlive).toBe(false); + expect(liveWs.close).not.toHaveBeenCalled(); }); }); describe('sendMessage', () => { it('should work with a message containing circular references', () => { - const server = new TaskBrokerWsServer(mock(), mock(), mock(), mock(), mock(), globalConfig); - const ws = mock(); + const server = createServer(); + const ws = mockWs(); server.runnerConnections.set('test-runner', ws); const messageData: Record = {}; diff --git a/packages/cli/src/task-runners/task-broker/__tests__/task-broker.service.test.ts b/packages/cli/src/task-runners/task-broker/__tests__/task-broker.service.test.ts index b7fa53c8a73..4485347cdd8 100644 --- a/packages/cli/src/task-runners/task-broker/__tests__/task-broker.service.test.ts +++ b/packages/cli/src/task-runners/task-broker/__tests__/task-broker.service.test.ts @@ -1,14 +1,16 @@ import type { Logger } from '@n8n/backend-common'; import type { TaskRunnersConfig } from '@n8n/config'; import { Time } from '@n8n/constants'; -import type { RunnerMessage, TaskResultData } from '@n8n/task-runner'; +import type { BrokerMessage, RunnerMessage, TaskResultData } from '@n8n/task-runner'; import { type INodeTypeBaseDescription } from 'n8n-workflow'; import { mock } from 'vitest-mock-extended'; import type { TaskRunnerLifecycleEvents } from '@/task-runners/task-runner-lifecycle-events'; import { TaskRejectError } from '../errors/task-reject.error'; +import { TaskRequesterAcceptTimeoutError } from '../errors/task-requester-accept-timeout.error'; import { TaskRunnerExecutionTimeoutError } from '../errors/task-runner-execution-timeout.error'; +import { TaskRunnerUnreachableError } from '../errors/task-runner-unreachable.error'; import { TaskBroker } from '../task-broker.service'; import type { TaskOffer, TaskRequest, TaskRunner } from '../task-broker.service'; @@ -18,10 +20,34 @@ describe('TaskBroker', () => { let taskBroker: TaskBroker; beforeEach(() => { - taskBroker = new TaskBroker(mock(), mock(), mock(), mock()); + // real timeout values, so flows that arm timers never call setTimeout(NaN) + taskBroker = new TaskBroker( + mock(), + mock({ taskRequestTimeout: 60, taskTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); vi.restoreAllMocks(); }); + describe('constructor', () => { + it('should reject a non-positive task request timeout', () => { + const config = mock({ taskRequestTimeout: 0 }); + + expect(() => new TaskBroker(mock(), config, mock(), mock())).toThrowError( + 'Task request timeout must be greater than 0', + ); + }); + + it('should reject a non-positive task accept timeout', () => { + const config = mock({ taskRequestTimeout: 60, taskAcceptTimeout: 0 }); + + expect(() => new TaskBroker(mock(), config, mock(), mock())).toThrowError( + 'Task accept timeout must be greater than 0', + ); + }); + }); + describe('expireTasks', () => { it('should remove expired task offers and keep valid task offers', () => { const validOffer: TaskOffer = { @@ -115,6 +141,190 @@ describe('TaskBroker', () => { }); }); + describe('unreachable runners', () => { + beforeEach(() => { + // fake timers so the matching flows arm no lingering real timers + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + const offerFrom = (runnerId: string): TaskOffer => ({ + offerId: `offer-${runnerId}`, + runnerId, + taskType: 'taskType1', + validFor: 1000, + validUntil: createValidUntil(1000), + }); + + const requestTaskFrom = (runnerId: string, isRunnerReachable: () => boolean) => { + const messageCallback = vi.fn(); + taskBroker.registerRunner( + mock({ id: runnerId }), + messageCallback, + isRunnerReachable, + ); + messageCallback.mockClear(); + taskBroker.setPendingTaskOffers([offerFrom(runnerId)]); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + }; + + taskBroker.taskRequested(request); + + return { request, messageCallback }; + }; + + it('should not match an offer from a registered but unreachable runner', () => { + const { request, messageCallback } = requestTaskFrom('deadRunner', () => false); + + expect(taskBroker.getPendingTaskOffers()).toHaveLength(0); + expect(request.acceptInProgress).toBeUndefined(); + expect(messageCallback).not.toHaveBeenCalled(); + }); + + it('should still match an offer from a reachable runner', () => { + const { request, messageCallback } = requestTaskFrom('liveRunner', () => true); + + expect(request.acceptInProgress).toBe(true); + expect(messageCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:taskofferaccept' }), + ); + }); + + it('should keep offers from runners that are not registered', () => { + taskBroker.setPendingTaskOffers([offerFrom('unregisteredRunner')]); + + taskBroker.settleTasks(); + + expect(taskBroker.getPendingTaskOffers()).toHaveLength(1); + }); + + const setupInFlightTask = (taskId: string, runnerId: string, isReachable: boolean) => { + const runnerCallback = vi.fn(); + const requesterCallback = vi.fn(); + + taskBroker.registerRunner( + mock({ id: runnerId }), + runnerCallback, + () => isReachable, + ); + taskBroker.registerRequester('requester1', requesterCallback); + taskBroker.setTasks({ + [taskId]: { id: taskId, runnerId, requesterId: 'requester1', taskType: 'taskType1' }, + }); + runnerCallback.mockClear(); + + return { runnerCallback, requesterCallback }; + }; + + it('should fail an in-flight task instead of sending settings to an unreachable runner', async () => { + const { runnerCallback, requesterCallback } = setupInFlightTask('task1', 'deadRunner', false); + + await expect(taskBroker.sendTaskSettings('task1', {})).rejects.toThrow( + TaskRunnerUnreachableError, + ); + + expect(runnerCallback).not.toHaveBeenCalled(); + expect(requesterCallback).toHaveBeenCalledWith({ + type: 'broker:taskerror', + taskId: 'task1', + error: expect.any(TaskRunnerUnreachableError), + }); + expect(taskBroker.getTasks().get('task1')).toBeUndefined(); + + const [[taskErrorMessage]] = requesterCallback.mock.calls as [ + [{ error: TaskRunnerUnreachableError }], + ]; + expect(taskErrorMessage.error.level).toBe('warning'); + }); + + it('should fail an in-flight task instead of relaying a requester response to an unreachable runner', async () => { + const { runnerCallback, requesterCallback } = setupInFlightTask('task1', 'deadRunner', false); + + await expect(taskBroker.handleRequesterDataResponse('task1', 'req1', {})).rejects.toThrow( + TaskRunnerUnreachableError, + ); + + expect(runnerCallback).not.toHaveBeenCalled(); + expect(requesterCallback).toHaveBeenCalledWith({ + type: 'broker:taskerror', + taskId: 'task1', + error: expect.any(TaskRunnerUnreachableError), + }); + expect(taskBroker.getTasks().get('task1')).toBeUndefined(); + }); + + it('should still send settings to a reachable runner', async () => { + const { runnerCallback } = setupInFlightTask('task1', 'liveRunner', true); + + await taskBroker.sendTaskSettings('task1', {}); + + expect(runnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:tasksettings', taskId: 'task1' }), + ); + expect(taskBroker.getTasks().get('task1')).toBeDefined(); + }); + + it('should settle the accept flow gracefully when the runner dies before receiving settings', async () => { + const loggerMock = mock(); + taskBroker = new TaskBroker( + loggerMock, + mock({ taskRequestTimeout: 60, taskTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + + let isReachable = true; + const runnerCallback = vi.fn((message: BrokerMessage.ToRunner.All) => { + if (message.type === 'broker:taskofferaccept') { + taskBroker.handleRunnerAccept(message.taskId); + isReachable = false; // runner dies right after acknowledging + } + }); + const requesterCallback = vi.fn((message: BrokerMessage.ToRequester.All) => { + if (message.type === 'broker:taskready') { + taskBroker.handleRequesterAccept(message.taskId, {}); + } + }); + + taskBroker.registerRunner( + mock({ id: 'runner1' }), + runnerCallback, + () => isReachable, + ); + taskBroker.registerRequester('requester1', requesterCallback); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + }; + taskBroker.setPendingTaskRequests([request]); + + await expect(taskBroker.acceptOffer(offerFrom('runner1'), request)).resolves.toBeUndefined(); + + expect(runnerCallback).not.toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:tasksettings' }), + ); + expect(requesterCallback).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'broker:taskerror', + error: expect.any(TaskRunnerUnreachableError), + }), + ); + expect(taskBroker.getTasks().size).toBe(0); + expect(loggerMock.warn).toHaveBeenCalledWith( + expect.stringContaining('Runner (runner1) became unreachable while processing task'), + ); + }); + }); + describe('registerRequester', () => { it('should add a requester to known requesters', () => { const requesterId = 'requester1'; @@ -183,7 +393,6 @@ describe('TaskBroker', () => { // eslint-disable-next-line @typescript-eslint/no-explicit-any const failSpy = vi.spyOn(taskBroker as any, 'failTask'); - const rejectSpy = vi.spyOn(taskBroker, 'handleRunnerReject'); taskBroker.registerRunner(runner, messageCallback); taskBroker.setTasks({ @@ -194,10 +403,7 @@ describe('TaskBroker', () => { taskBroker.deregisterRunner(runnerId, error); expect(failSpy).toBeCalledWith(taskId, error); - expect(rejectSpy).toBeCalledWith( - taskId, - `The Task Runner (${runnerId}) has disconnected: error`, - ); + expect(failSpy).not.toBeCalledWith('task2', expect.anything()); }); }); @@ -490,7 +696,7 @@ describe('TaskBroker', () => { const accept = vi.fn(); const reject = vi.fn(); - taskBroker.setRunnerAcceptRejects({ [taskId]: { accept, reject } }); + taskBroker.setRunnerAcceptRejects({ [taskId]: { accept, reject, runnerId } }); taskBroker.registerRunner(mock({ id: runnerId }), vi.fn()); await taskBroker.onRunnerMessage(runnerId, message); @@ -516,7 +722,7 @@ describe('TaskBroker', () => { const accept = vi.fn(); const reject = vi.fn(); - taskBroker.setRunnerAcceptRejects({ [taskId]: { accept, reject } }); + taskBroker.setRunnerAcceptRejects({ [taskId]: { accept, reject, runnerId } }); taskBroker.registerRunner(mock({ id: runnerId }), vi.fn()); await taskBroker.onRunnerMessage(runnerId, message); @@ -896,6 +1102,59 @@ describe('TaskBroker', () => { expect(taskBroker.getTasks().get(taskId)).toBeUndefined(); }); + it('on failing a task, we should clear timeout', async () => { + const taskId = 'task1'; + const runnerId = 'runner1'; + const requesterId = 'requester1'; + + taskBroker.registerRequester(requesterId, vi.fn()); + taskBroker.setTasks({ + [taskId]: { + id: taskId, + runnerId, + requesterId, + taskType: 'test', + timeout: setTimeout(() => {}, config.taskTimeout * Time.seconds.toMilliseconds), + }, + }); + + const armedTimers = vi.getTimerCount(); + + taskBroker.deregisterRunner(runnerId, new Error('Runner died')); + await Promise.resolve(); + + expect(vi.getTimerCount()).toBe(armedTimers - 1); + expect(taskBroker.getTasks().get(taskId)).toBeUndefined(); + }); + + it('on cancelling a task, we should clear timeout', async () => { + const taskId = 'task1'; + const runnerId = 'runner1'; + const requesterId = 'requester1'; + + taskBroker.registerRunner(mock({ id: runnerId }), vi.fn()); + taskBroker.setTasks({ + [taskId]: { + id: taskId, + runnerId, + requesterId, + taskType: 'test', + timeout: setTimeout(() => {}, config.taskTimeout * Time.seconds.toMilliseconds), + }, + }); + + const armedTimers = vi.getTimerCount(); + + await taskBroker.onRequesterMessage(requesterId, { + type: 'requester:taskcancel', + taskId, + reason: 'Cancelled by requester', + }); + + expect(vi.getTimerCount()).toBe(armedTimers - 1); + expect(taskBroker.getTasks().get(taskId)).toBeUndefined(); + }); + it('[internal mode] on timeout, we should emit `runner:timed-out-during-task` event and send error to requester', async () => { vi.spyOn(global, 'clearTimeout'); @@ -983,22 +1242,38 @@ describe('TaskBroker', () => { }); describe('task runner accept timeout', () => { + const ACCEPT_TIMEOUT_MS = 2100; + + // a failing assertion must not leak fake timers into later tests + afterEach(() => { + vi.useRealTimers(); + }); + + const offerFor = (runnerId: string, offerId: string): TaskOffer => ({ + offerId, + runnerId, + taskType: 'taskType1', + validFor: 10_000, + validUntil: createValidUntil(10_000), + }); + + const expectAcceptTimeout = async (offer: TaskOffer, request: TaskRequest) => { + vi.useFakeTimers(); + const acceptPromise = taskBroker.acceptOffer(offer, request); + vi.advanceTimersByTime(ACCEPT_TIMEOUT_MS); + await acceptPromise; + }; + it('broker should handle timeout when waiting for acknowledgment of offer accept', async () => { - const runnerId = 'runner1'; - const runner = mock({ id: runnerId }); - const messageCallback = vi.fn(); const loggerMock = mock(); - taskBroker = new TaskBroker(loggerMock, mock(), mock(), mock()); - taskBroker.registerRunner(runner, messageCallback); - - const offer: TaskOffer = { - offerId: 'offer1', - runnerId, - taskType: 'taskType1', - validFor: 1000, - validUntil: createValidUntil(1000), - }; + taskBroker = new TaskBroker( + loggerMock, + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); const request: TaskRequest = { requestId: 'request1', @@ -1006,26 +1281,390 @@ describe('TaskBroker', () => { taskType: 'taskType1', }; - vi.useFakeTimers(); - - const acceptPromise = taskBroker.acceptOffer(offer, request); - - vi.advanceTimersByTime(2100); - - await acceptPromise; + await expectAcceptTimeout(offerFor('runner1', 'offer1'), request); expect(request.acceptInProgress).toBe(false); expect(loggerMock.warn).toHaveBeenCalledWith( - expect.stringContaining( - `Runner (${runnerId}) took too long to acknowledge acceptance of task`, - ), + expect.stringContaining('Runner (runner1) took too long to acknowledge acceptance of task'), ); + }); + it('should discard the unresponsive runner remaining offers and retry the request', async () => { + const deadRunnerCallback = vi.fn(); + const liveRunnerCallback = vi.fn(); + + taskBroker = new TaskBroker( + mock(), + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'deadRunner' }), deadRunnerCallback); + taskBroker.registerRunner(mock({ id: 'liveRunner' }), liveRunnerCallback); + + const matchedOffer = offerFor('deadRunner', 'deadOffer1'); + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + + taskBroker.setPendingTaskRequests([request]); + taskBroker.setPendingTaskOffers([ + offerFor('deadRunner', 'deadOffer2'), + offerFor('liveRunner', 'liveOffer1'), + ]); + + await expectAcceptTimeout(matchedOffer, request); + + expect(taskBroker.getPendingTaskOffers()).toHaveLength(0); + expect(liveRunnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:taskofferaccept', offerId: 'liveOffer1' }), + ); + }); + + it('should release the task on the runner that failed to acknowledge', async () => { + const runnerCallback = vi.fn(); + + taskBroker = new TaskBroker( + mock(), + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'runner1' }), runnerCallback); + + await expectAcceptTimeout(offerFor('runner1', 'offer1'), { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + }); + + expect(runnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:taskcancel' }), + ); + }); + + it('should stop tracking the acknowledgment that timed out', async () => { + taskBroker = new TaskBroker( + mock(), + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); + + await expectAcceptTimeout(offerFor('runner1', 'offer1'), { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + }); + + expect(taskBroker.getRunnerAcceptRejects().size).toBe(0); + }); + + it('should wait for the configured acknowledgment window before timing out', async () => { + vi.useFakeTimers(); + + const loggerMock = mock(); + taskBroker = new TaskBroker( + loggerMock, + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 5 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); + + const acceptPromise = taskBroker.acceptOffer(offerFor('runner1', 'offer1'), { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + }); + + vi.advanceTimersByTime(ACCEPT_TIMEOUT_MS); + expect(loggerMock.warn).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(3000); + await acceptPromise; + expect(loggerMock.warn).toHaveBeenCalledWith( + expect.stringContaining('took too long to acknowledge'), + ); + }); + + it('should restart the request expiry window when the runner fails to acknowledge', async () => { + vi.useFakeTimers(); + + const config = mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }); + taskBroker = new TaskBroker(mock(), config, mock(), mock()); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); + + const requesterCallback = vi.fn(); + taskBroker.registerRequester('requester1', requesterCallback); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + timeout: taskBroker['createRequestTimeout']('request1'), + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + + // most of the window elapses before the matched runner fails to acknowledge + vi.advanceTimersByTime(55_000); + const acceptPromise = taskBroker.acceptOffer(offerFor('runner1', 'offer1'), request); + vi.advanceTimersByTime(ACCEPT_TIMEOUT_MS); + await acceptPromise; + + // without the restart, the original window would have expired by now + vi.advanceTimersByTime(30_000); + expect(taskBroker.getPendingTaskRequests()).toHaveLength(1); + + vi.advanceTimersByTime(60_000); + expect(taskBroker.getPendingTaskRequests()).toHaveLength(0); + expect(requesterCallback).toHaveBeenCalledWith({ + type: 'broker:requestexpired', + requestId: 'request1', + reason: 'timeout', + }); + }); + + it('should stop restarting the expiry window after repeated acknowledgment failures', async () => { + vi.useFakeTimers(); + + const config = mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }); + taskBroker = new TaskBroker(mock(), config, mock(), mock()); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + + const timeoutAfterAcceptFailure = async (offerId: string) => { + const acceptPromise = taskBroker.acceptOffer(offerFor('runner1', offerId), request); + vi.advanceTimersByTime(ACCEPT_TIMEOUT_MS); + await acceptPromise; + return request.timeout; + }; + + let previousTimeout = await timeoutAfterAcceptFailure('offer1'); + expect(previousTimeout).toBeDefined(); + + for (const offerId of ['offer2', 'offer3']) { + const refreshedTimeout = await timeoutAfterAcceptFailure(offerId); + expect(refreshedTimeout).not.toBe(previousTimeout); + previousTimeout = refreshedTimeout; + } + + expect(await timeoutAfterAcceptFailure('offer4')).toBe(previousTimeout); + }); + + it('should not restart the expiry window when the request already expired', async () => { + taskBroker = new TaskBroker( + mock(), + mock({ taskRequestTimeout: 60, taskAcceptTimeout: 2 }), + mock(), + mock(), + ); + taskBroker.registerRunner(mock({ id: 'runner1' }), vi.fn()); + + // the request expired during acceptance, so it is no longer pending + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + + await expectAcceptTimeout(offerFor('runner1', 'offer1'), request); + + expect(request.timeoutRefreshes).toBeUndefined(); + expect(request.timeout).toBeUndefined(); + expect(vi.getTimerCount()).toBe(0); + }); + }); + + describe('acceptOffer', () => { + // a failing assertion must not leak fake timers into later tests + afterEach(() => { vi.useRealTimers(); }); + + const offerFor = (runnerId: string, offerId: string): TaskOffer => ({ + offerId, + runnerId, + taskType: 'taskType1', + validFor: 10_000, + validUntil: createValidUntil(10_000), + }); + + const acknowledgingRunnerCallback = () => + vi.fn((message: BrokerMessage.ToRunner.All) => { + if (message.type === 'broker:taskofferaccept') { + taskBroker.handleRunnerAccept(message.taskId); + } + }); + + it('should cancel the task toward the runner when the request expired during acceptance', async () => { + const runnerCallback = acknowledgingRunnerCallback(); + taskBroker.registerRunner(mock({ id: 'runner1' }), runnerCallback); + + // the request is no longer pending by the time the runner acknowledges + await taskBroker.acceptOffer(offerFor('runner1', 'offer1'), { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }); + + expect(runnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:taskcancel', reason: 'Task request expired' }), + ); + expect(taskBroker.getTasks().size).toBe(0); + }); + + it('should not leave acknowledgment timers armed after a successful accept flow', async () => { + vi.useFakeTimers(); + + const config = mock({ + taskRequestTimeout: 60, + taskTimeout: 60, + taskAcceptTimeout: 2, + }); + taskBroker = new TaskBroker(mock(), config, mock(), mock()); + + const runnerCallback = acknowledgingRunnerCallback(); + const requesterCallback = vi.fn((message: BrokerMessage.ToRequester.All) => { + if (message.type === 'broker:taskready') { + taskBroker.handleRequesterAccept(message.taskId, {}); + } + }); + taskBroker.registerRunner(mock({ id: 'runner1' }), runnerCallback); + taskBroker.registerRequester('requester1', requesterCallback); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + + await taskBroker.acceptOffer(offerFor('runner1', 'offer1'), request); + + // only the task execution timeout remains armed + expect(vi.getTimerCount()).toBe(1); + }); + + it('should cancel the task and stop tracking it when the requester fails to acknowledge', async () => { + vi.useFakeTimers(); + + const runnerCallback = acknowledgingRunnerCallback(); + taskBroker.registerRunner(mock({ id: 'runner1' }), runnerCallback); + taskBroker.registerRequester('requester1', vi.fn()); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + + const acceptPromise = taskBroker.acceptOffer(offerFor('runner1', 'offer1'), request); + await vi.advanceTimersByTimeAsync(2100); + // resolves instead of leaking the rejection through the void'ed caller + await acceptPromise; + + expect(taskBroker.getTasks().size).toBe(0); + expect(taskBroker.getRequesterAcceptRejects().size).toBe(0); + expect(runnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'broker:taskcancel', + reason: 'Requester took too long to acknowledge the task', + }), + ); + expect(new TaskRequesterAcceptTimeoutError('task1', 'requester1').level).toBe('warning'); + }); + + it('should stop restarting the expiry window after repeated deferrals', async () => { + vi.useFakeTimers(); + + const deferringRunnerCallback = vi.fn((message: BrokerMessage.ToRunner.All) => { + if (message.type === 'broker:taskofferaccept') { + taskBroker.handleRunnerDeferred(message.taskId); + } + }); + taskBroker.registerRunner(mock({ id: 'runner1' }), deferringRunnerCallback); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + + const deferOnce = async (offerId: string) => { + await taskBroker.acceptOffer(offerFor('runner1', offerId), request); + return request.timeout; + }; + + let previousTimeout = await deferOnce('offer1'); + expect(previousTimeout).toBeDefined(); + + for (const offerId of ['offer2', 'offer3']) { + const refreshedTimeout = await deferOnce(offerId); + expect(refreshedTimeout).not.toBe(previousTimeout); + previousTimeout = refreshedTimeout; + } + + expect(await deferOnce('offer4')).toBe(previousTimeout); + }); + + it('should reject an acceptance awaiting acknowledgment when its runner deregisters', async () => { + vi.useFakeTimers(); + + const liveRunnerCallback = vi.fn(); + taskBroker.registerRunner(mock({ id: 'deadRunner' }), vi.fn()); + taskBroker.registerRunner(mock({ id: 'liveRunner' }), liveRunnerCallback); + + const request: TaskRequest = { + requestId: 'request1', + requesterId: 'requester1', + taskType: 'taskType1', + acceptInProgress: true, + }; + taskBroker.setPendingTaskRequests([request]); + taskBroker.setPendingTaskOffers([offerFor('liveRunner', 'liveOffer1')]); + + const acceptPromise = taskBroker.acceptOffer(offerFor('deadRunner', 'deadOffer1'), request); + taskBroker.deregisterRunner('deadRunner', new Error('connection lost')); + + // resolves without the acknowledgment window elapsing + await acceptPromise; + + const pendingAcceptanceRunnerIds = [...taskBroker.getRunnerAcceptRejects().values()].map( + ({ runnerId }) => runnerId, + ); + expect(pendingAcceptanceRunnerIds).toEqual(['liveRunner']); + expect(liveRunnerCallback).toHaveBeenCalledWith( + expect.objectContaining({ type: 'broker:taskofferaccept', offerId: 'liveOffer1' }), + ); + }); }); describe('request timeout', () => { + // a failing assertion must not leak fake timers into later tests + afterEach(() => { + vi.useRealTimers(); + }); + it('should time out request and send `broker:requestexpired` message', async () => { vi.useFakeTimers(); @@ -1055,8 +1694,6 @@ describe('TaskBroker', () => { requestId: 'request1', reason: 'timeout', }); - - vi.useRealTimers(); }); it('should clear timeout on request matched', async () => { @@ -1104,8 +1741,6 @@ describe('TaskBroker', () => { expect(requesterCallback).not.toHaveBeenCalledWith( expect.objectContaining({ type: 'broker:requestexpired' }), ); - - vi.useRealTimers(); }); it('should reset timeout on request deferred', async () => { @@ -1159,7 +1794,6 @@ describe('TaskBroker', () => { }); handleTimeoutSpy.mockRestore(); - vi.useRealTimers(); }); }); }); diff --git a/packages/cli/src/task-runners/task-broker/errors/task-requester-accept-timeout.error.ts b/packages/cli/src/task-runners/task-broker/errors/task-requester-accept-timeout.error.ts new file mode 100644 index 00000000000..1217d9acdcd --- /dev/null +++ b/packages/cli/src/task-runners/task-broker/errors/task-requester-accept-timeout.error.ts @@ -0,0 +1,7 @@ +import { OperationalError } from 'n8n-workflow'; + +export class TaskRequesterAcceptTimeoutError extends OperationalError { + constructor(taskId: string, requesterId: string) { + super(`Requester (${requesterId}) took too long to acknowledge task (${taskId})`); + } +} diff --git a/packages/cli/src/task-runners/task-broker/errors/task-runner-unreachable.error.ts b/packages/cli/src/task-runners/task-broker/errors/task-runner-unreachable.error.ts new file mode 100644 index 00000000000..e1a0016eaa4 --- /dev/null +++ b/packages/cli/src/task-runners/task-broker/errors/task-runner-unreachable.error.ts @@ -0,0 +1,7 @@ +import { OperationalError } from 'n8n-workflow'; + +export class TaskRunnerUnreachableError extends OperationalError { + constructor(taskId: string, runnerId: string) { + super(`Runner (${runnerId}) became unreachable while processing task (${taskId})`); + } +} diff --git a/packages/cli/src/task-runners/task-broker/task-broker-ws-server.ts b/packages/cli/src/task-runners/task-broker/task-broker-ws-server.ts index 012919486da..d909acc08b5 100644 --- a/packages/cli/src/task-runners/task-broker/task-broker-ws-server.ts +++ b/packages/cli/src/task-runners/task-broker/task-broker-ws-server.ts @@ -25,6 +25,17 @@ function heartbeat(this: WebSocket) { type WsStatusCode = (typeof WsStatusCodes)[keyof typeof WsStatusCodes]; +type RemoveConnectionOptions = { + reason?: DisconnectReason; + /** Close code sent to the runner. */ + code?: WsStatusCode; + /** + * The connection the caller intends to remove. If the runner is registered + * with a different connection by now, the removal is skipped as stale. + */ + expectedConnection?: WebSocket; +}; + /** * Responsible for handling WebSocket connections with task runners * and monitoring the connection liveness @@ -55,21 +66,32 @@ export class TaskBrokerWsServer { throw new UserError('Heartbeat interval must be greater than 0'); } - this.heartbeatTimer = setInterval(() => { - for (const [runnerId, connection] of this.runnerConnections.entries()) { - if (!connection.isAlive) { - void this.removeConnection( - runnerId, - 'failed-heartbeat-check', - WsStatusCodes.CloseProtocolError, - ); - this.runnerLifecycleEvents.emit('runner:failed-heartbeat-check'); - return; - } + this.heartbeatTimer = setInterval( + () => this.checkConnectionLiveness(), + heartbeatInterval * Time.seconds.toMilliseconds, + ); + } + + private checkConnectionLiveness() { + let anyDead = false; + + for (const [runnerId, connection] of this.runnerConnections) { + if (connection.isAlive) { connection.isAlive = false; connection.ping(); + } else { + anyDead = true; + void this.removeConnection(runnerId, { + reason: 'failed-heartbeat-check', + code: WsStatusCodes.CloseProtocolError, + expectedConnection: connection, + }); } - }, heartbeatInterval * Time.seconds.toMilliseconds); + } + + if (anyDead) { + this.runnerLifecycleEvents.emit('runner:failed-heartbeat-check'); + } } async stop() { @@ -111,29 +133,29 @@ export class TaskBrokerWsServer { buffer.toString('utf8'), ) as RunnerMessage.ToBroker.All; - if (!isConnected && message.type !== 'runner:info') { - return; - } else if (!isConnected && message.type === 'runner:info') { - await this.removeConnection(id); - isConnected = true; + if (!isConnected) { + if (message.type === 'runner:info') { + await this.removeConnection(id); + isConnected = true; - this.runnerConnections.set(id, connection); + this.runnerConnections.set(id, connection); - this.taskBroker.registerRunner( - { - id, - taskTypes: message.types, - lastSeen: new Date(), - name: message.name, - }, - this.sendMessage.bind(this, id) as MessageCallback, - ); + this.taskBroker.registerRunner( + { + id, + taskTypes: message.types, + lastSeen: new Date(), + name: message.name, + }, + this.sendMessage.bind(this, id) as MessageCallback, + () => this.isRunnerReachable(id, connection), + ); - this.logger.info(`Registered runner "${message.name}" (${id}) `); - return; + this.logger.info(`Registered runner "${message.name}" (${id}) `); + } + } else if (this.isCurrentConnection(id, connection)) { + void this.taskBroker.onRunnerMessage(id, message); } - - void this.taskBroker.onRunnerMessage(id, message); } catch (error) { this.logger.error(`Couldn't parse message from runner "${id}"`, { error: error as unknown, @@ -147,7 +169,7 @@ export class TaskBrokerWsServer { connection.once('close', async () => { connection.off('pong', heartbeat); connection.off('message', onMessage); - await this.removeConnection(id); + await this.removeConnection(id, { expectedConnection: connection }); }); connection.on('message', onMessage); @@ -158,20 +180,37 @@ export class TaskBrokerWsServer { async removeConnection( id: TaskRunner['id'], - reason: DisconnectReason = 'unknown', - code: WsStatusCode = WsStatusCodes.CloseNormal, + { + reason = 'unknown', + code = WsStatusCodes.CloseNormal, + expectedConnection, + }: RemoveConnectionOptions = {}, ) { const connection = this.runnerConnections.get(id); - if (connection) { + const isStaleRemoval = + expectedConnection !== undefined && !this.isCurrentConnection(id, expectedConnection); + + if (connection && !isStaleRemoval) { + // Stop routing to the runner before the disconnect analysis, which may be slow. + this.runnerConnections.delete(id); + connection.close(code); + + const inFlightTaskIds = this.taskBroker.getInFlightTaskIds(id); + const disconnectError = await this.disconnectAnalyzer.toDisconnectError({ runnerId: id, reason, heartbeatInterval: this.taskRunnersConfig.heartbeatInterval, }); - this.taskBroker.deregisterRunner(id, disconnectError); - this.logger.debug(`Deregistered runner "${id}"`); - connection.close(code); - this.runnerConnections.delete(id); + + const hasReconnected = this.runnerConnections.has(id); + + if (hasReconnected) { + this.taskBroker.failTasks(inFlightTaskIds, disconnectError); + } else { + this.taskBroker.deregisterRunner(id, disconnectError); + this.logger.debug(`Deregistered runner "${id}"`); + } } } @@ -179,13 +218,35 @@ export class TaskBrokerWsServer { this.add(req.query.id, req.ws); } + /** + * Whether `connection` is the connection the runner is currently registered with. A + * replaced connection keeps delivering the frames it already buffered, and nothing in + * those frames distinguishes them from the current connection's, since both speak for + * the same runner id. + */ + private isCurrentConnection(id: TaskRunner['id'], connection: WebSocket) { + return this.runnerConnections.get(id) === connection; + } + + /** + * Whether the runner behind `connection` can still be sent messages. A socket leaves + * `OPEN` as soon as it starts closing, while frames it already buffered keep arriving. + */ + private isRunnerReachable(id: TaskRunner['id'], connection: WebSocket) { + return this.isCurrentConnection(id, connection) && connection.readyState === connection.OPEN; + } + private async stopConnectedRunners() { await this.drainActiveTasks(); await Promise.all( - Array.from(this.runnerConnections.keys()).map( - async (id) => - await this.removeConnection(id, 'shutting-down', WsStatusCodes.CloseGoingAway), + Array.from(this.runnerConnections.entries()).map( + async ([id, connection]) => + await this.removeConnection(id, { + reason: 'shutting-down', + code: WsStatusCodes.CloseGoingAway, + expectedConnection: connection, + }), ), ); } diff --git a/packages/cli/src/task-runners/task-broker/task-broker.service.ts b/packages/cli/src/task-runners/task-broker/task-broker.service.ts index 626459ab3ae..cde3fbcc327 100644 --- a/packages/cli/src/task-runners/task-broker/task-broker.service.ts +++ b/packages/cli/src/task-runners/task-broker/task-broker.service.ts @@ -13,8 +13,10 @@ import { nanoid } from 'nanoid'; import { TaskDeferredError } from '@/task-runners/task-broker/errors/task-deferred.error'; import { TaskRejectError } from '@/task-runners/task-broker/errors/task-reject.error'; +import { TaskRequesterAcceptTimeoutError } from '@/task-runners/task-broker/errors/task-requester-accept-timeout.error'; import { TaskRunnerAcceptTimeoutError } from '@/task-runners/task-broker/errors/task-runner-accept-timeout.error'; import { TaskRunnerExecutionTimeoutError } from '@/task-runners/task-broker/errors/task-runner-execution-timeout.error'; +import { TaskRunnerUnreachableError } from '@/task-runners/task-broker/errors/task-runner-unreachable.error'; import { TaskRunnerLifecycleEvents } from '@/task-runners/task-runner-lifecycle-events'; export interface TaskRunner { @@ -42,12 +44,15 @@ export interface TaskOffer { validUntil: bigint; } +const MAX_REQUEST_TIMEOUT_REFRESHES = 3; + export interface TaskRequest { requestId: string; requesterId: string; taskType: string; timeout?: NodeJS.Timeout; acceptInProgress?: boolean; + timeoutRefreshes?: number; } export type MessageCallback = (message: BrokerMessage.ToRunner.All) => Promise | void; @@ -61,11 +66,17 @@ type RequesterAcceptCallback = ( ) => void; type TaskRejectCallback = (reason: TaskRejectError | TaskDeferredError) => void; +export type RunnerReachabilityCheck = () => boolean; + @Service() export class TaskBroker { private knownRunners: Map< TaskRunner['id'], - { runner: TaskRunner; messageCallback: MessageCallback } + { + runner: TaskRunner; + messageCallback: MessageCallback; + isRunnerReachable: RunnerReachabilityCheck; + } > = new Map(); private requesters: Map = new Map(); @@ -79,7 +90,7 @@ export class TaskBroker { private runnerAcceptRejects: Map< Task['id'], - { accept: RunnerAcceptCallback; reject: TaskRejectCallback } + { accept: RunnerAcceptCallback; reject: TaskRejectCallback; runnerId: TaskRunner['id'] } > = new Map(); private requesterAcceptRejects: Map< @@ -103,6 +114,12 @@ export class TaskBroker { if (this.taskRunnersConfig.taskTimeout <= 0) { throw new UserError('Task timeout must be greater than 0'); } + if (this.taskRunnersConfig.taskRequestTimeout <= 0) { + throw new UserError('Task request timeout must be greater than 0'); + } + if (this.taskRunnersConfig.taskAcceptTimeout <= 0) { + throw new UserError('Task accept timeout must be greater than 0'); + } } private createRequestTimeout(requestId: string): NodeJS.Timeout { @@ -111,6 +128,23 @@ export class TaskBroker { }, this.taskRunnersConfig.taskRequestTimeout * Time.seconds.toMilliseconds); } + /** + * Restarts the request's expiry window, so a request that lost part of its window + * to a failure the requester is not responsible for gets a full window for the + * next matching attempt. Capped per request, so a runner that keeps offering but + * never acknowledges cannot postpone the expiry indefinitely. No-op for a request + * that is no longer pending, so an already expired request arms no new timer. + */ + private refreshRequestTimeout(request: TaskRequest) { + const isPending = this.pendingTaskRequests.some((r) => r.requestId === request.requestId); + const refreshes = request.timeoutRefreshes ?? 0; + if (isPending && refreshes < MAX_REQUEST_TIMEOUT_REFRESHES) { + request.timeoutRefreshes = refreshes + 1; + clearTimeout(request.timeout); + request.timeout = this.createRequestTimeout(request.requestId); + } + } + private handleRequestTimeout(requestId: string) { const requestIndex = this.pendingTaskRequests.findIndex((r) => r.requestId === requestId); if (requestIndex === -1) return; @@ -128,42 +162,76 @@ export class TaskBroker { }); } - expireTasks() { - const now = process.hrtime.bigint(); + private discardOffers(shouldDiscard: (offer: TaskOffer) => boolean) { for (let i = this.pendingTaskOffers.length - 1; i >= 0; i--) { - const offer = this.pendingTaskOffers[i]; - if (offer.validFor === -1) continue; // non-expiring offer - if (offer.validUntil < now) { + if (shouldDiscard(this.pendingTaskOffers[i])) { this.pendingTaskOffers.splice(i, 1); } } } - registerRunner(runner: TaskRunner, messageCallback: MessageCallback) { - this.knownRunners.set(runner.id, { runner, messageCallback }); + expireTasks() { + const now = process.hrtime.bigint(); + this.discardOffers(({ validFor, validUntil }) => validFor !== -1 && validUntil < now); + } + + private discardOffersFrom(runnerId: TaskRunner['id']) { + this.discardOffers((offer) => offer.runnerId === runnerId); + } + + /** + * Discards offers from runners that are registered but no longer reachable, as matching + * one burns a task request on a runner that can never acknowledge it. Offers from + * unregistered runners are kept, since deregistration already discards them. + */ + private discardOffersFromUnreachableRunners() { + this.discardOffers( + (offer) => this.knownRunners.get(offer.runnerId)?.isRunnerReachable() === false, + ); + } + + /** + * Registers a runner as eligible for task offers. `isRunnerReachable` is consulted + * before matching any of its offers, and defaults to always reachable. Registration + * outlives reachability: a transport can die while messages it already buffered are + * still being delivered, so offers can enter the pool after the runner is gone. + */ + registerRunner( + runner: TaskRunner, + messageCallback: MessageCallback, + isRunnerReachable: RunnerReachabilityCheck = () => true, + ) { + this.knownRunners.set(runner.id, { runner, messageCallback, isRunnerReachable }); void this.knownRunners.get(runner.id)!.messageCallback({ type: 'broker:runnerregistered' }); } deregisterRunner(runnerId: string, error: Error) { this.knownRunners.delete(runnerId); - // Remove any pending offers - for (let i = this.pendingTaskOffers.length - 1; i >= 0; i--) { - if (this.pendingTaskOffers[i].runnerId === runnerId) { - this.pendingTaskOffers.splice(i, 1); - } - } + this.discardOffersFrom(runnerId); - // Fail any tasks - for (const task of this.tasks.values()) { - if (task.runnerId === runnerId) { - void this.failTask(task.id, error); + for (const [taskId, acceptReject] of this.runnerAcceptRejects) { + if (acceptReject.runnerId === runnerId) { this.handleRunnerReject( - task.id, + taskId, `The Task Runner (${runnerId}) has disconnected: ${error.message}`, ); } } + + this.failTasks(this.getInFlightTaskIds(runnerId), error); + } + + getInFlightTaskIds(runnerId: TaskRunner['id']): Array { + return [...this.tasks.values()] + .filter((task) => task.runnerId === runnerId) + .map((task) => task.id); + } + + failTasks(taskIds: Array, error: Error) { + for (const taskId of taskIds) { + void this.failTask(taskId, error); + } } registerRequester(requesterId: string, messageCallback: RequesterMessageCallback) { @@ -425,32 +493,45 @@ export class TaskBroker { } } - private async cancelTask(taskId: Task['id'], reason: string) { + /** + * Stops tracking the task and disarms its execution timeout. + * @returns the task that was being tracked, or `undefined` if it was already settled. + * @remarks A task's timeout outlives the task itself unless disarmed, so the two are dropped together. + */ + private untrackTask(taskId: Task['id']) { const task = this.tasks.get(taskId); - if (!task) { - return; - } - this.tasks.delete(taskId); - await this.messageRunner(task.runnerId, { - type: 'broker:taskcancel', - taskId, - reason, - }); + if (task) { + clearTimeout(task.timeout); + this.tasks.delete(taskId); + } + + return task; + } + + private async cancelTask(taskId: Task['id'], reason: string) { + const task = this.untrackTask(taskId); + + if (task) { + await this.messageRunner(task.runnerId, { + type: 'broker:taskcancel', + taskId, + reason, + }); + } } private async failTask(taskId: Task['id'], error: Error) { - const task = this.tasks.get(taskId); - if (!task) { - return; + const task = this.untrackTask(taskId); + + if (task) { + // TODO: special message type? + await this.messageRequester(task.requesterId, { + type: 'broker:taskerror', + taskId, + error, + }); } - this.tasks.delete(taskId); - // TODO: special message type? - await this.messageRequester(task.requesterId, { - type: 'broker:taskerror', - taskId, - error, - }); } private async getRunnerOrFailTask(taskId: Task['id']): Promise { @@ -466,6 +547,11 @@ export class TaskBroker { await this.failTask(taskId, error); throw error; } + if (!runner.isRunnerReachable()) { + const error = new TaskRunnerUnreachableError(taskId, task.runnerId); + await this.failTask(taskId, error); + throw error; + } return runner.runner; } @@ -514,44 +600,44 @@ export class TaskBroker { } async taskDoneHandler(taskId: Task['id'], data: TaskResultData) { - const task = this.tasks.get(taskId); - if (!task) return; + const task = this.untrackTask(taskId); - clearTimeout(task.timeout); - - await this.requesters.get(task.requesterId)?.({ - type: 'broker:taskdone', - taskId: task.id, - data, - }); - this.tasks.delete(task.id); + if (task) { + await this.requesters.get(task.requesterId)?.({ + type: 'broker:taskdone', + taskId: task.id, + data, + }); + } } async taskErrorHandler(taskId: Task['id'], error: unknown) { - const task = this.tasks.get(taskId); - if (!task) return; + const task = this.untrackTask(taskId); - clearTimeout(task.timeout); - - await this.requesters.get(task.requesterId)?.({ - type: 'broker:taskerror', - taskId: task.id, - error, - }); - this.tasks.delete(task.id); + if (task) { + await this.requesters.get(task.requesterId)?.({ + type: 'broker:taskerror', + taskId: task.id, + error, + }); + } } async acceptOffer(offer: TaskOffer, request: TaskRequest): Promise { const taskId = nanoid(8); + let runnerAcceptTimer: NodeJS.Timeout | undefined; try { const acceptPromise = new Promise((resolve, reject) => { - this.runnerAcceptRejects.set(taskId, { accept: resolve as () => void, reject }); + this.runnerAcceptRejects.set(taskId, { + accept: resolve as () => void, + reject, + runnerId: offer.runnerId, + }); - // TODO: customisable timeout - setTimeout(() => { + runnerAcceptTimer = setTimeout(() => { reject(new TaskRunnerAcceptTimeoutError(taskId, offer.runnerId)); - }, 2000); + }, this.taskRunnersConfig.taskAcceptTimeout * Time.seconds.toMilliseconds); }); await this.messageRunner(offer.runnerId, { @@ -565,19 +651,34 @@ export class TaskBroker { request.acceptInProgress = false; if (e instanceof TaskRejectError) { this.logger.info(`Task (${taskId}) rejected by Runner with reason "${e.reason}"`); + this.settleTasks(); return; } if (e instanceof TaskDeferredError) { this.logger.debug(`Task (${taskId}) deferred until runner is ready`); - clearTimeout(request.timeout); - request.timeout = this.createRequestTimeout(request.requestId); + this.refreshRequestTimeout(request); return; } if (e instanceof TaskRunnerAcceptTimeoutError) { this.logger.warn(e.message); + this.runnerAcceptRejects.delete(taskId); + this.refreshRequestTimeout(request); + // The runner may have created the task before its acknowledgement timed out, + // holding one of its concurrency slots until the task's own timeout elapses. + await this.messageRunner(offer.runnerId, { + type: 'broker:taskcancel', + taskId, + reason: 'Acknowledgement window elapsed', + }); + // A runner missing the acknowledgement window may be gone without its transport + // having reported it yet, so its remaining offers cannot be trusted either. + this.discardOffersFrom(offer.runnerId); + this.settleTasks(); return; } throw e; + } finally { + clearTimeout(runnerAcceptTimer); } clearTimeout(request.timeout); @@ -594,13 +695,16 @@ export class TaskBroker { (r) => r.requestId === request.requestId, ); if (requestIndex === -1) { - this.logger.error( - `Failed to find task request (${request.requestId}) after a task was accepted. This shouldn't happen, and might be a race condition.`, + this.logger.warn( + `Task request (${request.requestId}) expired while runner (${offer.runnerId}) was accepting task (${taskId}), canceling task`, ); + await this.cancelTask(taskId, 'Task request expired'); return; } this.pendingTaskRequests.splice(requestIndex, 1); + let requesterAcceptTimer: NodeJS.Timeout | undefined; + try { const acceptPromise = new Promise( (resolve, reject) => { @@ -611,10 +715,9 @@ export class TaskBroker { reject, }); - // TODO: customisable timeout - setTimeout(() => { - reject('Requester timed out'); - }, 2000); + requesterAcceptTimer = setTimeout(() => { + reject(new TaskRequesterAcceptTimeoutError(taskId, request.requesterId)); + }, this.taskRunnersConfig.taskAcceptTimeout * Time.seconds.toMilliseconds); }, ); @@ -632,8 +735,20 @@ export class TaskBroker { this.logger.info(`Task (${taskId}) rejected by Requester with reason "${e.reason}"`); return; } + if (e instanceof TaskRunnerUnreachableError) { + this.logger.warn(e.message); + return; + } + if (e instanceof TaskRequesterAcceptTimeoutError) { + this.logger.warn(e.message); + this.requesterAcceptRejects.delete(taskId); + await this.cancelTask(task.id, 'Requester took too long to acknowledge the task'); + return; + } await this.cancelTask(task.id, 'Unknown reason'); throw e; + } finally { + clearTimeout(requesterAcceptTimer); } } @@ -647,6 +762,7 @@ export class TaskBroker { // lists settleTasks() { this.expireTasks(); + this.discardOffersFromUnreachableRunners(); for (const request of this.pendingTaskRequests) { if (request.acceptInProgress) { @@ -737,6 +853,10 @@ export class TaskBroker { return this.runnerAcceptRejects; } + getRequesterAcceptRejects() { + return this.requesterAcceptRejects; + } + setTasks(tasks: Record) { this.tasks = new Map(Object.entries(tasks)); } @@ -752,7 +872,7 @@ export class TaskBroker { setRunnerAcceptRejects( runnerAcceptRejects: Record< string, - { accept: RunnerAcceptCallback; reject: TaskRejectCallback } + { accept: RunnerAcceptCallback; reject: TaskRejectCallback; runnerId: TaskRunner['id'] } >, ) { this.runnerAcceptRejects = new Map(Object.entries(runnerAcceptRejects)); diff --git a/packages/cli/src/task-runners/task-managers/task-requester.ts b/packages/cli/src/task-runners/task-managers/task-requester.ts index 89a31581232..3a44d4f7ca5 100644 --- a/packages/cli/src/task-runners/task-managers/task-requester.ts +++ b/packages/cli/src/task-runners/task-managers/task-requester.ts @@ -55,7 +55,10 @@ export type RunnerStatus = { available: true } | { available: false; reason?: st @Service() export abstract class TaskRequester { - requestAcceptRejects: Map = new Map(); + requestAcceptRejects: Map< + string, + { accept: RequestAccept; reject: RequestReject; requestedAt: number } + > = new Map(); taskAcceptRejects: Map = new Map(); @@ -136,6 +139,7 @@ export abstract class TaskRequester { this.requestAcceptRejects.set(request.requestId, { accept: resolve, reject, + requestedAt: Date.now(), }); }); @@ -301,14 +305,17 @@ export abstract class TaskRequester { const acceptReject = this.requestAcceptRejects.get(requestId); if (!acceptReject) return; + const elapsedSeconds = Math.round((Date.now() - acceptReject.requestedAt) / 1000); + const error = new TaskRequestTimeoutError({ - timeout: this.taskRunnersConfig.taskRequestTimeout, + elapsedSeconds, isSelfHosted: this.globalConfig.deployment.type !== 'cloud', }); this.errorReporter.error('Task request timed out', { extra: { requestId, + elapsedSeconds, timeout: this.taskRunnersConfig.taskRequestTimeout, deploymentType: this.globalConfig.deployment.type, },