fix(worker): bound the jobs stream and shed on overflow (#2204)

fluxer_api/src/api/worker/tests/JetStreamWorkerQueue.test.ts

@@ -0,0 +1,145 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {JetStreamConnectionManager} from '@pkgs/nats/src/JetStreamConnectionManager';
+import {DiscardPolicy, NatsError, RetentionPolicy, StorageType, type StreamConfig} from 'nats';
+import {describe, expect, it} from 'vitest';
+import {JetStreamWorkerQueue} from '../JetStreamWorkerQueue';
+import {WorkerQueueOverflowError} from '../WorkerQueueOverflowError';
+
+const EXPECTED_LIMITS = {
+	max_msgs: 2_000_000,
+	max_bytes: 8 * 1024 * 1024 * 1024,
+	max_msgs_per_subject: 250_000,
+	discard: DiscardPolicy.New,
+	discard_new_per_subject: true,
+};
+
+const LEGACY_CONFIG = {
+	name: 'JOBS',
+	subjects: ['jobs.>'],
+	retention: RetentionPolicy.Workqueue,
+	storage: StorageType.File,
+	max_msgs: -1,
+	max_bytes: -1,
+	max_msgs_per_subject: -1,
+	discard: DiscardPolicy.Old,
+	discard_new_per_subject: false,
+} as unknown as StreamConfig;
+
+function streamLimitError(description: string): NatsError {
+	const error = new NatsError('503', '503');
+	error.api_error = {code: 503, err_code: 10077, description};
+	return error;
+}
+
+function boundedPublisher(maxPerSubject: number): (subject: string) => {seq: number} {
+	const counts = new Map<string, number>();
+	let seq = 0;
+	return (subject) => {
+		const stored = counts.get(subject) ?? 0;
+		if (stored >= maxPerSubject) {
+			throw streamLimitError('maximum messages per subject exceeded');
+		}
+		counts.set(subject, stored + 1);
+		seq += 1;
+		return {seq};
+	};
+}
+
+function createQueue(params: {
+	existing?: StreamConfig | null;
+	updateError?: Error;
+	publish?: (subject: string) => {seq: number};
+}): {queue: JetStreamWorkerQueue; added: Array<Partial<StreamConfig>>; updated: Array<Partial<StreamConfig>>} {
+	const added: Array<Partial<StreamConfig>> = [];
+	const updated: Array<Partial<StreamConfig>> = [];
+	const connectionManager = {
+		getJetStreamManager: () =>
+			Promise.resolve({
+				streams: {
+					info: () => {
+						if (!params.existing) {
+							return Promise.reject(new Error('stream not found'));
+						}
+						return Promise.resolve({config: params.existing});
+					},
+					add: (config: Partial<StreamConfig>) => {
+						added.push(config);
+						return Promise.resolve({config});
+					},
+					update: (_name: string, config: Partial<StreamConfig>) => {
+						if (params.updateError) {
+							return Promise.reject(params.updateError);
+						}
+						updated.push(config);
+						return Promise.resolve({config});
+					},
+				},
+			}),
+		getJetStreamClient: () => ({
+			publish: (subject: string) => {
+				const publish = params.publish ?? (() => ({seq: 1}));
+				return Promise.resolve(publish(subject));
+			},
+		}),
+	} as unknown as JetStreamConnectionManager;
+	return {queue: new JetStreamWorkerQueue(connectionManager), added, updated};
+}
+
+describe('jobs stream limits', () => {
+	it('creates the stream bounded and discarding new messages', async () => {
+		const {queue, added, updated} = createQueue({existing: null});
+		await queue.ensureStream();
+		expect(updated).toHaveLength(0);
+		expect(added).toHaveLength(1);
+		expect(added[0]).toMatchObject(EXPECTED_LIMITS);
+		expect(added[0]).toMatchObject({retention: RetentionPolicy.Workqueue, storage: StorageType.File});
+	});
+
+	it('applies the limits to an existing unbounded stream', async () => {
+		const {queue, added, updated} = createQueue({existing: LEGACY_CONFIG});
+		await queue.ensureStream();
+		expect(added).toHaveLength(0);
+		expect(updated).toHaveLength(1);
+		expect(updated[0]).toEqual(EXPECTED_LIMITS);
+	});
+
+	it('leaves an already bounded stream alone', async () => {
+		const {queue, added, updated} = createQueue({existing: {...LEGACY_CONFIG, ...EXPECTED_LIMITS}});
+		await queue.ensureStream();
+		expect(added).toHaveLength(0);
+		expect(updated).toHaveLength(0);
+	});
+
+	it('keeps startup alive when the limit update is rejected', async () => {
+		const {queue} = createQueue({existing: LEGACY_CONFIG, updateError: new Error('stream update rejected')});
+		await expect(queue.ensureStream()).resolves.toBeUndefined();
+	});
+});
+
+describe('jobs stream enqueue shedding', () => {
+	it('rejects enqueues once the stream is at its cap', async () => {
+		const {queue} = createQueue({existing: LEGACY_CONFIG, publish: boundedPublisher(2)});
+		await expect(queue.enqueue('extractEmbeds', {})).resolves.toBe('1');
+		await expect(queue.enqueue('extractEmbeds', {})).resolves.toBe('2');
+		await expect(queue.enqueue('extractEmbeds', {})).rejects.toBeInstanceOf(WorkerQueueOverflowError);
+	});
+
+	it('caps each task type independently', async () => {
+		const {queue} = createQueue({existing: LEGACY_CONFIG, publish: boundedPublisher(1)});
+		await expect(queue.enqueue('extractEmbeds', {})).resolves.toBe('1');
+		await expect(queue.enqueue('extractEmbeds', {})).rejects.toBeInstanceOf(WorkerQueueOverflowError);
+		await expect(queue.enqueue('handleMentions', {})).resolves.toBe('2');
+	});
+
+	it('rethrows publish failures that are not stream limits', async () => {
+		const failure = new Error('no responders');
+		const {queue} = createQueue({
+			existing: LEGACY_CONFIG,
+			publish: () => {
+				throw failure;
+			},
+		});
+		await expect(queue.enqueue('extractEmbeds', {})).rejects.toBe(failure);
+	});
+});

fluxer_api/src/api/worker/tests/PerMessageJobLedger.test.ts

@@ -9,6 +9,7 @@ import {MessageMentionService} from '../../channel/services/message/MessageMenti
 import {EmbedService} from '../../infrastructure/EmbedService';
 import type {Message} from '../../models/Message';
 import type {WorkerTaskName} from '../WorkerLaneConfig';
+import {WorkerQueueOverflowError} from '../WorkerQueueOverflowError';
 
 class RecordingWorkerService implements IWorkerService<WorkerTaskName> {
 	readonly jobs: Array<{taskType: WorkerTaskName; options: WorkerJobOptions | undefined}> = [];
@@ -31,6 +32,24 @@ class RecordingWorkerService implements IWorkerService<WorkerTaskName> {
 	}
 }
 
+class OverflowingWorkerService implements IWorkerService<WorkerTaskName> {
+	async addJob<TPayload extends WorkerJobPayload = WorkerJobPayload>(
+		taskType: WorkerTaskName,
+		_payload: TPayload,
+		_options?: WorkerJobOptions,
+	): Promise<bigint> {
+		throw new WorkerQueueOverflowError(taskType, 'maximum messages per subject exceeded');
+	}
+
+	async cancelJob(): Promise<boolean> {
+		return false;
+	}
+
+	async retryDeadLetterJob(): Promise<boolean> {
+		return false;
+	}
+}
+
 function makeMentionMessage(): Message {
 	return {
 		id: createMessageID(2n),
@@ -71,4 +90,28 @@ describe('per-message worker jobs', () => {
 		expect(workerService.jobs[0]!.taskType).toBe('extractEmbeds');
 		expect(workerService.jobs[0]!.options?.skipLedger).toBe(true);
 	});
+
+	it('drops mention fanout instead of failing the send when the jobs stream is full', async () => {
+		const mentionService = new MessageMentionService(
+			null as never,
+			null as never,
+			null as never,
+			new OverflowingWorkerService(),
+			null as never,
+		);
+		await expect(
+			mentionService.handleMentionTasks({
+				guildId: null,
+				message: makeMentionMessage(),
+				authorId: createUserID(1n),
+			}),
+		).resolves.toBeUndefined();
+	});
+
+	it('drops embed extraction instead of failing the send when the jobs stream is full', async () => {
+		const embedService = new EmbedService(null as never, null as never, null as never, new OverflowingWorkerService());
+		await expect(
+			embedService.enqueueUrlEmbedExtraction(createChannelID(3n), createMessageID(2n), null, 'block'),
+		).resolves.toBeUndefined();
+	});
 });