feat: multi server check and app close

app/src/__tests__/app.helper.spec.ts

@@ -10,12 +10,32 @@ import {
 } from '@jest/globals';
 
 import { createMock, DeepMocked } from '@golevelup/ts-jest';
-import { DynamicModule, INestApplicationContext } from '@nestjs/common';
+import { DynamicModule, INestApplicationContext, Module } from '@nestjs/common';
 import { NestFactory } from '@nestjs/core';
 import {
   curriedExecuteStandaloneFunction,
   executeStandaloneFunction,
 } from '../app.helper.js';
+import { getBootstrappedApps } from '../app-bootstrap-base.helper.js';
+import { AppBootstrap } from '../app-bootstrap.helper.js';
+import { yalcBaseAppModuleMetadataFactory } from '../base-app-module.helper.js';
+import { EventModule } from '@nestjs-yalc/event-manager/event.module.js';
+
+@Module(
+  yalcBaseAppModuleMetadataFactory(TestModule1, 'test1', {
+    configFactory: () => ({}),
+    logger: true,
+  }),
+)
+class TestModule1 {}
+
+@Module(
+  yalcBaseAppModuleMetadataFactory(TestModule2, 'test2', {
+    configFactory: () => ({}),
+    logger: true,
+  }),
+)
+class TestModule2 {}
 
 describe('test standalone app functions', () => {
   let mockedModule: DeepMocked<DynamicModule>;
@@ -25,7 +45,9 @@ describe('test standalone app functions', () => {
     jest.resetAllMocks();
     mockedServiceFunction = jest.fn(() => 'Test');
 
-    mockedModule = createMock<DynamicModule>({});
+    mockedModule = createMock<DynamicModule>({
+      imports: [],
+    });
 
     const mockedCreateApplicationContext = jest.spyOn(
       NestFactory,
@@ -40,6 +62,11 @@ describe('test standalone app functions', () => {
     );
   });
 
+  afterEach(() => {
+    getBootstrappedApps().forEach((app) => app.closeApp());
+    getBootstrappedApps().clear();
+  });
+
   it('should run executeStandaloneFunction', async () => {
     const mockedFunction = jest.fn(async (service: any) => {
       return service;
@@ -50,7 +77,7 @@ describe('test standalone app functions', () => {
       mockedServiceFunction,
       mockedFunction,
       {},
-      {closeApp: true}
+      { closeApp: true },
     );
 
     expect(mockedFunction).toHaveBeenCalledTimes(1);
@@ -89,4 +116,38 @@ describe('test standalone app functions', () => {
 
     expect(mockedFunction).toHaveBeenCalledTimes(1);
   });
+
+  it('should trigger an error if we bootstrap multiple servers', async () => {
+    await new AppBootstrap('test1', TestModule1).initApp();
+
+    let error: any = null;
+    try {
+      await new AppBootstrap('test2', TestModule2, {}).initApp();
+    } catch (e: any) {
+      // eslint-disable-next-line no-console
+      console.log(e);
+      error = e;
+    }
+
+    await expect(error).not.toBe(null);
+  });
+
+  it('should not trigger an error if we bootstrap multiple servers with the skip option', async () => {
+    await new AppBootstrap('test1', TestModule1, {
+      skipMultiServerCheck: true,
+    }).initApp();
+
+    let error: any = null;
+    try {
+      await new AppBootstrap('test2', TestModule2, {
+        skipMultiServerCheck: true,
+      }).initApp();
+    } catch (e: any) {
+      // eslint-disable-next-line no-console
+      console.log(e);
+      error = e;
+    }
+
+    await expect(error).toBe(null);
+  });
 });

app/src/app-bootstrap-base.helper.ts

@@ -16,6 +16,7 @@ import { EventModule } from '@nestjs-yalc/event-manager/event.module.js';
 import { LoggerServiceFactory } from '@nestjs-yalc/logger/logger.service.js';
 import { FastifyInstance } from 'fastify';
 import { getEnvLoggerLevels } from '@nestjs-yalc/logger/logger.helper.js';
+import { globalPromiseTracker } from '@nestjs-yalc/utils/promise.helper.js';
 
 /**
  * Side effect to be executed as soon as the module is imported
@@ -26,14 +27,38 @@ export interface IGlobalOptions {
   extraImports?: NonNullable<DynamicModule['imports']>;
   eventModuleClass?: typeof EventModule;
   logger?: typeof LoggerServiceFactory;
+
+  /**
+   * This is used to avoid bootstrapping a multi-server app
+   */
+  skipMultiServerCheck?: boolean;
 }
 
+const bootstrappedApps: Set<any> = new Set();
+
+export const getBootstrappedApps = () => {
+  return bootstrappedApps;
+};
+
+export const getMainBootstrappedApp = <
+  TApp extends BaseAppBootstrap<
+    NestFastifyApplication | INestApplicationContext
+  >,
+>(): TApp | null => {
+  if (getBootstrappedApps().size === 0) {
+    return null;
+  }
+
+  return getBootstrappedApps().entries().next().value;
+};
+
 export abstract class BaseAppBootstrap<
   TAppType extends NestFastifyApplication | INestApplicationContext,
 > {
   protected app?: TAppType;
   protected loggerService!: LoggerService;
   protected module: Type<any> | DynamicModule;
+  private isClosed = false;
 
   constructor(
     protected appAlias: string,
@@ -45,6 +70,16 @@ export abstract class BaseAppBootstrap<
       [appModule, ...(options?.globalsOptions?.extraImports ?? [])],
       options?.globalsOptions,
     );
+
+    const bootstrappedApp = getMainBootstrappedApp();
+
+    if (bootstrappedApp && !options?.globalsOptions?.skipMultiServerCheck) {
+      throw new Error(
+        'You are trying to bootstrap multiple servers in the same process. This is not allowed. Use a different process for each server',
+      );
+    }
+
+    getBootstrappedApps().add(this);
   }
 
   async initApp(options?: {
@@ -58,9 +93,34 @@ export abstract class BaseAppBootstrap<
   setApp(app: TAppType) {
     this.app = app;
 
+    /**
+     * Monkey patch the close method to set the isClosed flag
+     */
+    const originalCloseFn = this.app.close.bind(this.app);
+    this.app.close = async () => {
+      this.isClosed = true;
+      const closeRes = await originalCloseFn();
+      getBootstrappedApps().delete(this);
+      return closeRes;
+    };
+
+    /**
+     * Monkey patch the init method to set the isClosed flag
+     */
+    const originalInitFn = this.app.init.bind(this.app);
+    this.app.init = async () => {
+      this.isClosed = false;
+      getBootstrappedApps().add(this);
+      return originalInitFn();
+    };
+
     return this;
   }
 
+  isAppClosed() {
+    return this.isClosed;
+  }
+
   getAppAlias() {
     return this.appAlias;
   }
@@ -78,6 +138,24 @@ export abstract class BaseAppBootstrap<
     return this.app;
   }
 
+  async closeApp() {
+    await this.cleanup();
+
+    await this.app?.close();
+
+    getBootstrappedApps().delete(this);
+    this.isClosed = true;
+  }
+
+  async cleanup() {
+    /**
+     * When running behind a lambda, we have to await for all the promises that have been added to the global promise tracker
+     * to avoid them being killed by the lambda
+     * @see - https://stackoverflow.com/questions/64688812/running-tasks-in-aws-lambda-background
+     */
+    await globalPromiseTracker.waitForAll();
+  }
+
   /**
    *
    * @returns The main module of the business logic (the one that is passed in the constructor)

app/src/app-bootstrap.helper.ts

@@ -55,7 +55,7 @@ export class AppBootstrap<
     await this.initApp(options);
 
     if (envIsTrue(process.env.APP_DRY_RUN) === true) {
-      await this.getApp().close();
+      await this.closeApp();
       process.exit(0);
     }
 
@@ -91,7 +91,7 @@ export class AppBootstrap<
 
     if (envIsTrue(process.env.APP_DRY_RUN) === true) {
       this.loggerService?.log('Dry run, exiting...');
-      await this.getApp().close();
+      await this.closeApp();
       process.exit(0);
     }
 
@@ -231,7 +231,7 @@ export class AppBootstrap<
       // eslint-disable-next-line no-console
       console.debug('Hot reload enabled. Reloading...');
       hmr.accept();
-      hmr.dispose(() => this.getApp().close());
+      hmr.dispose(() => this.closeApp());
     }
   }
 }

app/src/app.helper.ts

@@ -1,26 +1,30 @@
 import { ClassType } from '@nestjs-yalc/types/globals.d.js';
-import { DynamicModule, INestApplicationContext, Type } from '@nestjs/common';
+import { DynamicModule, Type } from '@nestjs/common';
 import { StandaloneAppBootstrap } from './app-bootstrap-standalone.helper.js';
 import lodash from 'lodash';
-import { IGlobalOptions } from './app-bootstrap-base.helper.js';
+import {
+  BaseAppBootstrap,
+  IGlobalOptions,
+} from './app-bootstrap-base.helper.js';
 const { curry } = lodash;
 
 export function isDynamicModule(module: any): module is DynamicModule {
   return module.module !== undefined;
 }
 
 export const executeFunctionForApp = async (
-  app: INestApplicationContext,
+  app: BaseAppBootstrap<any>,
   serviceType: any,
   fn: { (service: any): Promise<any> },
   options: { closeApp?: boolean },
 ): Promise<void> => {
-  await app.init();
+  const nestApp = await app.getApp();
+  await nestApp.init();
 
-  const service = await app.resolve(serviceType);
+  const service = await nestApp.resolve(serviceType);
 
   await fn(service).finally(async () => {
-    if (options.closeApp) await app.close();
+    if (options.closeApp) await app.closeApp();
   });
 };
 
@@ -38,13 +42,11 @@ export const curriedExecuteStandaloneFunction = async <
   options?: TOptions,
 ) =>
   curry(executeFunctionForApp)(
-    (
-      await new StandaloneAppBootstrap(
-        isDynamicModule(module) ? module.module.name : module.name,
-        module,
-        options,
-      ).initApp()
-    ).getApp(),
+    await new StandaloneAppBootstrap(
+      isDynamicModule(module) ? module.module.name : module.name,
+      module,
+      options,
+    ).initApp(),
   );
 
 /**