Merge pull request nestjs#3889 from nestjs/fix/lifecycle-hooks-3870

fix(core): fix lifecycle hooks for middleware, injectables

Author: kamilmysliwiec

integration/docker-compose.yml

@@ -48,7 +48,7 @@ services:
   zookeeper:
     container_name: test-zookeeper
     hostname: zookeeper
-    image: confluentinc/cp-zookeeper:5.4.0
+    image: confluentinc/cp-zookeeper:5.3.2
     ports:
       - "2181:2181"
     environment:
@@ -57,7 +57,7 @@ services:
   kafka:
     container_name: test-kafka
     hostname: kafka
-    image: confluentinc/cp-kafka:5.4.0
+    image: confluentinc/cp-kafka:5.3.2
     depends_on:
       - zookeeper
     ports:

packages/core/hooks/on-app-bootstrap.hook.ts

@@ -43,7 +43,12 @@ export async function callModuleBootstrapHook(module: Module): Promise<any> {
   // Module (class) instance is the first element of the providers array
   // Lifecycle hook has to be called once all classes are properly initialized
   const [_, { instance: moduleClassInstance }] = providers.shift();
-  const instances = [...module.controllers, ...providers];
+  const instances = [
+    ...module.controllers,
+    ...providers,
+    ...module.injectables,
+    ...module.middlewares,
+  ];
 
   const nonTransientInstances = getNonTransientInstances(instances);
   await Promise.all(callOperator(nonTransientInstances));

packages/core/hooks/on-app-shutdown.hook.ts

@@ -51,7 +51,12 @@ export async function callAppShutdownHook(
   // Module (class) instance is the first element of the providers array
   // Lifecycle hook has to be called once all classes are properly initialized
   const [_, { instance: moduleClassInstance }] = providers.shift();
-  const instances = [...module.controllers, ...providers];
+  const instances = [
+    ...module.controllers,
+    ...providers,
+    ...module.injectables,
+    ...module.middlewares,
+  ];
 
   const nonTransientInstances = getNonTransientInstances(instances);
   await Promise.all(callOperator(nonTransientInstances, signal));

packages/core/hooks/on-module-destroy.hook.ts

@@ -43,7 +43,12 @@ export async function callModuleDestroyHook(module: Module): Promise<any> {
   // Module (class) instance is the first element of the providers array
   // Lifecycle hook has to be called once all classes are properly destroyed
   const [_, { instance: moduleClassInstance }] = providers.shift();
-  const instances = [...module.controllers, ...providers];
+  const instances = [
+    ...module.controllers,
+    ...providers,
+    ...module.injectables,
+    ...module.middlewares,
+  ];
 
   const nonTransientInstances = getNonTransientInstances(instances);
   await Promise.all(callOperator(nonTransientInstances));

packages/core/hooks/on-module-init.hook.ts

@@ -39,7 +39,12 @@ export async function callModuleInitHook(module: Module): Promise<void> {
   // Module (class) instance is the first element of the providers array
   // Lifecycle hook has to be called once all classes are properly initialized
   const [_, { instance: moduleClassInstance }] = providers.shift();
-  const instances = [...module.controllers, ...providers];
+  const instances = [
+    ...module.controllers,
+    ...providers,
+    ...module.injectables,
+    ...module.middlewares,
+  ];
 
   const nonTransientInstances = getNonTransientInstances(instances);
   await Promise.all(callOperator(nonTransientInstances));

packages/core/injector/module.ts

@@ -39,6 +39,7 @@ export class Module {
   private readonly _imports = new Set<Module>();
   private readonly _providers = new Map<any, InstanceWrapper<Injectable>>();
   private readonly _injectables = new Map<any, InstanceWrapper<Injectable>>();
+  private readonly _middlewares = new Map<any, InstanceWrapper<Injectable>>();
   private readonly _controllers = new Map<
     string,
     InstanceWrapper<Controller>
@@ -51,7 +52,7 @@ export class Module {
     private readonly _scope: Type<any>[],
     private readonly container: NestContainer,
   ) {
-    this.addCoreProviders(container);
+    this.addCoreProviders();
     this._id = randomStringGenerator();
   }
 
@@ -67,6 +68,10 @@ export class Module {
     return this._providers;
   }
 
+  get middlewares(): Map<any, InstanceWrapper<Injectable>> {
+    return this._middlewares;
+  }
+
   get imports(): Set<Module> {
     return this._imports;
   }
@@ -124,7 +129,7 @@ export class Module {
     this._distance = value;
   }
 
-  public addCoreProviders(container: NestContainer) {
+  public addCoreProviders() {
     this.addModuleAsProvider();
     this.addModuleRef();
     this.addApplicationConfig();

packages/core/middleware/container.ts

@@ -1,6 +1,7 @@
 import { Scope, Type } from '@nestjs/common';
 import { SCOPE_OPTIONS_METADATA } from '@nestjs/common/constants';
 import { MiddlewareConfiguration } from '@nestjs/common/interfaces/middleware/middleware-configuration.interface';
+import { NestContainer } from '../injector';
 import { InstanceWrapper } from '../injector/instance-wrapper';
 
 export class MiddlewareContainer {
@@ -10,17 +11,28 @@ export class MiddlewareContainer {
     Set<MiddlewareConfiguration>
   >();
 
-  public getMiddlewareCollection(module: string): Map<string, InstanceWrapper> {
-    return this.middleware.get(module) || new Map();
+  constructor(private readonly container: NestContainer) {}
+
+  public getMiddlewareCollection(
+    moduleKey: string,
+  ): Map<string, InstanceWrapper> {
+    if (!this.middleware.has(moduleKey)) {
+      const moduleRef = this.container.getModuleByKey(moduleKey);
+      this.middleware.set(moduleKey, moduleRef.middlewares);
+    }
+    return this.middleware.get(moduleKey);
   }
 
   public getConfigurations(): Map<string, Set<MiddlewareConfiguration>> {
     return this.configurationSets;
   }
 
-  public insertConfig(configList: MiddlewareConfiguration[], module: string) {
-    const middleware = this.getTargetMiddleware(module);
-    const targetConfig = this.getTargetConfig(module);
+  public insertConfig(
+    configList: MiddlewareConfiguration[],
+    moduleKey: string,
+  ) {
+    const middleware = this.getMiddlewareCollection(moduleKey);
+    const targetConfig = this.getTargetConfig(moduleKey);
 
     const configurations = configList || [];
     const insertMiddleware = <T extends Type<any>>(metatype: T) => {
@@ -40,13 +52,6 @@ export class MiddlewareContainer {
     });
   }
 
-  private getTargetMiddleware(module: string) {
-    if (!this.middleware.has(module)) {
-      this.middleware.set(module, new Map<string, InstanceWrapper>());
-    }
-    return this.middleware.get(module);
-  }
-
   private getTargetConfig(module: string) {
     if (!this.configurationSets.has(module)) {
       this.configurationSets.set(module, new Set<MiddlewareConfiguration>());

packages/core/nest-application.ts

@@ -44,11 +44,12 @@ export class NestApplication extends NestApplicationContext
   implements INestApplication {
   private readonly logger = new Logger(NestApplication.name, true);
   private readonly middlewareModule = new MiddlewareModule();
-  private readonly middlewareContainer = new MiddlewareContainer();
-  private readonly microservicesModule = MicroservicesModule
-    ? new MicroservicesModule()
-    : null;
-  private readonly socketModule = SocketModule ? new SocketModule() : null;
+  private readonly middlewareContainer = new MiddlewareContainer(
+    this.container,
+  );
+  private readonly microservicesModule =
+    MicroservicesModule && new MicroservicesModule();
+  private readonly socketModule = SocketModule && new SocketModule();
   private readonly routesResolver: Resolver;
   private readonly microservices: any[] = [];
   private httpServer: any;

packages/core/test/middleware/container.spec.ts

@@ -5,10 +5,14 @@ import { RequestMapping } from '../../../common/decorators/http/request-mapping.
 import { RequestMethod } from '../../../common/enums/request-method.enum';
 import { MiddlewareConfiguration } from '../../../common/interfaces/middleware/middleware-configuration.interface';
 import { NestMiddleware } from '../../../common/interfaces/middleware/nest-middleware.interface';
+import { NestContainer } from '../../injector';
 import { InstanceWrapper } from '../../injector/instance-wrapper';
+import { Module } from '../../injector/module';
 import { MiddlewareContainer } from '../../middleware/container';
 
 describe('MiddlewareContainer', () => {
+  class ExampleModule {}
+
   @Controller('test')
   class TestRoute {
     @RequestMapping({ path: 'test' })
@@ -26,7 +30,13 @@ describe('MiddlewareContainer', () => {
   let container: MiddlewareContainer;
 
   beforeEach(() => {
-    container = new MiddlewareContainer();
+    const nestContainer = new NestContainer();
+    const modules = nestContainer.getModules();
+
+    modules.set('Module', new Module(ExampleModule, [], nestContainer));
+    modules.set('Test', new Module(ExampleModule, [], nestContainer));
+
+    container = new MiddlewareContainer(nestContainer);
   });
 
   it('should store expected configurations for given module', () => {
@@ -36,7 +46,7 @@ describe('MiddlewareContainer', () => {
         forRoutes: [TestRoute, 'test'],
       },
     ];
-    container.insertConfig(config, 'Module' as any);
+    container.insertConfig(config, 'Module');
     expect([...container.getConfigurations().get('Module')]).to.deep.equal(
       config,
     );
@@ -50,7 +60,7 @@ describe('MiddlewareContainer', () => {
       },
     ];
 
-    const key = 'Test' as any;
+    const key = 'Test';
     container.insertConfig(config, key);
 
     const collection = container.getMiddlewareCollection(key);

…st/middleware/middlewares-module.spec.ts …est/middleware/middleware-module.spec.tspackages/core/test/middleware/middlewares-module.spec.ts renamed to packages/core/test/middleware/middleware-module.spec.ts packages/core/test/middleware/middlewares-module.spec.ts renamed to packages/core/test/middleware/middleware-module.spec.ts

@@ -56,7 +56,7 @@ describe('MiddlewareModule', () => {
       };
 
       await middlewareModule.loadConfiguration(
-        new MiddlewareContainer(),
+        new MiddlewareContainer(new NestContainer()),
         mockModule as any,
         'Test' as any,
       );
@@ -71,27 +71,35 @@ describe('MiddlewareModule', () => {
   });
 
   describe('registerRouteMiddleware', () => {
+    class TestModule {}
+
+    let nestContainer: NestContainer;
+
+    beforeEach(() => {
+      nestContainer = new NestContainer();
+      nestContainer
+        .getModules()
+        .set('Test', new Module(TestModule, [], nestContainer));
+    });
     it('should throw "RuntimeException" exception when middleware is not stored in container', () => {
       const route = { path: 'Test' };
       const configuration = {
         middleware: [TestMiddleware],
         forRoutes: [BaseController],
       };
-
       const useSpy = sinon.spy();
       const app = { use: useSpy };
 
-      const nestContainer = new NestContainer();
       // tslint:disable-next-line:no-string-literal
       middlewareModule['container'] = nestContainer;
 
       expect(
         middlewareModule.registerRouteMiddleware(
-          new MiddlewareContainer(),
+          new MiddlewareContainer(nestContainer),
           route as any,
           configuration,
-          'Test' as any,
-          app as any,
+          'Test',
+          app,
         ),
       ).to.eventually.be.rejectedWith(RuntimeException);
     });
@@ -109,8 +117,8 @@ describe('MiddlewareModule', () => {
       const useSpy = sinon.spy();
       const app = { use: useSpy };
 
-      const container = new MiddlewareContainer();
-      const moduleKey = 'Test' as any;
+      const container = new MiddlewareContainer(nestContainer);
+      const moduleKey = 'Test';
       container.insertConfig([configuration], moduleKey);
 
       const instance = new InvalidMiddleware();
@@ -125,7 +133,7 @@ describe('MiddlewareModule', () => {
           route as any,
           configuration,
           moduleKey,
-          app as any,
+          app,
         ),
       ).to.be.rejectedWith(InvalidMiddlewareException);
     });
@@ -143,7 +151,7 @@ describe('MiddlewareModule', () => {
       const app = {
         createMiddlewareFactory: createMiddlewareFactoryStub,
       };
-      const container = new MiddlewareContainer();
+      const container = new MiddlewareContainer(new NestContainer());
       const moduleKey = 'Test';
       container.insertConfig([configuration], moduleKey);
 
@@ -155,7 +163,6 @@ describe('MiddlewareModule', () => {
           instance,
         }),
       );
-      const nestContainer = new NestContainer();
       sinon
         .stub(nestContainer, 'getModuleByKey')
         .callsFake(() => new Module(class {}, [], nestContainer));