diff --git a/fluxer_api/src/api/worker/WorkerRunner.ts b/fluxer_api/src/api/worker/WorkerRunner.ts index 806e05c75..c920db248 100644 --- a/fluxer_api/src/api/worker/WorkerRunner.ts +++ b/fluxer_api/src/api/worker/WorkerRunner.ts @@ -11,6 +11,7 @@ import {isJsonRecord, parseJsonRecord} from '../utils/JsonBoundaryUtils'; const MAX_DLQ_PUBLISH_ATTEMPTS = 3; const MIN_ACK_HEARTBEAT_MS = 1000; +const RESUBSCRIBE_DELAY_MS = 5000; interface WorkerRunnerJetStreamClient { consumers: { @@ -89,16 +90,8 @@ export class WorkerRunner { } this.running = true; Logger.info({workerId: this.workerId, lane: this.laneName, concurrency: this.concurrency}, 'Worker starting'); - const js = this.queue.getConnectionManager().getJetStreamClient(); - const consumer = await js.consumers.get(this.queue.getStreamName(), this.consumerName); - const prefetch = Math.max(this.concurrency * 2, 16); - this.consumerMessages = await consumer.consume({ - max_messages: prefetch, - idle_heartbeat: 5000, - }); - this.processingLoop = this.processMessages().catch((error) => { - Logger.error({workerId: this.workerId, err: error}, 'Worker message processing failed unexpectedly'); - }); + this.consumerMessages = await this.openConsumerMessages(); + this.processingLoop = this.consumeUntilStopped(this.consumerMessages); } async stop(): Promise { @@ -117,11 +110,54 @@ export class WorkerRunner { Logger.info({workerId: this.workerId}, 'Worker stopped'); } - private async processMessages(): Promise { - if (this.consumerMessages === null) { - return; + private async openConsumerMessages(): Promise { + const js = this.queue.getConnectionManager().getJetStreamClient(); + const consumer = await js.consumers.get(this.queue.getStreamName(), this.consumerName); + const prefetch = Math.max(this.concurrency * 2, 16); + return await consumer.consume({ + max_messages: prefetch, + idle_heartbeat: 5000, + }); + } + + private async consumeUntilStopped(initialMessages: ConsumerMessages): Promise { + let messages: ConsumerMessages | null = initialMessages; + while (this.running) { + if (messages === null) { + await new Promise((resolve) => setTimeout(resolve, RESUBSCRIBE_DELAY_MS)); + if (!this.running) { + break; + } + try { + messages = await this.openConsumerMessages(); + this.consumerMessages = messages; + } catch (error) { + Logger.error( + {workerId: this.workerId, lane: this.laneName, err: error}, + 'Failed to resubscribe the worker consumer', + ); + } + continue; + } + try { + await this.processMessages(messages); + } catch (error) { + Logger.error({workerId: this.workerId, err: error}, 'Worker message processing failed unexpectedly'); + } + messages = null; + this.consumerMessages = null; + if (this.running) { + Logger.error( + {workerId: this.workerId, lane: this.laneName}, + 'Worker message stream ended while running, resubscribing', + ); + } } - for await (const msg of this.consumerMessages) { + Logger.info({workerId: this.workerId}, 'Worker message iterator ended'); + } + + private async processMessages(consumerMessages: ConsumerMessages): Promise { + for await (const msg of consumerMessages) { if (!this.running) { break; } @@ -159,7 +195,6 @@ export class WorkerRunner { this.inFlightJobs.add(jobPromise); } await Promise.allSettled(this.inFlightJobs); - Logger.info({workerId: this.workerId}, 'Worker message iterator ended'); } protected async processJob(taskType: string, msg: JsMsg): Promise { diff --git a/fluxer_api/src/api/worker/tests/WorkerRunnerResubscribe.test.ts b/fluxer_api/src/api/worker/tests/WorkerRunnerResubscribe.test.ts new file mode 100644 index 000000000..6845a16f7 --- /dev/null +++ b/fluxer_api/src/api/worker/tests/WorkerRunnerResubscribe.test.ts @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import type {ConsumerMessages, JsMsg} from 'nats'; +import {afterEach, beforeAll, describe, expect, it, vi} from 'vitest'; +import type {IJobLedgerRepository} from '../../jobs/IJobLedgerRepository'; +import {setInjectedWorkerService} from '../../middleware/ServiceRegistry'; +import {NoopWorkerService} from '../../test/NoopWorkerService'; +import {WorkerRunner} from '../WorkerRunner'; + +const TASK_TYPE = 'processInactivityDeletions'; +const RESUBSCRIBE_WINDOW_MS = 30000; + +class FakeConsumerMessages { + private readonly pending: Array = []; + private notify: (() => void) | null = null; + private closed = false; + private failure: Error | null = null; + + push(msg: JsMsg): void { + this.pending.push(msg); + this.wake(); + } + + abort(error: Error): void { + this.failure = error; + this.wake(); + } + + end(): void { + this.closed = true; + this.wake(); + } + + async close(): Promise { + this.end(); + } + + async *[Symbol.asyncIterator](): AsyncGenerator { + while (true) { + while (this.pending.length > 0) { + yield this.pending.shift()!; + } + if (this.failure !== null) { + throw this.failure; + } + if (this.closed) { + return; + } + await new Promise((resolve) => { + this.notify = resolve; + }); + } + } + + private wake(): void { + const notify = this.notify; + this.notify = null; + notify?.(); + } +} + +function createRunner(streams: Array): {runner: WorkerRunner; consumeCount: () => number} { + let consumed = 0; + const queue = { + getConnectionManager: () => ({ + getJetStreamClient: () => ({ + consumers: { + get: async () => ({ + consume: async () => { + const stream = streams[consumed]; + consumed += 1; + if (!stream) { + throw new Error('no more streams'); + } + return stream as unknown as ConsumerMessages; + }, + }), + }, + }), + }), + getStreamName: () => 'JOBS', + publishToDlq: vi.fn(), + }; + const runner = new WorkerRunner({ + tasks: {[TASK_TYPE]: async () => {}}, + queue, + consumerName: 'workers_batch', + laneName: 'batch', + ledger: {} as IJobLedgerRepository, + concurrency: 12, + maxDeliver: 25, + ackWaitMs: 120000, + }); + return {runner, consumeCount: () => consumed}; +} + +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(), + }; +} + +describe('Worker runner resubscribe', () => { + beforeAll(() => { + setInjectedWorkerService(new NoopWorkerService()); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('resubscribes after the message stream fails', async () => { + vi.useFakeTimers(); + const streams = [new FakeConsumerMessages(), new FakeConsumerMessages()]; + const {runner, consumeCount} = createRunner(streams); + + await runner.start(); + expect(consumeCount()).toBe(1); + + streams[0]!.abort(new Error('consumer deleted')); + await vi.advanceTimersByTimeAsync(RESUBSCRIBE_WINDOW_MS); + expect(consumeCount()).toBe(2); + + const msg = createJobMessage(); + streams[1]!.push(msg as unknown as JsMsg); + await vi.advanceTimersByTimeAsync(0); + expect(msg.ack).toHaveBeenCalledTimes(1); + + await runner.stop(); + }); + + it('resubscribes after the message stream ends without a stop', async () => { + vi.useFakeTimers(); + const streams = [new FakeConsumerMessages(), new FakeConsumerMessages()]; + const {runner, consumeCount} = createRunner(streams); + + await runner.start(); + streams[0]!.end(); + await vi.advanceTimersByTimeAsync(RESUBSCRIBE_WINDOW_MS); + expect(consumeCount()).toBe(2); + + const msg = createJobMessage(); + streams[1]!.push(msg as unknown as JsMsg); + await vi.advanceTimersByTimeAsync(0); + expect(msg.ack).toHaveBeenCalledTimes(1); + + await runner.stop(); + }); + + it('does not resubscribe after the runner is stopped', async () => { + vi.useFakeTimers(); + const streams = [new FakeConsumerMessages(), new FakeConsumerMessages()]; + const {runner, consumeCount} = createRunner(streams); + + await runner.start(); + await runner.stop(); + await vi.advanceTimersByTimeAsync(RESUBSCRIBE_WINDOW_MS); + + expect(consumeCount()).toBe(1); + }); +});