diff --git a/fluxer_api/src/api/worker/WorkerRunner.ts b/fluxer_api/src/api/worker/WorkerRunner.ts index 37ab3acab..d48e838c4 100644 --- a/fluxer_api/src/api/worker/WorkerRunner.ts +++ b/fluxer_api/src/api/worker/WorkerRunner.ts @@ -11,6 +11,7 @@ import {getWorkerService} from '../middleware/ServiceRegistry'; import {isJsonRecord, parseJsonRecord} from '../utils/JsonBoundaryUtils'; const MAX_DLQ_PUBLISH_ATTEMPTS = 3; +const MIN_ACK_HEARTBEAT_MS = 1000; interface WorkerRunnerJetStreamClient { consumers: { @@ -274,6 +275,7 @@ export class WorkerRunner { } }, }; + const ackHeartbeat = this.startAckHeartbeat(taskType, msg); try { await task(jobPayload, helpers); if (ledgerJobId !== null) { @@ -351,6 +353,25 @@ export class WorkerRunner { msg.nak(5000); } return false; + } finally { + clearInterval(ackHeartbeat); } } + + private startAckHeartbeat(taskType: string, msg: JsMsg): ReturnType { + const heartbeat = setInterval( + () => { + try { + msg.working(); + } catch (err) { + Logger.warn({workerId: this.workerId, taskType, seq: msg.seq, err}, 'Failed to extend job ack deadline'); + } + }, + Math.max(MIN_ACK_HEARTBEAT_MS, Math.floor(this.ackWaitMs / 2)), + ); + if (typeof heartbeat === 'object' && heartbeat && 'unref' in heartbeat) { + (heartbeat as {unref(): void}).unref(); + } + return heartbeat; + } } diff --git a/fluxer_api/src/api/worker/tests/WorkerAckHeartbeat.test.ts b/fluxer_api/src/api/worker/tests/WorkerAckHeartbeat.test.ts new file mode 100644 index 000000000..29ab7ba18 --- /dev/null +++ b/fluxer_api/src/api/worker/tests/WorkerAckHeartbeat.test.ts @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import {JobCancelledError} from '@pkgs/worker/src/contracts/WorkerTask'; +import type {JsMsg} from 'nats'; +import {afterEach, beforeAll, describe, expect, it, vi} from 'vitest'; +import {KVScheduledJobQueueService} from '../../infrastructure/KVScheduledJobQueueService'; +import type {IJobLedgerRepository} from '../../jobs/IJobLedgerRepository'; +import {setInjectedWorkerService} from '../../middleware/ServiceRegistry'; +import {MockKVProvider} from '../../test/mocks/MockKVProvider'; +import {NoopWorkerService} from '../../test/NoopWorkerService'; +import {WorkerRunner} from '../WorkerRunner'; + +const LANE_ACK_WAIT_MS = 120000; +const LANE_MAX_DELIVER = 25; +const TASK_TYPE = 'processInactivityDeletions'; + +const queueStub = { + getConnectionManager: () => { + throw new Error('WorkerRunner tests never consume messages'); + }, + getStreamName: () => 'JOBS', + publishToDlq: vi.fn(), +}; + +class TestWorkerRunner extends WorkerRunner { + async runJob(taskType: string, msg: JsMsg): Promise { + return await this.processJob(taskType, msg); + } +} + +function createRunner(task: () => Promise): TestWorkerRunner { + return new TestWorkerRunner({ + tasks: {[TASK_TYPE]: task}, + queue: queueStub, + scheduledJobQueue: new KVScheduledJobQueueService(new MockKVProvider()), + consumerName: 'workers_batch', + laneName: 'batch', + ledger: {} as IJobLedgerRepository, + concurrency: 12, + maxDeliver: LANE_MAX_DELIVER, + ackWaitMs: LANE_ACK_WAIT_MS, + }); +} + +function createJobMessage() { + const envelope = { + payload: {}, + max_attempts: 5, + priority: 0, + created_at: new Date().toISOString(), + }; + return { + seq: 1, + subject: `jobs.${TASK_TYPE}`, + redelivered: false, + data: new TextEncoder().encode(JSON.stringify(envelope)), + info: {deliveryCount: 1}, + ack: vi.fn(), + nak: vi.fn(), + term: vi.fn(), + working: vi.fn(), + }; +} + +function createDeferredTask(): {task: () => Promise; settle: (error?: Error) => void} { + let settle: (error?: Error) => void = () => {}; + const pending = new Promise((resolve, reject) => { + settle = (error?: Error) => { + if (error) { + reject(error); + } else { + resolve(); + } + }; + }); + return {task: () => pending, settle: (error?: Error) => settle(error)}; +} + +describe('Worker ack heartbeat', () => { + beforeAll(() => { + setInjectedWorkerService(new NoopWorkerService()); + }); + + afterEach(() => { + vi.useRealTimers(); + queueStub.publishToDlq.mockClear(); + }); + + it('holds the ack for a task running past the lane ack wait', async () => { + vi.useFakeTimers(); + const {task, settle} = createDeferredTask(); + const runner = createRunner(task); + const msg = createJobMessage(); + + const job = runner.runJob(TASK_TYPE, msg as unknown as JsMsg); + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS * 3); + + expect(msg.working.mock.calls.length).toBeGreaterThanOrEqual(1); + expect(msg.ack).not.toHaveBeenCalled(); + expect(msg.nak).not.toHaveBeenCalled(); + + settle(); + await expect(job).resolves.toBe(true); + + expect(msg.ack).toHaveBeenCalledTimes(1); + const heartbeatsAtCompletion = msg.working.mock.calls.length; + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS * 3); + expect(msg.working.mock.calls.length).toBe(heartbeatsAtCompletion); + }); + + it('stops the heartbeat when the task throws', async () => { + vi.useFakeTimers(); + const {task, settle} = createDeferredTask(); + const runner = createRunner(task); + const msg = createJobMessage(); + + const job = runner.runJob(TASK_TYPE, msg as unknown as JsMsg); + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS); + expect(msg.working.mock.calls.length).toBeGreaterThanOrEqual(1); + + settle(new Error('scan failed')); + await expect(job).resolves.toBe(false); + + expect(msg.nak).toHaveBeenCalledTimes(1); + const heartbeatsAtCompletion = msg.working.mock.calls.length; + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS * 3); + expect(msg.working.mock.calls.length).toBe(heartbeatsAtCompletion); + }); + + it('stops the heartbeat when the task is cancelled', async () => { + vi.useFakeTimers(); + const {task, settle} = createDeferredTask(); + const runner = createRunner(task); + const msg = createJobMessage(); + + const job = runner.runJob(TASK_TYPE, msg as unknown as JsMsg); + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS); + expect(msg.working.mock.calls.length).toBeGreaterThanOrEqual(1); + + settle(new JobCancelledError('cancelled by admin')); + await expect(job).resolves.toBe(false); + + expect(msg.ack).toHaveBeenCalledTimes(1); + const heartbeatsAtCompletion = msg.working.mock.calls.length; + await vi.advanceTimersByTimeAsync(LANE_ACK_WAIT_MS * 3); + expect(msg.working.mock.calls.length).toBe(heartbeatsAtCompletion); + }); + + it('does not heartbeat a parked scheduled job', async () => { + const scheduledJobQueue = new KVScheduledJobQueueService(new MockKVProvider()); + const runner = new TestWorkerRunner({ + tasks: {[TASK_TYPE]: async () => {}}, + queue: queueStub, + scheduledJobQueue, + consumerName: 'workers_batch', + laneName: 'batch', + ledger: {} as IJobLedgerRepository, + concurrency: 12, + maxDeliver: LANE_MAX_DELIVER, + ackWaitMs: LANE_ACK_WAIT_MS, + }); + const runAt = new Date(Date.now() + 60 * 60 * 1000); + const envelope = { + payload: {}, + run_at: runAt.toISOString(), + max_attempts: 5, + priority: 0, + created_at: new Date().toISOString(), + }; + const msg = { + ...createJobMessage(), + data: new TextEncoder().encode(JSON.stringify(envelope)), + }; + + await runner.runJob(TASK_TYPE, msg as unknown as JsMsg); + + expect(msg.ack).toHaveBeenCalledTimes(1); + expect(msg.working).not.toHaveBeenCalled(); + expect(await scheduledJobQueue.getQueueSize()).toBe(1); + }); +});