fix(websockets): Prevent HTTP server early close in Socket.IO shutdown

nestjs · Nov 21, 2024 · 40dbb27 · 40dbb27
1 parent 440e9cf
commit 40dbb27
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/packages/websockets/adapters/ws-adapter.ts b/packages/websockets/adapters/ws-adapter.ts
@@ -17,6 +17,7 @@ export abstract class AbstractWsAdapter<
 > implements WebSocketAdapter<TServer, TClient, TOptions>
 {
   protected readonly httpServer: any;
+  private _forceCloseConnections: boolean;
 
   constructor(appOrHttpServer?: INestApplicationContext | any) {
     if (appOrHttpServer && appOrHttpServer instanceof NestApplication) {
@@ -26,6 +27,10 @@ export abstract class AbstractWsAdapter<
     }
   }
 
+  public set forceCloseConnections(value: boolean) {
+    this._forceCloseConnections = value;
+  }
+
   public bindClientConnect(server: TServer, callback: Function) {
     server.on(CONNECTION_EVENT, callback);
   }
@@ -35,6 +40,9 @@ export abstract class AbstractWsAdapter<
   }
 
   public async close(server: TServer) {
+    if (this._forceCloseConnections) {
+      return;
+    }
     const isCallable = server && isFunction(server.close);
     isCallable && (await new Promise(resolve => server.close(resolve)));
   }

diff --git a/packages/websockets/socket-module.ts b/packages/websockets/socket-module.ts
@@ -1,6 +1,6 @@
+import { NestApplicationOptions } from '@nestjs/common';
 import { InjectionToken } from '@nestjs/common/interfaces';
 import { Injectable } from '@nestjs/common/interfaces/injectable.interface';
-import { NestApplicationContextOptions } from '@nestjs/common/interfaces/nest-application-context-options.interface';
 import { ApplicationConfig } from '@nestjs/core/application-config';
 import { GuardsConsumer } from '@nestjs/core/guards/guards-consumer';
 import { GuardsContextCreator } from '@nestjs/core/guards/guards-context-creator';
@@ -25,8 +25,7 @@ import { WebSocketsController } from './web-sockets-controller';
 
 export class SocketModule<
   THttpServer = any,
-  TAppOptions extends
-    NestApplicationContextOptions = NestApplicationContextOptions,
+  TAppOptions extends NestApplicationOptions = NestApplicationOptions,
 > {
   private readonly socketsContainer = new SocketsContainer();
   private applicationConfig: ApplicationConfig;
@@ -113,8 +112,11 @@ export class SocketModule<
   }
 
   private initializeAdapter() {
+    const forceCloseConnections = this.appOptions.forceCloseConnections;
     const adapter = this.applicationConfig.getIoAdapter();
     if (adapter) {
+      (adapter as AbstractWsAdapter).forceCloseConnections =
+        forceCloseConnections;
       this.isAdapterInitialized = true;
       return;
     }
@@ -124,6 +126,7 @@ export class SocketModule<
       () => require('@nestjs/platform-socket.io'),
     );
     const ioAdapter = new IoAdapter(this.httpServer);
+    ioAdapter.forceCloseConnections = forceCloseConnections;
     this.applicationConfig.setIoAdapter(ioAdapter);
 
     this.isAdapterInitialized = true;