diff --git a/packages/backend-defaults/src/entrypoints/scheduler/lib/TaskWorker.ts b/packages/backend-defaults/src/entrypoints/scheduler/lib/TaskWorker.ts index e33e112f43..4f32dc6863 100644 --- a/packages/backend-defaults/src/entrypoints/scheduler/lib/TaskWorker.ts +++ b/packages/backend-defaults/src/entrypoints/scheduler/lib/TaskWorker.ts @@ -263,11 +263,16 @@ export class TaskWorker { const timeoutHandle = setTimeout(() => { taskAbortController.abort(); }, Duration.fromISO(taskSettings.timeoutAfterDuration).as('milliseconds')); - const livenessHandle: { ref?: ReturnType } = {}; - livenessHandle.ref = setInterval( - () => this.checkLiveness(ticket, taskAbortController, livenessHandle), - this.workCheckFrequency.as('milliseconds'), - ); + let livenessHandle: ReturnType | undefined; + const scheduleLivenessCheck = () => { + livenessHandle = setTimeout(async () => { + await this.checkLiveness(ticket, taskAbortController); + if (!taskAbortController.signal.aborted) { + scheduleLivenessCheck(); + } + }, this.workCheckFrequency.as('milliseconds')); + }; + scheduleLivenessCheck(); try { this.#workerState = { @@ -284,7 +289,7 @@ export class TaskWorker { status: 'idle', }; clearTimeout(timeoutHandle); - clearInterval(livenessHandle.ref); + clearTimeout(livenessHandle); } await this.tryReleaseTask(ticket, taskSettings); @@ -379,7 +384,6 @@ export class TaskWorker { private async checkLiveness( ticket: string, taskAbortController: AbortController, - livenessHandle: { ref?: ReturnType }, ): Promise { try { const [row] = await this.knex(DB_TASKS_TABLE) @@ -390,7 +394,6 @@ export class TaskWorker { this.logger.info( `Task ticket for "${this.taskId}" is no longer valid; aborting execution`, ); - clearInterval(livenessHandle.ref); taskAbortController.abort(); } } catch (e) { diff --git a/packages/backend-test-utils/src/services/MockSchedulerService.test.ts b/packages/backend-test-utils/src/services/MockSchedulerService.test.ts index faf4e3e26d..1c451da749 100644 --- a/packages/backend-test-utils/src/services/MockSchedulerService.test.ts +++ b/packages/backend-test-utils/src/services/MockSchedulerService.test.ts @@ -206,6 +206,68 @@ describe('MockSchedulerService', () => { await expect(isDone()).resolves.toBe(true); }); + it('should cancel a running task and allow re-triggering with a fresh signal', async () => { + const scheduler = new MockSchedulerService(); + const signals: AbortSignal[] = []; + + scheduler.scheduleTask({ + ...baseOpts, + id: 'test', + fn: async signal => { + signals.push(signal); + // Simulate long-running work that respects cancellation + await new Promise((resolve, reject) => { + if (signal.aborted) { + reject(new Error('aborted')); + return; + } + signal.addEventListener('abort', () => reject(new Error('aborted'))); + setTimeout(1).then(resolve); + }); + }, + }); + + // First run completes normally + await scheduler.triggerTask('test'); + expect(signals).toHaveLength(1); + expect(signals[0].aborted).toBe(false); + + // Start a task that will block until cancelled + const blockingScheduler = new MockSchedulerService(); + let resolveBlock: (() => void) | undefined; + blockingScheduler.scheduleTask({ + ...baseOpts, + id: 'blocking', + fn: async signal => { + signals.push(signal); + await new Promise((resolve, reject) => { + signal.addEventListener('abort', () => reject(new Error('aborted'))); + resolveBlock = resolve; + }); + }, + }); + + const triggerPromise = blockingScheduler.triggerTask('blocking'); + // Give the task fn time to start + await setTimeout(1); + + await blockingScheduler.cancelTask('blocking'); + await triggerPromise.catch(() => {}); + + expect(signals).toHaveLength(2); + expect(signals[1].aborted).toBe(true); + + // Re-trigger should get a fresh non-aborted signal + resolveBlock = undefined; + const triggerPromise2 = blockingScheduler.triggerTask('blocking'); + await setTimeout(1); + resolveBlock?.(); + await triggerPromise2; + + expect(signals).toHaveLength(3); + expect(signals[2].aborted).toBe(false); + }); + it('should abort tasks when shutting down', async () => { let taskSignal: AbortSignal | undefined; diff --git a/packages/backend-test-utils/src/services/MockSchedulerService.ts b/packages/backend-test-utils/src/services/MockSchedulerService.ts index a47859a96f..3309ebc954 100644 --- a/packages/backend-test-utils/src/services/MockSchedulerService.ts +++ b/packages/backend-test-utils/src/services/MockSchedulerService.ts @@ -105,6 +105,7 @@ export class MockSchedulerService implements SchedulerService { throw new ConflictError(`Task ${id} is not running`); } task.abortControllers.abort(); + task.abortControllers = new AbortController(); } async triggerTask(id: string): Promise {