fix(worker): fit the job streams to the jetstream budget (#2458)

fluxer_api/src/api/worker/JetStreamWorkerQueue.ts

@@ -26,8 +26,13 @@ const DLQ_SUBJECT_PREFIX = 'dlq.';
 const DLQ_MAX_AGE_MS = 30 * 24 * 60 * 60 * 1000;
 const STREAM_MAX_MSGS = 2_000_000;
 const STREAM_MAX_BYTES = 8 * 1024 * 1024 * 1024;
+const STREAM_MIN_BYTES = 64 * 1024 * 1024;
 const STREAM_MAX_MSGS_PER_SUBJECT = 250_000;
-const STREAM_STORE_ERR_CODE = 10077;
+const DLQ_MAX_BYTES = 64 * 1024 * 1024;
+const DLQ_MIN_BYTES = 8 * 1024 * 1024;
+const STREAM_FULL_ERR_CODES = new Set([10023, 10077]);
+const STREAM_NO_STORAGE_ERR_CODE = 10047;
+const STREAM_NAME_IN_USE_ERR_CODE = 10058;
 
 const STREAM_LIMITS = {
 	max_msgs: STREAM_MAX_MSGS,
@@ -37,12 +42,19 @@ const STREAM_LIMITS = {
 	discard_new_per_subject: true,
 } satisfies Partial<StreamConfig>;
 
+function jsErrorCode(error: unknown): number | null {
+	if (!(error instanceof NatsError)) {
+		return null;
+	}
+	return error.jsError()?.err_code ?? null;
+}
+
 function describeStreamRejection(error: unknown): string | null {
 	if (!(error instanceof NatsError)) {
 		return null;
 	}
 	const apiError = error.jsError();
-	if (apiError?.err_code !== STREAM_STORE_ERR_CODE) {
+	if (apiError === null || !STREAM_FULL_ERR_CODES.has(apiError.err_code ?? 0)) {
 		return null;
 	}
 	return apiError.description ?? 'stream rejected the publish';
@@ -63,45 +75,116 @@ export class JetStreamWorkerQueue {
 			return;
 		}
 		const jsm = await this.connectionManager.getJetStreamManager();
-		let existingConfig: StreamConfig | null = null;
-		try {
-			existingConfig = (await jsm.streams.info(STREAM_NAME)).config;
-		} catch {
-			existingConfig = null;
-		}
+		const existingConfig = await this.readStreamConfig(jsm, STREAM_NAME);
 		if (existingConfig === null) {
-			await jsm.streams.add({
-				name: STREAM_NAME,
-				subjects: [`${SUBJECT_PREFIX}>`],
-				retention: RetentionPolicy.Workqueue,
-				storage: StorageType.File,
-				max_age: nanos(MAX_AGE_MS),
-				duplicate_window: nanos(2 * 60 * 1000),
-				num_replicas: 1,
-				...STREAM_LIMITS,
-			});
+			await this.addStream(jsm);
 		} else if (!this.hasStreamLimits(existingConfig)) {
-			await this.applyStreamLimits(jsm);
+			await this.applyStreamLimits(jsm, existingConfig);
 		}
 		this.streamReady = true;
 	}
 
+	private async readStreamConfig(jsm: JetStreamManager, name: string): Promise<StreamConfig | null> {
+		try {
+			return (await jsm.streams.info(name)).config;
+		} catch {
+			return null;
+		}
+	}
+
+	private async fitToStorageBudget(
+		startBytes: number,
+		minBytes: number,
+		apply: (maxBytes: number) => Promise<unknown>,
+	): Promise<number> {
+		let maxBytes = startBytes;
+		for (;;) {
+			try {
+				await apply(maxBytes);
+				return maxBytes;
+			} catch (error) {
+				if (jsErrorCode(error) !== STREAM_NO_STORAGE_ERR_CODE || maxBytes <= minBytes) {
+					throw error;
+				}
+				maxBytes = Math.floor(maxBytes / 2);
+			}
+		}
+	}
+
+	private async adoptConcurrentStream(jsm: JetStreamManager, name: string, error: unknown): Promise<boolean> {
+		if (jsErrorCode(error) !== STREAM_NAME_IN_USE_ERR_CODE) {
+			return false;
+		}
+		const config = await this.readStreamConfig(jsm, name);
+		if (config === null) {
+			return false;
+		}
+		Logger.info({stream: name, max_bytes: config.max_bytes}, 'Stream was created concurrently, adopting it');
+		return true;
+	}
+
+	private async addStream(jsm: JetStreamManager): Promise<void> {
+		let maxBytes: number;
+		try {
+			maxBytes = await this.fitToStorageBudget(STREAM_MAX_BYTES, STREAM_MIN_BYTES, (bytes) =>
+				jsm.streams.add({
+					name: STREAM_NAME,
+					subjects: [`${SUBJECT_PREFIX}>`],
+					retention: RetentionPolicy.Workqueue,
+					storage: StorageType.File,
+					max_age: nanos(MAX_AGE_MS),
+					duplicate_window: nanos(2 * 60 * 1000),
+					num_replicas: 1,
+					...STREAM_LIMITS,
+					max_bytes: bytes,
+				}),
+			);
+		} catch (error) {
+			if (await this.adoptConcurrentStream(jsm, STREAM_NAME, error)) {
+				return;
+			}
+			if (jsErrorCode(error) === STREAM_NO_STORAGE_ERR_CODE) {
+				Logger.error(
+					{err: error, stream: STREAM_NAME, max_bytes: STREAM_MIN_BYTES},
+					'Jobs stream does not fit the JetStream storage budget at its smallest size, free disk space on the NATS store',
+				);
+			}
+			throw error;
+		}
+		if (maxBytes !== STREAM_MAX_BYTES) {
+			Logger.warn(
+				{stream: STREAM_NAME, max_bytes: maxBytes},
+				'Created the jobs stream below its target size to fit the JetStream storage budget',
+			);
+		}
+	}
+
 	private hasStreamLimits(config: StreamConfig): boolean {
 		return (
 			config.max_msgs === STREAM_MAX_MSGS &&
-			config.max_bytes === STREAM_MAX_BYTES &&
+			config.max_bytes > 0 &&
 			config.max_msgs_per_subject === STREAM_MAX_MSGS_PER_SUBJECT &&
 			config.discard === DiscardPolicy.New &&
 			config.discard_new_per_subject
 		);
 	}
 
-	private async applyStreamLimits(jsm: JetStreamManager): Promise<void> {
+	private async applyStreamLimits(jsm: JetStreamManager, existingConfig: StreamConfig): Promise<void> {
+		const startBytes = existingConfig.max_bytes > 0 ? existingConfig.max_bytes : STREAM_MAX_BYTES;
 		try {
-			await jsm.streams.update(STREAM_NAME, {...STREAM_LIMITS});
-			Logger.info({stream: STREAM_NAME, ...STREAM_LIMITS}, 'Applied jobs stream limits');
+			const maxBytes = await this.fitToStorageBudget(startBytes, STREAM_MIN_BYTES, (bytes) =>
+				jsm.streams.update(STREAM_NAME, {...STREAM_LIMITS, max_bytes: bytes}),
+			);
+			Logger.info({stream: STREAM_NAME, ...STREAM_LIMITS, max_bytes: maxBytes}, 'Applied jobs stream limits');
 		} catch (error) {
-			Logger.error({err: error, stream: STREAM_NAME}, 'Failed to apply jobs stream limits, stream stays unbounded');
+			if (jsErrorCode(error) !== STREAM_NO_STORAGE_ERR_CODE) {
+				Logger.error({err: error, stream: STREAM_NAME}, 'Failed to apply jobs stream limits');
+				return;
+			}
+			Logger.warn(
+				{err: error, stream: STREAM_NAME, max_bytes: existingConfig.max_bytes},
+				'JetStream has no room to bound the jobs stream, it keeps the limits it already has',
+			);
 		}
 	}
 
@@ -110,21 +193,45 @@ export class JetStreamWorkerQueue {
 			return;
 		}
 		const jsm = await this.connectionManager.getJetStreamManager();
-		try {
-			await jsm.streams.info(DLQ_STREAM_NAME);
-		} catch {
-			await jsm.streams.add({
-				name: DLQ_STREAM_NAME,
-				subjects: [`${DLQ_SUBJECT_PREFIX}>`],
-				retention: RetentionPolicy.Limits,
-				storage: StorageType.File,
-				max_age: nanos(DLQ_MAX_AGE_MS),
-				num_replicas: 1,
-			});
+		if ((await this.readStreamConfig(jsm, DLQ_STREAM_NAME)) === null) {
+			if (!(await this.addDlqStream(jsm))) {
+				return;
+			}
 		}
 		this.dlqStreamReady = true;
 	}
 
+	private async addDlqStream(jsm: JetStreamManager): Promise<boolean> {
+		try {
+			const maxBytes = await this.fitToStorageBudget(DLQ_MAX_BYTES, DLQ_MIN_BYTES, (bytes) =>
+				jsm.streams.add({
+					name: DLQ_STREAM_NAME,
+					subjects: [`${DLQ_SUBJECT_PREFIX}>`],
+					retention: RetentionPolicy.Limits,
+					storage: StorageType.File,
+					max_age: nanos(DLQ_MAX_AGE_MS),
+					num_replicas: 1,
+					max_bytes: bytes,
+					discard: DiscardPolicy.Old,
+				}),
+			);
+			Logger.info({stream: DLQ_STREAM_NAME, max_bytes: maxBytes}, 'Dead-letter stream created');
+			return true;
+		} catch (error) {
+			if (await this.adoptConcurrentStream(jsm, DLQ_STREAM_NAME, error)) {
+				return true;
+			}
+			if (jsErrorCode(error) !== STREAM_NO_STORAGE_ERR_CODE) {
+				throw error;
+			}
+			Logger.warn(
+				{err: error, stream: DLQ_STREAM_NAME},
+				'JetStream has no room for the dead-letter stream, failed jobs stay in the jobs stream until they expire',
+			);
+			return false;
+		}
+	}
+
 	async ensureConsumers(lanes: ReadonlyArray<WorkerLaneDefinition>): Promise<void> {
 		if (this.consumersReady) {
 			return;

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

@@ -7,9 +7,12 @@ import {JetStreamWorkerQueue} from '../JetStreamWorkerQueue';
 import {WORKER_LANES} from '../WorkerLaneConfig';
 import {WorkerQueueOverflowError} from '../WorkerQueueOverflowError';
 
+const GIB = 1024 * 1024 * 1024;
+const MIB = 1024 * 1024;
+
 const EXPECTED_LIMITS = {
 	max_msgs: 2_000_000,
-	max_bytes: 8 * 1024 * 1024 * 1024,
+	max_bytes: 8 * GIB,
 	max_msgs_per_subject: 250_000,
 	discard: DiscardPolicy.New,
 	discard_new_per_subject: true,
@@ -38,6 +41,22 @@ function streamLimitError(description: string): NatsError {
 	return error;
 }
 
+function serverResourceError(): NatsError {
+	const error = new NatsError('503', '503');
+	error.api_error = {code: 503, err_code: 10023, description: 'insufficient resources'};
+	return error;
+}
+
+function noStorageError(): NatsError {
+	const error = new NatsError('503', '503');
+	error.api_error = {code: 503, err_code: 10047, description: 'insufficient storage resources available'};
+	return error;
+}
+
+function storageBudget(budget: number): (config: Partial<StreamConfig>) => Error | null {
+	return (config) => ((config.max_bytes ?? 0) > budget ? noStorageError() : null);
+}
+
 function boundedPublisher(maxPerSubject: number): (subject: string) => {seq: number} {
 	const counts = new Map<string, number>();
 	let seq = 0;
@@ -54,15 +73,19 @@ function boundedPublisher(maxPerSubject: number): (subject: string) => {seq: num
 
 function createQueue(params: {
 	existing?: StreamConfig | null;
+	dlqExists?: boolean;
+	reject?: (config: Partial<StreamConfig>) => Error | null;
 	updateError?: Error;
 	publish?: (subject: string) => {seq: number};
 }): {
 	queue: JetStreamWorkerQueue;
 	added: Array<Partial<StreamConfig>>;
+	dlqAdded: Array<Partial<StreamConfig>>;
 	updated: Array<Partial<StreamConfig>>;
 	consumerAdds: Array<ConsumerAddConfig>;
 } {
 	const added: Array<Partial<StreamConfig>> = [];
+	const dlqAdded: Array<Partial<StreamConfig>> = [];
 	const updated: Array<Partial<StreamConfig>> = [];
 	const consumerAdds: Array<ConsumerAddConfig> = [];
 	const connectionManager = {
@@ -77,14 +100,23 @@ function createQueue(params: {
 					info: () => Promise.reject(new Error('consumer not found')),
 				},
 				streams: {
-					info: () => {
+					info: (name: string) => {
+						if (name === 'JOBS_DLQ') {
+							return params.dlqExists
+								? Promise.resolve({config: {name} as StreamConfig})
+								: Promise.reject(new Error('stream not found'));
+						}
 						if (!params.existing) {
 							return Promise.reject(new Error('stream not found'));
 						}
 						return Promise.resolve({config: params.existing});
 					},
 					add: (config: Partial<StreamConfig>) => {
-						added.push(config);
+						(config.name === 'JOBS_DLQ' ? dlqAdded : added).push(config);
+						const rejection = params.reject?.(config) ?? null;
+						if (rejection !== null) {
+							return Promise.reject(rejection);
+						}
 						return Promise.resolve({config});
 					},
 					update: (_name: string, config: Partial<StreamConfig>) => {
@@ -92,6 +124,10 @@ function createQueue(params: {
 							return Promise.reject(params.updateError);
 						}
 						updated.push(config);
+						const rejection = params.reject?.(config) ?? null;
+						if (rejection !== null) {
+							return Promise.reject(rejection);
+						}
 						return Promise.resolve({config});
 					},
 				},
@@ -103,7 +139,7 @@ function createQueue(params: {
 			},
 		}),
 	} as unknown as JetStreamConnectionManager;
-	return {queue: new JetStreamWorkerQueue(connectionManager), added, updated, consumerAdds};
+	return {queue: new JetStreamWorkerQueue(connectionManager), added, dlqAdded, updated, consumerAdds};
 }
 
 describe('jobs stream limits', () => {
@@ -116,6 +152,28 @@ describe('jobs stream limits', () => {
 		expect(added[0]).toMatchObject({retention: RetentionPolicy.Workqueue, storage: StorageType.File});
 	});
 
+	it('lowers the new stream until the server storage budget accepts it', async () => {
+		const {queue, added} = createQueue({existing: null, reject: storageBudget(2 * GIB)});
+		await queue.ensureStream();
+		expect(added.map((config) => config.max_bytes)).toEqual([8 * GIB, 4 * GIB, 2 * GIB]);
+		expect(added[2]).toMatchObject({...EXPECTED_LIMITS, max_bytes: 2 * GIB});
+		expect(added[2]).toMatchObject({retention: RetentionPolicy.Workqueue, storage: StorageType.File});
+	});
+
+	it('fails the boot when even the smallest jobs stream does not fit', async () => {
+		const {queue, added} = createQueue({existing: null, reject: storageBudget(0)});
+		await expect(queue.ensureStream()).rejects.toBeInstanceOf(NatsError);
+		expect(added).toHaveLength(8);
+		expect(added[7]?.max_bytes).toBe(64 * MIB);
+	});
+
+	it('rethrows stream creation failures that are not storage rejections', async () => {
+		const failure = new Error('stream name already in use');
+		const {queue, added} = createQueue({existing: null, reject: () => failure});
+		await expect(queue.ensureStream()).rejects.toBe(failure);
+		expect(added).toHaveLength(1);
+	});
+
 	it('applies the limits to an existing unbounded stream', async () => {
 		const {queue, added, updated} = createQueue({existing: LEGACY_CONFIG});
 		await queue.ensureStream();
@@ -124,6 +182,12 @@ describe('jobs stream limits', () => {
 		expect(updated[0]).toEqual(EXPECTED_LIMITS);
 	});
 
+	it('lowers the limit update until the server storage budget accepts it', async () => {
+		const {queue, updated} = createQueue({existing: LEGACY_CONFIG, reject: storageBudget(1 * GIB)});
+		await queue.ensureStream();
+		expect(updated.map((config) => config.max_bytes)).toEqual([8 * GIB, 4 * GIB, 2 * GIB, 1 * GIB]);
+	});
+
 	it('leaves an already bounded stream alone', async () => {
 		const {queue, added, updated} = createQueue({existing: {...LEGACY_CONFIG, ...EXPECTED_LIMITS}});
 		await queue.ensureStream();
@@ -131,10 +195,39 @@ describe('jobs stream limits', () => {
 		expect(updated).toHaveLength(0);
 	});
 
+	it('leaves a stream that was lowered to fit the server alone', async () => {
+		const {queue, added, updated} = createQueue({
+			existing: {...LEGACY_CONFIG, ...EXPECTED_LIMITS, max_bytes: 1536 * MIB},
+		});
+		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();
 	});
+
+	it('keeps startup alive when no limit update fits the server', async () => {
+		const {queue, updated} = createQueue({existing: LEGACY_CONFIG, reject: storageBudget(0)});
+		await expect(queue.ensureStream()).resolves.toBeUndefined();
+		expect(updated).toHaveLength(8);
+	});
+});
+
+describe('dead-letter stream', () => {
+	it('keeps startup alive when the dead-letter stream does not fit', async () => {
+		const {queue, dlqAdded} = createQueue({dlqExists: false, reject: () => noStorageError()});
+		await expect(queue.ensureDlqStream()).resolves.toBeUndefined();
+		expect(dlqAdded).toHaveLength(4);
+	});
+
+	it('rethrows dead-letter creation failures that are not storage rejections', async () => {
+		const failure = new Error('no responders');
+		const {queue} = createQueue({dlqExists: false, reject: () => failure});
+		await expect(queue.ensureDlqStream()).rejects.toBe(failure);
+	});
 });
 
 describe('jobs stream enqueue shedding', () => {
@@ -152,6 +245,16 @@ describe('jobs stream enqueue shedding', () => {
 		await expect(queue.enqueue('handleMentions', {})).resolves.toBe('2');
 	});
 
+	it('sheds enqueues the server refuses for lack of resources', async () => {
+		const {queue} = createQueue({
+			existing: LEGACY_CONFIG,
+			publish: () => {
+				throw serverResourceError();
+			},
+		});
+		await expect(queue.enqueue('handleMentions', {})).rejects.toBeInstanceOf(WorkerQueueOverflowError);
+	});
+
 	it('rethrows publish failures that are not stream limits', async () => {
 		const failure = new Error('no responders');
 		const {queue} = createQueue({