From 3b996a7da0137a75c3047656a4bc8cc336ebfc1e Mon Sep 17 00:00:00 2001
From: Tomi Turtiainen <10324676+tomi@users.noreply.github.com>
Date: Mon, 8 Jan 2024 17:46:45 +0200
Subject: [PATCH] feat(core): Validate shutdown handlers on startup (#8260)

---
 packages/cli/src/commands/BaseCommand.ts      |  8 ++++-
 packages/cli/src/shutdown/Shutdown.service.ts | 20 ++++++++++++
 .../unit/shutdown/Shutdown.service.test.ts    | 32 +++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/packages/cli/src/commands/BaseCommand.ts b/packages/cli/src/commands/BaseCommand.ts
index f6091eb01..ff664505f 100644
--- a/packages/cli/src/commands/BaseCommand.ts
+++ b/packages/cli/src/commands/BaseCommand.ts
@@ -9,7 +9,7 @@ import { Logger } from '@/Logger';
 import config from '@/config';
 import * as Db from '@/Db';
 import * as CrashJournal from '@/CrashJournal';
-import { LICENSE_FEATURES, inTest } from '@/constants';
+import { LICENSE_FEATURES, inDevelopment, inTest } from '@/constants';
 import { initErrorHandling } from '@/ErrorReporting';
 import { ExternalHooks } from '@/ExternalHooks';
 import { NodeTypes } from '@/NodeTypes';
@@ -63,6 +63,12 @@ export abstract class BaseCommand extends Command {
 			this.exitWithCrash('There was an error initializing DB', error),
 		);
 
+		// This needs to happen after DB.init() or otherwise DB Connection is not
+		// available via the dependency Container that services depend on.
+		if (inDevelopment || inTest) {
+			this.shutdownService.validate();
+		}
+
 		await this.server?.init();
 
 		await Db.migrate().catch(async (error: Error) =>
diff --git a/packages/cli/src/shutdown/Shutdown.service.ts b/packages/cli/src/shutdown/Shutdown.service.ts
index 3c6608ecc..14b3367c4 100644
--- a/packages/cli/src/shutdown/Shutdown.service.ts
+++ b/packages/cli/src/shutdown/Shutdown.service.ts
@@ -39,6 +39,26 @@ export class ShutdownService {
 		this.handlersByPriority[priority].push(handler);
 	}
 
+	/** Validates that all the registered shutdown handlers are properly configured */
+	validate() {
+		const handlers = this.handlersByPriority.flat();
+
+		for (const { serviceClass, methodName } of handlers) {
+			if (!Container.has(serviceClass)) {
+				throw new ApplicationError(
+					`Component "${serviceClass.name}" is not registered with the DI container. Any component using @OnShutdown() must be decorated with @Service()`,
+				);
+			}
+
+			const service = Container.get(serviceClass);
+			if (!service[methodName]) {
+				throw new ApplicationError(
+					`Component "${serviceClass.name}" does not have a "${methodName}" method`,
+				);
+			}
+		}
+	}
+
 	/** Signals all registered listeners that the application is shutting down */
 	shutdown() {
 		if (this.shutdownPromise) {
diff --git a/packages/cli/test/unit/shutdown/Shutdown.service.test.ts b/packages/cli/test/unit/shutdown/Shutdown.service.test.ts
index d0f761524..b253929cb 100644
--- a/packages/cli/test/unit/shutdown/Shutdown.service.test.ts
+++ b/packages/cli/test/unit/shutdown/Shutdown.service.test.ts
@@ -124,4 +124,36 @@ describe('ShutdownService', () => {
 			expect(shutdownService.isShuttingDown()).toBe(false);
 		});
 	});
+
+	describe('validate', () => {
+		it('should throw error if component is not registered with the DI container', () => {
+			class UnregisteredComponent {
+				onShutdown() {}
+			}
+
+			shutdownService.register(10, {
+				serviceClass: UnregisteredComponent as unknown as ServiceClass,
+				methodName: 'onShutdown',
+			});
+
+			expect(() => shutdownService.validate()).toThrow(
+				'Component "UnregisteredComponent" is not registered with the DI container. Any component using @OnShutdown() must be decorated with @Service()',
+			);
+		});
+
+		it('should throw error if component is missing the shutdown method', () => {
+			class TestComponent {}
+
+			shutdownService.register(10, {
+				serviceClass: TestComponent as unknown as ServiceClass,
+				methodName: 'onShutdown',
+			});
+
+			Container.set(TestComponent, new TestComponent());
+
+			expect(() => shutdownService.validate()).toThrow(
+				'Component "TestComponent" does not have a "onShutdown" method',
+			);
+		});
+	});
 });