refactor(core): Simplify subscriber handler setters (#10896)

packages/cli/src/scaling/__tests__/subscriber.service.test.ts

@@ -9,6 +9,7 @@ import { Subscriber } from '../pubsub/subscriber.service';
 describe('Subscriber', () => {
 	beforeEach(() => {
 		config.set('executions.mode', 'queue');
+		jest.restoreAllMocks();
 	});
 
 	const client = mock<SingleNodeClient>();
@@ -47,14 +48,16 @@ describe('Subscriber', () => {
 		});
 	});
 
-	describe('setHandler', () => {
-		it('should set handler function', () => {
+	describe('setMessageHandler', () => {
+		it('should set message handler function for channel', () => {
 			const subscriber = new Subscriber(mock(), redisClientService);
+			const channel = 'n8n.commands';
 			const handlerFn = jest.fn();
 
-			subscriber.addMessageHandler(handlerFn);
+			subscriber.setMessageHandler(channel, handlerFn);
 
-			expect(client.on).toHaveBeenCalledWith('message', handlerFn);
+			// @ts-expect-error Private field
+			expect(subscriber.handlers).toEqual(new Map([[channel, handlerFn]]));
 		});
 	});
 });

packages/cli/src/scaling/constants.ts

@@ -2,6 +2,8 @@ export const QUEUE_NAME = 'jobs';
 
 export const JOB_TYPE_NAME = 'job';
 
+/** Pubsub channel for commands sent by a main process to workers or to other main processes. */
 export const COMMAND_PUBSUB_CHANNEL = 'n8n.commands';
 
+/** Pubsub channel for messages sent by workers in response to commands from main processes. */
 export const WORKER_RESPONSE_PUBSUB_CHANNEL = 'n8n.worker-response';

packages/cli/src/scaling/pubsub/pubsub.types.ts

@@ -4,15 +4,11 @@ import type { IWorkflowDb } from '@/interfaces';
 
 import type { COMMAND_PUBSUB_CHANNEL, WORKER_RESPONSE_PUBSUB_CHANNEL } from '../constants';
 
-/**
- * Pubsub channel used by scaling mode:
- *
- * - `n8n.commands` for messages sent by a main process to command workers or other main processes
- * - `n8n.worker-response` for messages sent by workers in response to commands from main processes
- */
-export type ScalingPubSubChannel =
-	| typeof COMMAND_PUBSUB_CHANNEL
-	| typeof WORKER_RESPONSE_PUBSUB_CHANNEL;
+/** Pubsub channel used by scaling mode. */
+export type PubSubChannel = typeof COMMAND_PUBSUB_CHANNEL | typeof WORKER_RESPONSE_PUBSUB_CHANNEL;
+
+/** Handler function for every message received via a `PubSubChannel`. */
+export type PubSubHandlerFn = (msg: string) => void;
 
 export type PubSubMessageMap = {
 	// #region Lifecycle

packages/cli/src/scaling/pubsub/subscriber.service.ts

@@ -5,7 +5,7 @@ import config from '@/config';
 import { Logger } from '@/logger';
 import { RedisClientService } from '@/services/redis-client.service';
 
-import type { ScalingPubSubChannel } from './pubsub.types';
+import type { PubSubHandlerFn, PubSubChannel } from './pubsub.types';
 
 /**
  * Responsible for subscribing to the pubsub channels used by scaling mode.
@@ -14,6 +14,8 @@ import type { ScalingPubSubChannel } from './pubsub.types';
 export class Subscriber {
 	private readonly client: SingleNodeClient | MultiNodeClient;
 
+	private readonly handlers: Map<PubSubChannel, PubSubHandlerFn> = new Map();
+
 	// #region Lifecycle
 
 	constructor(
@@ -26,6 +28,10 @@ export class Subscriber {
 		this.client = this.redisClientService.createClient({ type: 'subscriber(n8n)' });
 
 		this.client.on('error', (error) => this.logger.error(error.message));
+
+		this.client.on('message', (channel: PubSubChannel, message) => {
+			this.handlers.get(channel)?.(message);
+		});
 	}
 
 	getClient() {
@@ -41,7 +47,7 @@ export class Subscriber {
 
 	// #region Subscribing
 
-	async subscribe(channel: ScalingPubSubChannel) {
+	async subscribe(channel: PubSubChannel) {
 		await this.client.subscribe(channel, (error) => {
 			if (error) {
 				this.logger.error('Failed to subscribe to channel', { channel, cause: error });
@@ -52,8 +58,9 @@ export class Subscriber {
 		});
 	}
 
-	addMessageHandler(handlerFn: (channel: string, msg: string) => void) {
-		this.client.on('message', handlerFn);
+	/** Set the message handler function for a channel. */
+	setMessageHandler(channel: PubSubChannel, handlerFn: PubSubHandlerFn) {
+		this.handlers.set(channel, handlerFn);
 	}
 
 	// #endregion