From afc4a54798bbed080dad6b68231655a34d4ed14d Mon Sep 17 00:00:00 2001 From: Emilia <100027345+sovietspaceship@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:41:31 +0000 Subject: [PATCH] fix(core): Wait for detached sub-workflows before the worker stops its task runner (#37737) Co-authored-by: Claude Opus 5 --- .../src/repositories/execution.repository.ts | 9 + .../src/__tests__/active-executions.test.ts | 165 ++++++++++- packages/cli/src/active-executions.ts | 66 ++++- .../cli/src/commands/__tests__/worker.test.ts | 96 +++++-- packages/cli/src/commands/worker.ts | 21 +- .../scaling/__tests__/scaling.service.test.ts | 258 +++++++++++++++++- packages/cli/src/scaling/scaling.service.ts | 88 +++++- .../repositories/execution.repository.test.ts | 70 +++++ 8 files changed, 724 insertions(+), 49 deletions(-) diff --git a/packages/@n8n/db/src/repositories/execution.repository.ts b/packages/@n8n/db/src/repositories/execution.repository.ts index f16852d91e0..82000fe9a40 100644 --- a/packages/@n8n/db/src/repositories/execution.repository.ts +++ b/packages/@n8n/db/src/repositories/execution.repository.ts @@ -754,6 +754,15 @@ export class ExecutionRepository extends BaseRepository { ); } + async cancelManyRunning(executionIds: string[]) { + await this.update( + // The caller's ID list is a snapshot. The status match stops an execution that + // reached a terminal status in the meantime from being recorded as cancelled. + { id: In(executionIds), status: 'running' }, + { status: 'canceled', stoppedAt: new Date(), waitTill: null }, + ); + } + // ---------------------------------- // new API // ---------------------------------- diff --git a/packages/cli/src/__tests__/active-executions.test.ts b/packages/cli/src/__tests__/active-executions.test.ts index 8c9edbab139..ff43e51f69a 100644 --- a/packages/cli/src/__tests__/active-executions.test.ts +++ b/packages/cli/src/__tests__/active-executions.test.ts @@ -6,6 +6,7 @@ import type { ExecutionRepository } from '@n8n/db'; import type { IDeferredPromise } from '@n8n/utils/promise/deferred-promise'; import type { Response } from 'express'; import type { + ExecutionStatus, IExecuteResponsePromiseData, IRun, IWorkflowExecutionDataProcess, @@ -37,6 +38,7 @@ vi.mock('@n8n/utils/sleep', () => ({ const FAKE_EXECUTION_ID = '15'; const FAKE_SECOND_EXECUTION_ID = '20'; +const logger = mock(); const executionRepository = mock(); const executionPersistence = mock(); @@ -81,7 +83,7 @@ describe('ActiveExecutions', () => { beforeEach(() => { activeExecutions = new ActiveExecutions( - mock(), + logger, executionRepository, executionPersistence, concurrencyControl, @@ -89,6 +91,7 @@ describe('ActiveExecutions', () => { executionsConfig, ); + executionRepository.cancelManyRunning.mockResolvedValue(); executionPersistence.create.mockResolvedValue(FAKE_EXECUTION_ID); executionPersistence.updateExistingExecution.mockResolvedValue(true); executionRepository.setRunning.mockResolvedValue(new Date()); @@ -550,6 +553,166 @@ describe('ActiveExecutions', () => { }); }); + describe('getRunningExecutionIds and cancelRunningExecutions', () => { + const addExecutionWithStatus = async (status: ExecutionStatus) => { + const executionId = await activeExecutions.add(executionData); + activeExecutions.setStatus(executionId, status); + return executionId; + }; + + beforeEach(() => { + let i = 2000; + executionPersistence.create.mockImplementation(async () => `${i++}`); + }); + + test('Should list only the executions with a running status', async () => { + const runningExecutionId = await addExecutionWithStatus('running'); + await addExecutionWithStatus('waiting'); + await addExecutionWithStatus('new'); + + expect(activeExecutions.getRunningExecutionIds()).toEqual([runningExecutionId]); + }); + + test('Should cancel only the executions with a running status', async () => { + const runningExecutionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(runningExecutionId, workflowExecution); + const runningPostExecutePromise = activeExecutions.getPostExecutePromise(runningExecutionId); + + const waitingExecutionId = await addExecutionWithStatus('waiting'); + + await expect(activeExecutions.cancelRunningExecutions()).resolves.toEqual([ + runningExecutionId, + ]); + + await expect(runningPostExecutePromise).rejects.toThrow( + SystemShutdownExecutionCancelledError, + ); + expect(workflowExecution.cancel).toHaveBeenCalled(); + expect(activeExecutions.has(waitingExecutionId)).toBe(true); + }); + + test('Should not cancel a running execution with no workflow execution attached yet', async () => { + const attachedExecutionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(attachedExecutionId, workflowExecution); + const attachedPostExecutePromise = + activeExecutions.getPostExecutePromise(attachedExecutionId); + + const unattachedExecutionId = await addExecutionWithStatus('running'); + + await expect(activeExecutions.cancelRunningExecutions()).resolves.toEqual([ + attachedExecutionId, + ]); + + await expect(attachedPostExecutePromise).rejects.toThrow( + SystemShutdownExecutionCancelledError, + ); + expect(activeExecutions.has(unattachedExecutionId)).toBe(true); + expect(activeExecutions.getRunningExecutionIds()).toEqual([unattachedExecutionId]); + }); + + test('Should list and cancel nothing when no execution is running', async () => { + const waitingExecutionId = await addExecutionWithStatus('waiting'); + + expect(activeExecutions.getRunningExecutionIds()).toEqual([]); + await expect(activeExecutions.cancelRunningExecutions()).resolves.toEqual([]); + expect(activeExecutions.has(waitingExecutionId)).toBe(true); + expect(executionRepository.cancelManyRunning).not.toHaveBeenCalled(); + }); + + test('Should record the executions as cancelled before returning', async () => { + const executionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(executionId, workflowExecution); + const postExecute = activeExecutions.getPostExecutePromise(executionId); + + let hasRecorded = false; + executionRepository.cancelManyRunning.mockImplementation(async () => { + hasRecorded = true; + }); + + await activeExecutions.cancelRunningExecutions(); + + expect(hasRecorded).toBe(true); + expect(executionRepository.cancelManyRunning).toHaveBeenCalledWith([executionId]); + await expect(postExecute).rejects.toThrow(SystemShutdownExecutionCancelledError); + }); + + test('Should still cancel and report the executions when the recording fails', async () => { + const executionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(executionId, workflowExecution); + const postExecute = activeExecutions.getPostExecutePromise(executionId); + + executionRepository.cancelManyRunning.mockRejectedValue(new Error('Connection terminated')); + + await expect(activeExecutions.cancelRunningExecutions()).resolves.toEqual([executionId]); + + await expect(postExecute).rejects.toThrow(SystemShutdownExecutionCancelledError); + expect(logger.error).toHaveBeenCalledWith( + 'Failed to record 1 cancelled executions: Connection terminated', + { executionIds: [executionId] }, + ); + }); + + test('Should give up on a recording that does not settle within its deadline', async () => { + vi.useFakeTimers(); + + const executionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(executionId, workflowExecution); + + executionRepository.cancelManyRunning.mockReturnValue(new Promise(() => {})); + + let hasCancelled = false; + const cancelling = activeExecutions.cancelRunningExecutions().then((ids) => { + hasCancelled = true; + return ids; + }); + + await vi.advanceTimersByTimeAsync(2_999); + + expect(hasCancelled).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + + await expect(cancelling).resolves.toEqual([executionId]); + expect(workflowExecution.cancel).toHaveBeenCalled(); + expect(logger.error).toHaveBeenCalledWith( + 'Failed to record 1 cancelled executions: Timed out writing the cancelled status', + { executionIds: [executionId] }, + ); + + vi.useRealTimers(); + }); + + test('Should give up on a recording that outlives the deadline the caller passed', async () => { + vi.useFakeTimers(); + + const executionId = await addExecutionWithStatus('running'); + activeExecutions.attachWorkflowExecution(executionId, workflowExecution); + + executionRepository.cancelManyRunning.mockReturnValue(new Promise(() => {})); + + let hasCancelled = false; + const cancelling = activeExecutions.cancelRunningExecutions(1_000).then((ids) => { + hasCancelled = true; + return ids; + }); + + await vi.advanceTimersByTimeAsync(999); + + expect(hasCancelled).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + + await expect(cancelling).resolves.toEqual([executionId]); + expect(workflowExecution.cancel).toHaveBeenCalled(); + expect(logger.error).toHaveBeenCalledWith( + 'Failed to record 1 cancelled executions: Timed out writing the cancelled status', + { executionIds: [executionId] }, + ); + + vi.useRealTimers(); + }); + }); + describe('shutdown', () => { let newExecutionId1: string, newExecutionId2: string; let waitingExecutionId1: string, waitingExecutionId2: string; diff --git a/packages/cli/src/active-executions.ts b/packages/cli/src/active-executions.ts index 0a13c92948e..6219481f50d 100644 --- a/packages/cli/src/active-executions.ts +++ b/packages/cli/src/active-executions.ts @@ -1,8 +1,10 @@ import { Logger } from '@n8n/backend-common'; import { ExecutionsConfig } from '@n8n/config'; +import { Time } from '@n8n/constants'; import type { CreateExecutionPayload, IExecutionDb } from '@n8n/db'; import { ExecutionRepository } from '@n8n/db'; import { Service } from '@n8n/di'; +import { ensureError } from '@n8n/utils/errors/ensure-error'; import { createDeferredPromise, type IDeferredPromise } from '@n8n/utils/promise/deferred-promise'; import type { IExecuteResponsePromiseData, @@ -12,7 +14,11 @@ import type { StructuredChunk, WebhookResponseMode, } from 'n8n-workflow'; -import { ExecutionCancelledError, SystemShutdownExecutionCancelledError } from 'n8n-workflow'; +import { + ExecutionCancelledError, + OperationalError, + SystemShutdownExecutionCancelledError, +} from 'n8n-workflow'; import { sleep } from '@n8n/utils/sleep'; import { strict as assert } from 'node:assert'; import type PCancelable from 'p-cancelable'; @@ -32,6 +38,8 @@ import { ConcurrencyCapacityReservation } from './concurrency/concurrency-capaci import { ConcurrencyControlService } from './concurrency/concurrency-control.service'; import { EventService } from './events/event.service'; +const DEFAULT_CANCEL_WRITE_TIMEOUT_MS = 3 * Time.seconds.toMilliseconds; + @Service() export class ActiveExecutions { /** @@ -311,6 +319,62 @@ export class ActiveExecutions { return returnData; } + getRunningExecutionIds(): string[] { + return Object.keys(this.activeExecutions).filter( + (executionId) => this.activeExecutions[executionId].status === 'running', + ); + } + + /** + * @param writeDeadlineMs - How long to wait for the cancelled status to be recorded. + * Pass what the caller's own shutdown window can still afford. + */ + async cancelRunningExecutions( + writeDeadlineMs = DEFAULT_CANCEL_WRITE_TIMEOUT_MS, + ): Promise { + // An execution is registered before its workflow execution is attached. To + // cancel inside that window records the execution as failed, not cancelled. + const executionIds = this.getRunningExecutionIds().filter( + (executionId) => this.activeExecutions[executionId].workflowExecution !== undefined, + ); + + if (executionIds.length === 0) return executionIds; + + // The engine's own write for a cancel is fire-and-forget and conditional on the + // execution not being canceled, so recording first makes it a no-op, not a race. + await this.recordAsCancelled(executionIds, writeDeadlineMs); + + for (const executionId of executionIds) { + this.stopExecution(executionId, new SystemShutdownExecutionCancelledError(executionId)); + } + + return executionIds; + } + + /** Record executions as cancelled, bounded so a stalling database cannot hold up shutdown. */ + private async recordAsCancelled(executionIds: string[], writeDeadlineMs: number) { + let timeout: NodeJS.Timeout | undefined; + + try { + await Promise.race([ + this.executionRepository.cancelManyRunning(executionIds), + new Promise((_, reject) => { + timeout = setTimeout( + () => reject(new OperationalError('Timed out writing the cancelled status')), + writeDeadlineMs, + ); + }), + ]); + } catch (error) { + this.logger.error( + `Failed to record ${executionIds.length} cancelled executions: ${ensureError(error).message}`, + { executionIds }, + ); + } finally { + clearTimeout(timeout); + } + } + setStatus(executionId: string, status: ExecutionStatus) { this.getExecutionOrFail(executionId).status = status; } diff --git a/packages/cli/src/commands/__tests__/worker.test.ts b/packages/cli/src/commands/__tests__/worker.test.ts index 3ba2e52fc61..f32c6afef2d 100644 --- a/packages/cli/src/commands/__tests__/worker.test.ts +++ b/packages/cli/src/commands/__tests__/worker.test.ts @@ -76,6 +76,43 @@ describe('Worker', () => { vi.clearAllMocks(); }); + /** Worker with the init steps that go beyond `super.init()` stubbed, as in start.test.ts. */ + const createWorkerForInit = (globalConfigOverrides: Record = {}) => { + const worker = new Worker(); + + // @ts-expect-error - Overriding readonly property for testing + worker.globalConfig = { + executions: { mode: 'regular' }, + multiMainSetup: { enabled: false }, + endpoints: { metrics: { enable: false }, health: '/health' }, + database: { type: 'sqlite' }, + sentry: { backendDsn: '' }, + cache: { backend: 'memory' }, + taskRunners: {}, + outboundProxy: { mode: 'main-only' }, + expressionEngine: { engine: 'legacy', poolSize: 1, maxCodeCacheSize: 1024 }, + queue: { bull: { gracefulShutdownTimeout: 20 }, workerPool: { enabled: false, name: '' } }, + generic: { gracefulShutdownTimeout: 30 }, + ...globalConfigOverrides, + }; + + worker.setConcurrency = vi.fn().mockResolvedValue(undefined); + worker.initLicense = vi.fn().mockResolvedValue(undefined); + worker.initBinaryDataService = vi.fn().mockResolvedValue(undefined); + // @ts-expect-error - Accessing protected method for testing + worker.initDataDeduplicationService = vi.fn().mockResolvedValue(undefined); + worker.initExternalHooks = vi.fn().mockResolvedValue(undefined); + worker.initEventBus = vi.fn().mockResolvedValue(undefined); + worker.initScalingService = vi.fn().mockResolvedValue(undefined); + worker.initOrchestration = vi.fn().mockResolvedValue(undefined); + // @ts-expect-error - Accessing protected property for testing + worker.moduleRegistry = { initModules: vi.fn().mockResolvedValue(undefined) }; + // @ts-expect-error - Accessing protected property for testing + worker.executionContextHookRegistry = { init: vi.fn().mockResolvedValue(undefined) }; + + return worker; + }; + describe('initOrchestration', () => { it('should instantiate WorkerStatusService during orchestration setup', async () => { const containerGetSpy = vi.spyOn(Container, 'get'); @@ -143,42 +180,41 @@ describe('Worker', () => { it('should install env-proxy global agents on `init()`, via the base command', async () => { vi.stubEnv('HTTPS_PROXY', 'http://proxy.host.invalid:3128'); - const worker = new Worker(); - // @ts-expect-error - Overriding readonly property for testing - worker.globalConfig = { - executions: { mode: 'regular' }, - multiMainSetup: { enabled: false }, - endpoints: { metrics: { enable: false }, health: '/health' }, - database: { type: 'sqlite' }, - sentry: { backendDsn: '' }, - cache: { backend: 'memory' }, - taskRunners: {}, - outboundProxy: { mode: 'all' }, - expressionEngine: { engine: 'legacy', poolSize: 1, maxCodeCacheSize: 1024 }, - queue: { workerPool: { enabled: false, name: '' } }, - }; - // Stub the init steps that go beyond `super.init()`, as in start.test.ts - worker.setConcurrency = vi.fn().mockResolvedValue(undefined); - worker.initLicense = vi.fn().mockResolvedValue(undefined); - worker.initBinaryDataService = vi.fn().mockResolvedValue(undefined); - // @ts-expect-error - Accessing protected method for testing - worker.initDataDeduplicationService = vi.fn().mockResolvedValue(undefined); - worker.initExternalHooks = vi.fn().mockResolvedValue(undefined); - worker.initEventBus = vi.fn().mockResolvedValue(undefined); - worker.initScalingService = vi.fn().mockResolvedValue(undefined); - worker.initOrchestration = vi.fn().mockResolvedValue(undefined); - // @ts-expect-error - Accessing protected property for testing - worker.moduleRegistry = { initModules: vi.fn().mockResolvedValue(undefined) }; - // @ts-expect-error - Accessing protected property for testing - worker.executionContextHookRegistry = { init: vi.fn().mockResolvedValue(undefined) }; - - await worker.init(); + await createWorkerForInit({ outboundProxy: { mode: 'all' } }).init(); expect(http.globalAgent.constructor.name).toBe('EnvProxyHttpAgent'); expect(https.globalAgent.constructor.name).toBe('EnvProxyHttpsAgent'); }); }); + describe('init', () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it.each([ + { envValue: '45', expected: 45, case: 'the parsed value' }, + { envValue: 'not-a-number', expected: 20, case: 'the queue default when unparseable' }, + { envValue: '45seconds', expected: 20, case: 'the queue default when partly numeric' }, + { envValue: '0', expected: 20, case: 'the queue default when zero' }, + { envValue: '-45', expected: 20, case: 'the queue default when negative' }, + ])( + 'should apply QUEUE_WORKER_TIMEOUT as $case to both the shutdown and the drain timeout', + async ({ envValue, expected }) => { + vi.stubEnv('QUEUE_WORKER_TIMEOUT', envValue); + + const worker = createWorkerForInit(); + + await worker.init(); + + // @ts-expect-error - Accessing protected property for testing + expect(worker.gracefulShutdownTimeoutInS).toBe(expected); + // @ts-expect-error - Accessing protected property for testing + expect(worker.globalConfig.generic.gracefulShutdownTimeout).toBe(expected); + }, + ); + }); + describe('stopProcess', () => { it('should keep the DB connection open until in-flight executions have persisted', async () => { // In-process executions on a worker e.g. sub-workflows started by diff --git a/packages/cli/src/commands/worker.ts b/packages/cli/src/commands/worker.ts index 4340371b3a9..1c44ad1b61c 100644 --- a/packages/cli/src/commands/worker.ts +++ b/packages/cli/src/commands/worker.ts @@ -65,8 +65,7 @@ export class Worker extends BaseCommand> { try { await this.externalHooks?.run('n8n.stop'); - // Wait for in-process executions not tracked as Bull jobs, - // which are instead drained by `ScalingService.stopWorker` + // Final wait for any execution still active after the shutdown hooks. await Container.get(ActiveExecutions).shutdown(); } catch (error) { await this.exitWithCrash('Error shutting down worker', error); @@ -88,8 +87,22 @@ export class Worker extends BaseCommand> { async init() { const { QUEUE_WORKER_TIMEOUT } = process.env; if (QUEUE_WORKER_TIMEOUT) { - this.gracefulShutdownTimeoutInS = - parseInt(QUEUE_WORKER_TIMEOUT, 10) || this.globalConfig.queue.bull.gracefulShutdownTimeout; + // Accept a whole positive number only, the way the config layer parses + // `N8N_GRACEFUL_SHUTDOWN_TIMEOUT`, so a malformed value cannot shorten the + // shutdown window. + const parsed = Number(QUEUE_WORKER_TIMEOUT); + const isValid = Number.isInteger(parsed) && parsed > 0; + + if (!isValid) { + this.logger.warn( + `Invalid QUEUE_WORKER_TIMEOUT value "${QUEUE_WORKER_TIMEOUT}". Using the configured queue timeout.`, + ); + } + + const timeout = isValid ? parsed : this.globalConfig.queue.bull.gracefulShutdownTimeout; + // One field arms the force-exit timer, the other sizes the shutdown drains. + this.gracefulShutdownTimeoutInS = timeout; + this.globalConfig.generic.gracefulShutdownTimeout = timeout; this.logger.warn( 'QUEUE_WORKER_TIMEOUT has been deprecated. Rename it to N8N_GRACEFUL_SHUTDOWN_TIMEOUT.', ); diff --git a/packages/cli/src/scaling/__tests__/scaling.service.test.ts b/packages/cli/src/scaling/__tests__/scaling.service.test.ts index a094ff1bd98..67d09eef5a5 100644 --- a/packages/cli/src/scaling/__tests__/scaling.service.test.ts +++ b/packages/cli/src/scaling/__tests__/scaling.service.test.ts @@ -1,3 +1,4 @@ +import type { Logger } from '@n8n/backend-common'; import { mockLogger, mockInstance } from '@n8n/backend-test-utils'; import { GlobalConfig, WorkerPoolConfig } from '@n8n/config'; import type { ExecutionRepository } from '@n8n/db'; @@ -84,9 +85,16 @@ describe('ScalingService', () => { keepLastFailed: 0, }, }, + generic: { + gracefulShutdownTimeout: 30, + }, }); const instanceSettings = Container.get(InstanceSettings); + // The service scopes its logger on construction, so assertions go to the scoped mock. + const scopedLogger = mock(); + const logger = mock({ scoped: () => scopedLogger }); + const activeExecutions = mock(); const jobProcessor = mock(); const executionRepository = mock(); const executionPersistence = mock(); @@ -117,11 +125,15 @@ describe('ScalingService', () => { // @ts-expect-error readonly property instanceSettings.instanceType = 'main'; instanceSettings.markAsLeader(); + activeExecutions.getRunningExecutionIds.mockReturnValue([]); + activeExecutions.cancelRunningExecutions.mockResolvedValue([]); + jobProcessor.getRunningJobsSummary.mockReturnValue([]); + globalConfig.generic.gracefulShutdownTimeout = 30; scalingService = new ScalingService( - mockLogger(), - mock(), + logger, mock(), + activeExecutions, jobProcessor, globalConfig, executionRepository, @@ -151,6 +163,10 @@ describe('ScalingService', () => { stopQueueMetricsSpy = vi.spyOn(scalingService, 'stopQueueMetrics'); }); + afterEach(() => { + vi.useRealTimers(); + }); + describe('setupQueue', () => { describe('if leader main', () => { it('should set up queue + listeners + queue recovery', async () => { @@ -320,20 +336,252 @@ describe('ScalingService', () => { }); it('should log the execution IDs it is waiting for while draining', async () => { + vi.useFakeTimers(); // @ts-expect-error readonly property instanceSettings.instanceType = 'worker'; await scalingService.setupQueue(); jobProcessor.getRunningJobIds.mockReturnValueOnce(['1']).mockReturnValue([]); jobProcessor.getRunningJobsSummary.mockReturnValue([mock({ executionId: 'exec-1' })]); - await scalingService.stop(); + const stopped = scalingService.stop(); + await vi.advanceTimersByTimeAsync(500); + await stopped; - // @ts-expect-error private property - expect(scalingService.logger.info).toHaveBeenCalledWith( + expect(scopedLogger.info).toHaveBeenCalledWith( 'Waiting for 1 active executions to finish... (execution IDs: exec-1)', { executionIds: ['exec-1'] }, ); }); + + it('should keep waiting for an in-process execution that has no queue job', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + + let inProcessExecutionIds = ['exec-1']; + activeExecutions.getRunningExecutionIds.mockImplementation(() => inProcessExecutionIds); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(2_000); + + expect(hasStopped).toBe(false); + expect(scopedLogger.info).toHaveBeenCalledWith( + 'Waiting for 1 in-process executions to finish... (execution IDs: exec-1)', + { executionIds: ['exec-1'] }, + ); + + inProcessExecutionIds = []; + await vi.advanceTimersByTimeAsync(500); + await stopped; + + expect(hasStopped).toBe(true); + expect(scopedLogger.warn).not.toHaveBeenCalled(); + }); + + it('should stop waiting and cancel the executions once the drain budget is spent', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + // The budget is 80% of the shutdown window, so 4s of the 5s here. + globalConfig.generic.gracefulShutdownTimeout = 5; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + activeExecutions.getRunningExecutionIds.mockReturnValue(['exec-1']); + activeExecutions.cancelRunningExecutions.mockResolvedValue(['exec-1']); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(3_500); + + expect(hasStopped).toBe(false); + expect(activeExecutions.cancelRunningExecutions).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(500); + await stopped; + + expect(hasStopped).toBe(true); + expect(activeExecutions.cancelRunningExecutions).toHaveBeenCalled(); + expect(scopedLogger.warn).toHaveBeenCalledWith( + 'Drain timeout reached after 4s, shutting down with executions still active...', + ); + expect(scopedLogger.warn).toHaveBeenCalledWith( + 'Cancelled 1 in-process executions that could not finish before shutdown (execution IDs: exec-1)', + { executionIds: ['exec-1'] }, + ); + }); + + it('should not finish the drain until the cancellation has settled', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + globalConfig.generic.gracefulShutdownTimeout = 5; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + activeExecutions.getRunningExecutionIds.mockReturnValue(['exec-1']); + + let finishCancellation: (executionIds: string[]) => void = () => {}; + activeExecutions.cancelRunningExecutions.mockReturnValue( + new Promise((resolve) => (finishCancellation = resolve)), + ); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(4_000); + + expect(activeExecutions.cancelRunningExecutions).toHaveBeenCalled(); + expect(hasStopped).toBe(false); + + finishCancellation(['exec-1']); + await stopped; + + expect(hasStopped).toBe(true); + expect(scopedLogger.warn).toHaveBeenCalledWith( + 'Cancelled 1 in-process executions that could not finish before shutdown (execution IDs: exec-1)', + { executionIds: ['exec-1'] }, + ); + }); + + it('should still drain for part of a one-second shutdown window', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + globalConfig.generic.gracefulShutdownTimeout = 1; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + + let inProcessExecutionIds = ['exec-1']; + activeExecutions.getRunningExecutionIds.mockImplementation(() => inProcessExecutionIds); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(0); + + expect(hasStopped).toBe(false); + + inProcessExecutionIds = []; + await vi.advanceTimersByTimeAsync(500); + await stopped; + + expect(hasStopped).toBe(true); + expect(activeExecutions.cancelRunningExecutions).not.toHaveBeenCalled(); + expect(scopedLogger.warn).not.toHaveBeenCalled(); + }); + + // The two warnings are the drain timeout and the cancellation summary. + it.each([ + { inProcessExecutionIds: [], expectedCancelCalls: 0, expectedWarnings: 0 }, + { inProcessExecutionIds: ['exec-1'], expectedCancelCalls: 1, expectedWarnings: 2 }, + ])( + 'should wait for queued jobs past the drain budget, then cancel in-process executions only if any are left (in-process: $inProcessExecutionIds)', + async ({ inProcessExecutionIds, expectedCancelCalls, expectedWarnings }) => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + globalConfig.generic.gracefulShutdownTimeout = 2; + await scalingService.setupQueue(); + + let runningJobIds = ['1']; + jobProcessor.getRunningJobIds.mockImplementation(() => runningJobIds); + activeExecutions.getRunningExecutionIds.mockReturnValue(inProcessExecutionIds); + activeExecutions.cancelRunningExecutions.mockResolvedValue(inProcessExecutionIds); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(10_000); + + expect(hasStopped).toBe(false); + expect(activeExecutions.cancelRunningExecutions).not.toHaveBeenCalled(); + expect(scopedLogger.warn).not.toHaveBeenCalled(); + + runningJobIds = []; + await vi.advanceTimersByTimeAsync(500); + await stopped; + + expect(hasStopped).toBe(true); + expect(activeExecutions.cancelRunningExecutions).toHaveBeenCalledTimes( + expectedCancelCalls, + ); + expect(scopedLogger.warn).toHaveBeenCalledTimes(expectedWarnings); + }, + ); + + it('should cancel within the shutdown window when the window is short', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + // The budget is 800ms, which the force-exit timer at 1s must not beat. + globalConfig.generic.gracefulShutdownTimeout = 1; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + activeExecutions.getRunningExecutionIds.mockReturnValue(['exec-1']); + activeExecutions.cancelRunningExecutions.mockResolvedValue(['exec-1']); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(800); + + expect(hasStopped).toBe(true); + expect(activeExecutions.cancelRunningExecutions).toHaveBeenCalled(); + + await stopped; + }); + + it('should not drain or warn when the shutdown window is zero', async () => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + globalConfig.generic.gracefulShutdownTimeout = 0; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + activeExecutions.getRunningExecutionIds.mockReturnValue(['exec-1']); + + let hasStopped = false; + const stopped = scalingService.stop().then(() => (hasStopped = true)); + + await vi.advanceTimersByTimeAsync(0); + await stopped; + + expect(hasStopped).toBe(true); + expect(activeExecutions.cancelRunningExecutions).not.toHaveBeenCalled(); + expect(scopedLogger.warn).not.toHaveBeenCalled(); + }); + + it.each([ + { shutdownTimeout: 30, expectedDeadlineMs: 3_000, case: 'the ceiling on a wide window' }, + { + shutdownTimeout: 10, + expectedDeadlineMs: 1_000, + case: 'half of a short window remainder', + }, + { shutdownTimeout: 1, expectedDeadlineMs: 100, case: 'half of a tiny window remainder' }, + ])( + 'should give the cancellation write $case', + async ({ shutdownTimeout, expectedDeadlineMs }) => { + vi.useFakeTimers(); + // @ts-expect-error readonly property + instanceSettings.instanceType = 'worker'; + globalConfig.generic.gracefulShutdownTimeout = shutdownTimeout; + await scalingService.setupQueue(); + jobProcessor.getRunningJobIds.mockReturnValue([]); + activeExecutions.getRunningExecutionIds.mockReturnValue(['exec-1']); + activeExecutions.cancelRunningExecutions.mockResolvedValue(['exec-1']); + + const stopped = scalingService.stop(); + await vi.advanceTimersByTimeAsync(shutdownTimeout * 1_000); + await stopped; + + expect(activeExecutions.cancelRunningExecutions).toHaveBeenCalledWith(expectedDeadlineMs); + }, + ); }); }); diff --git a/packages/cli/src/scaling/scaling.service.ts b/packages/cli/src/scaling/scaling.service.ts index fafe46befda..723c7403529 100644 --- a/packages/cli/src/scaling/scaling.service.ts +++ b/packages/cli/src/scaling/scaling.service.ts @@ -34,6 +34,14 @@ import type { } from './scaling.types'; import { decodeRelayedWebhookResponse, WebhookResponseRelay } from './webhook-response-relay'; +const DRAIN_POLL_INTERVAL_MS = 500; + +/** Share of the time left at the drain deadline that the cancellation write may use. */ +const CANCEL_WRITE_BUDGET_SHARE = 0.5; + +/** Ceiling for the cancellation write, so a long shutdown window does not stall on it. */ +const MAX_CANCEL_WRITE_TIMEOUT_MS = 3 * Time.seconds.toMilliseconds; + @Service() export class ScalingService { /** Bull queues keyed by queue name. Pool queues are created lazily. */ @@ -226,19 +234,83 @@ export class ScalingService { private async stopWorker() { await this.pauseAllQueues(); + const shutdownWindowMs = + this.globalConfig.generic.gracefulShutdownTimeout * Time.seconds.toMilliseconds; + + // The budget bounds only the in-process wait. The queued-job wait stays + // unbounded, so a long queued execution still runs to completion. + const drainTimeoutMs = shutdownWindowMs * 0.8; + + const start = Date.now(); + + const hasQueuedJobsToDrain = () => this.getRunningJobsCount() !== 0; + const hasInProcessExecutionsToDrain = () => + this.activeExecutions.getRunningExecutionIds().length !== 0; + const isWithinDrainBudget = () => Date.now() - start < drainTimeoutMs; + let count = 0; - while (this.getRunningJobsCount() !== 0) { - if (count++ % 4 === 0) { - const summaries = this.jobProcessor.getRunningJobsSummary(); - const executionIds = summaries.map((summary) => summary.executionId); - this.logger.info( - `Waiting for ${executionIds.length} active executions to finish... (execution IDs: ${executionIds.join(', ')})`, - { executionIds }, + while (hasQueuedJobsToDrain() || (hasInProcessExecutionsToDrain() && isWithinDrainBudget())) { + if (count++ % 4 === 0) this.logExecutionsToDrain(); + + // Cap the last sleep at what is left of the budget, so the tick cannot + // overshoot it and leave no time to cancel before the force-exit timer. + const remainingBudgetMs = drainTimeoutMs - (Date.now() - start); + const sleepMs = hasQueuedJobsToDrain() + ? DRAIN_POLL_INTERVAL_MS + : Math.max(1, Math.min(DRAIN_POLL_INTERVAL_MS, remainingBudgetMs)); + + await sleep(sleepMs); + } + + // Cancel the stragglers rather than leave them to run. The task runner stops + // next, so they cannot make progress. + if (drainTimeoutMs > 0 && hasInProcessExecutionsToDrain() && !isWithinDrainBudget()) { + const elapsed = Math.round((Date.now() - start) / Time.seconds.toMilliseconds); + + this.logger.warn( + `Drain timeout reached after ${elapsed}s, shutting down with executions still active...`, + ); + this.logExecutionsToDrain(); + + // The force-exit timer is armed at the full window, so the write gets a share of + // what is left of it. The rest stays for the shutdown hooks that run after this one. + const remainingWindowMs = Math.max(0, shutdownWindowMs - (Date.now() - start)); + const writeDeadlineMs = Math.min( + MAX_CANCEL_WRITE_TIMEOUT_MS, + Math.round(remainingWindowMs * CANCEL_WRITE_BUDGET_SHARE), + ); + + const cancelledExecutionIds = + await this.activeExecutions.cancelRunningExecutions(writeDeadlineMs); + + if (cancelledExecutionIds.length > 0) { + this.logger.warn( + `Cancelled ${cancelledExecutionIds.length} in-process executions that could not finish before shutdown (execution IDs: ${cancelledExecutionIds.join(', ')})`, + { executionIds: cancelledExecutionIds }, ); } + } + } - await sleep(500); + private logExecutionsToDrain() { + const summaries = this.jobProcessor.getRunningJobsSummary(); + const executionIds = summaries.map((summary) => summary.executionId); + + if (executionIds.length > 0) { + this.logger.info( + `Waiting for ${executionIds.length} active executions to finish... (execution IDs: ${executionIds.join(', ')})`, + { executionIds }, + ); + } + + const inProcessExecutionIds = this.activeExecutions.getRunningExecutionIds(); + + if (inProcessExecutionIds.length > 0) { + this.logger.info( + `Waiting for ${inProcessExecutionIds.length} in-process executions to finish... (execution IDs: ${inProcessExecutionIds.join(', ')})`, + { executionIds: inProcessExecutionIds }, + ); } } diff --git a/packages/cli/test/integration/database/repositories/execution.repository.test.ts b/packages/cli/test/integration/database/repositories/execution.repository.test.ts index d0b47c35033..26b9937ebb4 100644 --- a/packages/cli/test/integration/database/repositories/execution.repository.test.ts +++ b/packages/cli/test/integration/database/repositories/execution.repository.test.ts @@ -280,6 +280,76 @@ describe('ExecutionRepository', () => { }); }); + describe('cancelManyRunning', () => { + const createExecution = async (status: ExecutionStatus, extra: { waitTill?: Date } = {}) => { + const workflow = await createWorkflow(); + const { identifiers } = await Container.get(ExecutionRepository).insert({ + workflowId: workflow.id, + mode: 'manual', + startedAt: new Date(), + status, + finished: status === 'success', + createdAt: new Date(), + ...extra, + }); + return identifiers[0].id as string; + }; + + it('should cancel a running execution and clear its waitTill', async () => { + const executionRepo = Container.get(ExecutionRepository); + const runningId = await createExecution('running', { + waitTill: new Date(Date.now() + 1000 * 60 * 60), + }); + + await executionRepo.cancelManyRunning([runningId]); + + const runningExec = await executionRepo.findOneBy({ id: runningId }); + expect(runningExec?.status).toBe('canceled'); + expect(runningExec?.stoppedAt).toBeInstanceOf(Date); + expect(runningExec?.waitTill).toBeNull(); + }); + + it('should not overwrite executions in a terminal status', async () => { + const executionRepo = Container.get(ExecutionRepository); + const successId = await createExecution('success'); + const errorId = await createExecution('error'); + const crashedId = await createExecution('crashed'); + + await executionRepo.cancelManyRunning([successId, errorId, crashedId]); + + const [successExec, errorExec, crashedExec] = await Promise.all([ + executionRepo.findOneBy({ id: successId }), + executionRepo.findOneBy({ id: errorId }), + executionRepo.findOneBy({ id: crashedId }), + ]); + + expect(successExec?.status).toBe('success'); + expect(errorExec?.status).toBe('error'); + expect(crashedExec?.status).toBe('crashed'); + }); + + it('should cancel only the running executions in a mixed batch', async () => { + const executionRepo = Container.get(ExecutionRepository); + const waitTill = new Date(Date.now() + 1000 * 60 * 60); + const runningId = await createExecution('running'); + const waitingId = await createExecution('waiting', { waitTill }); + const successId = await createExecution('success'); + + await executionRepo.cancelManyRunning([runningId, waitingId, successId]); + + const [runningExec, waitingExec, successExec] = await Promise.all([ + executionRepo.findOneBy({ id: runningId }), + executionRepo.findOneBy({ id: waitingId }), + executionRepo.findOneBy({ id: successId }), + ]); + + expect(runningExec?.status).toBe('canceled'); + expect(waitingExec?.status).toBe('waiting'); + expect(waitingExec?.waitTill?.getTime()).toBe(waitTill.getTime()); + expect(successExec?.status).toBe('success'); + }); + }); + describe('getWorkflowIdsWithExecutionsSince', () => { const insertExecution = async (workflowId: string, startedAt: Date) => await Container.get(ExecutionRepository).insert({