fix(worker): resubscribe when the job stream ends unexpectedly (#2282)

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<void> {
@@ -117,11 +110,54 @@ export class WorkerRunner {
 		Logger.info({workerId: this.workerId}, 'Worker stopped');
 	}
 
-	private async processMessages(): Promise<void> {
-		if (this.consumerMessages === null) {
-			return;
+	private async openConsumerMessages(): Promise<ConsumerMessages> {
+		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<void> {
+		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<void> {
+		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<boolean> {

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<JsMsg> = [];
+	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<void> {
+		this.end();
+	}
+
+	async *[Symbol.asyncIterator](): AsyncGenerator<JsMsg> {
+		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<void>((resolve) => {
+				this.notify = resolve;
+			});
+		}
+	}
+
+	private wake(): void {
+		const notify = this.notify;
+		this.notify = null;
+		notify?.();
+	}
+}
+
+function createRunner(streams: Array<FakeConsumerMessages>): {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);
+	});
+});