refactor(core): Introduce @OnLeaderTakeover and @OnLeaderStepdown (#14940)

packages/@n8n/decorators/src/__tests__/on-multi-main-event.test.ts

@@ -0,0 +1,193 @@
+import { Container } from '@n8n/di';
+import { Service } from '@n8n/di';
+import { EventEmitter } from 'node:events';
+
+import { MultiMainMetadata } from '../multi-main-metadata';
+import { LEADER_TAKEOVER_EVENT_NAME, LEADER_STEPDOWN_EVENT_NAME } from '../multi-main-metadata';
+import { NonMethodError, OnLeaderStepdown, OnLeaderTakeover } from '../on-multi-main-event';
+
+class MockMultiMainSetup extends EventEmitter {
+	registerEventHandlers() {
+		const handlers = Container.get(MultiMainMetadata).getHandlers();
+
+		for (const { eventHandlerClass, methodName, eventName } of handlers) {
+			const instance = Container.get(eventHandlerClass);
+			this.on(eventName, async () => {
+				return await instance[methodName].call(instance);
+			});
+		}
+	}
+}
+
+let multiMainSetup: MockMultiMainSetup;
+let metadata: MultiMainMetadata;
+
+beforeEach(() => {
+	Container.reset();
+
+	metadata = new MultiMainMetadata();
+	Container.set(MultiMainMetadata, metadata);
+
+	multiMainSetup = new MockMultiMainSetup();
+});
+
+it('should register methods decorated with @OnLeaderTakeover', () => {
+	jest.spyOn(metadata, 'register');
+
+	@Service()
+	class TestService {
+		@OnLeaderTakeover()
+		async handleLeaderTakeover() {}
+	}
+
+	expect(metadata.register).toHaveBeenCalledWith({
+		eventName: LEADER_TAKEOVER_EVENT_NAME,
+		methodName: 'handleLeaderTakeover',
+		eventHandlerClass: TestService,
+	});
+});
+
+it('should register methods decorated with @OnLeaderStepdown', () => {
+	jest.spyOn(metadata, 'register');
+
+	@Service()
+	class TestService {
+		@OnLeaderStepdown()
+		async handleLeaderStepdown() {}
+	}
+
+	expect(metadata.register).toHaveBeenCalledTimes(1);
+	expect(metadata.register).toHaveBeenCalledWith({
+		eventName: LEADER_STEPDOWN_EVENT_NAME,
+		methodName: 'handleLeaderStepdown',
+		eventHandlerClass: TestService,
+	});
+});
+
+it('should throw an error if the decorated target is not a method', () => {
+	expect(() => {
+		@Service()
+		class TestService {
+			// @ts-expect-error Testing invalid code
+			@OnLeaderTakeover()
+			notAFunction = 'string';
+		}
+
+		new TestService();
+	}).toThrowError(NonMethodError);
+});
+
+it('should call decorated methods when events are emitted', async () => {
+	@Service()
+	class TestService {
+		takeoverCalled = false;
+
+		stepdownCalled = false;
+
+		@OnLeaderTakeover()
+		async handleLeaderTakeover() {
+			this.takeoverCalled = true;
+		}
+
+		@OnLeaderStepdown()
+		async handleLeaderStepdown() {
+			this.stepdownCalled = true;
+		}
+	}
+
+	const testService = Container.get(TestService);
+	jest.spyOn(testService, 'handleLeaderTakeover');
+	jest.spyOn(testService, 'handleLeaderStepdown');
+
+	multiMainSetup.registerEventHandlers();
+
+	multiMainSetup.emit(LEADER_TAKEOVER_EVENT_NAME);
+	multiMainSetup.emit(LEADER_STEPDOWN_EVENT_NAME);
+
+	expect(testService.handleLeaderTakeover).toHaveBeenCalledTimes(1);
+	expect(testService.handleLeaderStepdown).toHaveBeenCalledTimes(1);
+	expect(testService.takeoverCalled).toBe(true);
+	expect(testService.stepdownCalled).toBe(true);
+});
+
+it('should register multiple handlers for the same event', async () => {
+	@Service()
+	class TestService {
+		firstHandlerCalled = false;
+
+		secondHandlerCalled = false;
+
+		@OnLeaderTakeover()
+		async firstHandler() {
+			this.firstHandlerCalled = true;
+		}
+
+		@OnLeaderTakeover()
+		async secondHandler() {
+			this.secondHandlerCalled = true;
+		}
+	}
+
+	const testService = Container.get(TestService);
+
+	multiMainSetup.registerEventHandlers();
+
+	multiMainSetup.emit(LEADER_TAKEOVER_EVENT_NAME);
+
+	expect(testService.firstHandlerCalled).toBe(true);
+	expect(testService.secondHandlerCalled).toBe(true);
+});
+
+it('should register handlers from multiple service classes', async () => {
+	@Service()
+	class FirstService {
+		handlerCalled = false;
+
+		@OnLeaderTakeover()
+		async handleTakeover() {
+			this.handlerCalled = true;
+		}
+	}
+
+	@Service()
+	class SecondService {
+		handlerCalled = false;
+
+		@OnLeaderTakeover()
+		async handleTakeover() {
+			this.handlerCalled = true;
+		}
+	}
+
+	const firstService = Container.get(FirstService);
+	const secondService = Container.get(SecondService);
+
+	multiMainSetup.registerEventHandlers();
+
+	multiMainSetup.emit(LEADER_TAKEOVER_EVENT_NAME);
+
+	expect(firstService.handlerCalled).toBe(true);
+	expect(secondService.handlerCalled).toBe(true);
+});
+
+it('should handle async methods correctly', async () => {
+	@Service()
+	class TestService {
+		result = '';
+
+		@OnLeaderTakeover()
+		async handleLeaderTakeover() {
+			await new Promise((resolve) => setTimeout(resolve, 10));
+			this.result = 'completed';
+		}
+	}
+
+	const testService = Container.get(TestService);
+
+	multiMainSetup.registerEventHandlers();
+	multiMainSetup.emit(LEADER_TAKEOVER_EVENT_NAME);
+
+	await new Promise((resolve) => setTimeout(resolve, 20));
+
+	expect(testService.result).toBe('completed');
+});

packages/@n8n/decorators/src/index.ts

@@ -18,3 +18,5 @@ export { BaseN8nModule, N8nModule } from './module';
 export { Debounce } from './debounce';
 export type { AccessScope, Controller, RateLimit } from './types';
 export type { ShutdownHandler } from './types';
+export { MultiMainMetadata } from './multi-main-metadata';
+export { OnLeaderTakeover, OnLeaderStepdown } from './on-multi-main-event';

packages/@n8n/decorators/src/module.ts

@@ -1,10 +1,4 @@
 import { Container, Service, type Constructable } from '@n8n/di';
-import type EventEmitter from 'node:events';
-
-/**
- * @TODO Temporary dummy type until `MultiMainSetup` registers listeners via decorators.
- */
-type MultiMainSetup = EventEmitter;
 
 /**
  * @TODO Temporary dummy type until `ExecutionLifecycleHooks` registers hooks via decorators.
@@ -14,7 +8,6 @@ export type ExecutionLifecycleHooks = object;
 export interface BaseN8nModule {
 	initialize?(): void;
 	registerLifecycleHooks?(hooks: ExecutionLifecycleHooks): void;
-	registerMultiMainListeners?(multiMainSetup: MultiMainSetup): void;
 }
 
 type Module = Constructable<BaseN8nModule>;
@@ -47,10 +40,4 @@ export class ModuleRegistry {
 			}
 		}
 	}
-
-	registerMultiMainListeners(multiMainSetup: MultiMainSetup) {
-		for (const ModuleClass of registry.keys()) {
-			Container.get(ModuleClass).registerMultiMainListeners?.(multiMainSetup);
-		}
-	}
 }

packages/@n8n/decorators/src/multi-main-metadata.ts

@@ -0,0 +1,36 @@
+import { Service } from '@n8n/di';
+
+import type { Class } from './types';
+
+export const LEADER_TAKEOVER_EVENT_NAME = 'leader-takeover';
+export const LEADER_STEPDOWN_EVENT_NAME = 'leader-stepdown';
+
+export type MultiMainEvent = typeof LEADER_TAKEOVER_EVENT_NAME | typeof LEADER_STEPDOWN_EVENT_NAME;
+
+type EventHandlerFn = () => Promise<void> | void;
+
+export type EventHandlerClass = Class<Record<string, EventHandlerFn>>;
+
+type EventHandler = {
+	/** Class holding the method to call on a multi-main event. */
+	eventHandlerClass: EventHandlerClass;
+
+	/** Name of the method to call on a multi-main event. */
+	methodName: string;
+
+	/** Name of the multi-main event to listen to. */
+	eventName: MultiMainEvent;
+};
+
+@Service()
+export class MultiMainMetadata {
+	private readonly handlers: EventHandler[] = [];
+
+	register(handler: EventHandler) {
+		this.handlers.push(handler);
+	}
+
+	getHandlers(): EventHandler[] {
+		return this.handlers;
+	}
+}

packages/@n8n/decorators/src/on-multi-main-event.ts

@@ -0,0 +1,66 @@
+import { Container } from '@n8n/di';
+import { UnexpectedError } from 'n8n-workflow';
+
+import type { EventHandlerClass, MultiMainEvent } from './multi-main-metadata';
+import {
+	LEADER_TAKEOVER_EVENT_NAME,
+	LEADER_STEPDOWN_EVENT_NAME,
+	MultiMainMetadata,
+} from './multi-main-metadata';
+
+export class NonMethodError extends UnexpectedError {
+	constructor(name: string) {
+		super(`${name} must be a method on a class to use this decorator`);
+	}
+}
+
+const OnMultiMainEvent =
+	(eventName: MultiMainEvent): MethodDecorator =>
+	(prototype, propertyKey, descriptor) => {
+		const eventHandlerClass = prototype.constructor as EventHandlerClass;
+		const methodName = String(propertyKey);
+
+		if (typeof descriptor?.value !== 'function') {
+			throw new NonMethodError(`${eventHandlerClass.name}.${methodName}()`);
+		}
+
+		Container.get(MultiMainMetadata).register({
+			eventHandlerClass,
+			methodName,
+			eventName,
+		});
+	};
+
+/**
+ * Decorator that registers a method to be called when this main instance becomes the leader.
+ *
+ * @example
+ *
+ * ```ts
+ * @Service()
+ * class MyService {
+ *   @OnLeaderTakeover()
+ *   async startDoingThings() {
+ *     // ...
+ *   }
+ * }
+ * ```
+ */
+export const OnLeaderTakeover = () => OnMultiMainEvent(LEADER_TAKEOVER_EVENT_NAME);
+
+/**
+ * Decorator that registers a method to be called when this main instance stops being the leader.
+ *
+ * @example
+ *
+ * ```ts
+ * @Service()
+ * class MyService {
+ *   @OnLeaderStepdown()
+ *   async stopDoingThings() {
+ *     // ...
+ *   }
+ * }
+ * ```
+ */
+export const OnLeaderStepdown = () => OnMultiMainEvent(LEADER_STEPDOWN_EVENT_NAME);

packages/@n8n/decorators/src/types.ts

@@ -50,7 +50,7 @@ export type Controller = Constructable<object> &
 
 type RouteHandlerFn = () => Promise<void> | void;
 
-type Class<T = object, A extends unknown[] = unknown[]> = new (...args: A) => T;
+export type Class<T = object, A extends unknown[] = unknown[]> = new (...args: A) => T;
 
 export type ServiceClass = Class<Record<string, RouteHandlerFn>>;