Merge pull request #4918 from novuhq/nv-3205-do-not-perform-fetchsock…

…ets-when-service-is-shutting-down

refactor: do not perform fetchsockets when service is shutting down

apps/api/src/bootstrap.ts

@@ -126,7 +126,6 @@ export async function bootstrap(expressApp?): Promise<INestApplication> {
     await app.listen(process.env.PORT);
   }
 
-  // Starts listening for shutdown hooks
   app.enableShutdownHooks();
 
   Logger.log(`Started application in NODE_ENV=${process.env.NODE_ENV} on port ${process.env.PORT}`);

apps/worker/src/app/workflow/services/execution-log.worker.ts

@@ -2,7 +2,6 @@ import { Inject, Injectable, Logger } from '@nestjs/common';
 const nr = require('newrelic');
 import {
   getExecutionLogWorkerOptions,
-  INovuWorker,
   PinoLogger,
   storage,
   Store,
@@ -18,13 +17,12 @@ import { ObservabilityBackgroundTransactionEnum } from '@novu/shared';
 const LOG_CONTEXT = 'ExecutionLogWorker';
 
 @Injectable()
-export class ExecutionLogWorker extends ExecutionLogWorkerService implements INovuWorker {
+export class ExecutionLogWorker extends ExecutionLogWorkerService {
   constructor(
     private createExecutionDetails: CreateExecutionDetails,
     public workflowInMemoryProviderService: WorkflowInMemoryProviderService
   ) {
     super(new BullMqService(workflowInMemoryProviderService));
-
     this.initWorker(this.getWorkerProcessor(), this.getWorkerOptions());
   }
 

apps/worker/src/app/workflow/services/standard.worker.ts

@@ -4,7 +4,6 @@ import { IJobData, ObservabilityBackgroundTransactionEnum } from '@novu/shared';
 import {
   BullMqService,
   getStandardWorkerOptions,
-  INovuWorker,
   Job,
   PinoLogger,
   StandardWorkerService,
@@ -29,7 +28,7 @@ import {
 const LOG_CONTEXT = 'StandardWorker';
 
 @Injectable()
-export class StandardWorker extends StandardWorkerService implements INovuWorker {
+export class StandardWorker extends StandardWorkerService {
   constructor(
     private handleLastFailedJob: HandleLastFailedJob,
     private runJob: RunJob,

apps/worker/src/app/workflow/services/workflow.worker.ts

@@ -1,8 +1,7 @@
-import { Inject, Injectable, Logger } from '@nestjs/common';
+import { Injectable, Logger } from '@nestjs/common';
 const nr = require('newrelic');
 import {
   getWorkflowWorkerOptions,
-  INovuWorker,
   PinoLogger,
   storage,
   Store,
@@ -19,7 +18,7 @@ import { ObservabilityBackgroundTransactionEnum } from '@novu/shared';
 const LOG_CONTEXT = 'WorkflowWorker';
 
 @Injectable()
-export class WorkflowWorker extends WorkflowWorkerService implements INovuWorker {
+export class WorkflowWorker extends WorkflowWorkerService {
   constructor(
     private triggerEventUsecase: TriggerEvent,
     public workflowInMemoryProviderService: WorkflowInMemoryProviderService

apps/ws/src/bootstrap.ts

@@ -45,6 +45,8 @@ export async function bootstrap() {
 
   app.useWebSocketAdapter(redisIoAdapter);
 
+  app.enableShutdownHooks();
+
   await app.init();
 
   try {

apps/ws/src/health/health.module.ts

@@ -4,12 +4,12 @@ import { TerminusModule } from '@nestjs/terminus';
 import { HealthController } from './health.controller';
 import { SharedModule } from '../shared/shared.module';
 import { WSServerHealthIndicator } from '../socket/services';
-import { WSGateway } from '../socket/ws.gateway';
+import { SocketModule } from '../socket/socket.module';
 
-const PROVIDERS = [WSServerHealthIndicator, WSGateway];
+const PROVIDERS = [WSServerHealthIndicator];
 
 @Module({
-  imports: [TerminusModule, SharedModule],
+  imports: [TerminusModule, SharedModule, SocketModule],
   providers: PROVIDERS,
   controllers: [HealthController],
 })

apps/ws/src/shared/subscriber-online/subscriber-online.service.ts

@@ -20,6 +20,7 @@ export class SubscriberOnlineService {
   async handleDisconnection(subscriber: ISubscriberJwt, activeConnections: number) {
     const lastOnlineAt = new Date().toISOString();
     let isOnline = false;
+
     if (activeConnections > 1) {
       isOnline = true;
     }

apps/ws/src/socket/services/web-socket.worker.ts

@@ -4,7 +4,6 @@ import { Injectable, Logger } from '@nestjs/common';
 import {
   BullMqService,
   getWebSocketWorkerOptions,
-  INovuWorker,
   WebSocketsWorkerService,
   WorkerOptions,
   WorkflowInMemoryProviderService,
@@ -16,7 +15,7 @@ import { ObservabilityBackgroundTransactionEnum } from '@novu/shared';
 const LOG_CONTEXT = 'WebSocketWorker';
 
 @Injectable()
-export class WebSocketWorker extends WebSocketsWorkerService implements INovuWorker {
+export class WebSocketWorker extends WebSocketsWorkerService {
   constructor(
     private externalServicesRoute: ExternalServicesRoute,
     private workflowInMemoryProviderService: WorkflowInMemoryProviderService

apps/ws/src/socket/usecases/external-services-route/external-services-route.usecase.ts

@@ -60,12 +60,6 @@ export class ExternalServicesRoute {
       return;
     }
 
-    const connection = await this.connectionExist(command);
-
-    if (!connection) {
-      return;
-    }
-
     let unreadCount = this.extractCount(command.payload?.unreadCount);
 
     if (unreadCount === undefined) {
@@ -91,12 +85,6 @@ export class ExternalServicesRoute {
       return;
     }
 
-    const connection = await this.connectionExist(command);
-
-    if (!connection) {
-      return;
-    }
-
     let unseenCount = this.extractCount(command.payload?.unseenCount);
 
     if (unseenCount === undefined) {

apps/ws/src/socket/ws.gateway.ts

@@ -1,17 +1,20 @@
 const nr = require('newrelic');
+import { Server, Socket } from 'socket.io';
 import { JwtService } from '@nestjs/jwt';
 import { OnGatewayConnection, OnGatewayDisconnect, WebSocketGateway, WebSocketServer } from '@nestjs/websockets';
 import { Logger } from '@nestjs/common';
-import { Server, Socket } from 'socket.io';
 
 import { ISubscriberJwt, ObservabilityBackgroundTransactionEnum } from '@novu/shared';
+import { IDestroy } from '@novu/application-generic';
 
 import { SubscriberOnlineService } from '../shared/subscriber-online';
 
 const LOG_CONTEXT = 'WSGateway';
 
 @WebSocketGateway()
-export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect {
+export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect, IDestroy {
+  private isShutdown = false;
+
   constructor(private jwtService: JwtService, private subscriberOnlineService: SubscriberOnlineService) {}
 
   @WebSocketServer()
@@ -82,7 +85,18 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect {
     }
   }
 
+  /*
+   * This method is called when a client disconnects from the server.
+   * * When a shutdown is in progress, we opt out of updating the subscriber status,
+   * assuming that when the current instance goes down, another instance will take its place and handle the subscriber status update.
+   */
   private async processDisconnectionRequest(connection: Socket) {
+    if (!this.isShutdown) {
+      await this.handlerSubscriberDisconnection(connection);
+    }
+  }
+
+  private async handlerSubscriberDisconnection(connection: Socket) {
     const token = this.extractToken(connection);
 
     if (!token || token === 'null') {
@@ -98,6 +112,12 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect {
     await this.subscriberOnlineService.handleDisconnection(subscriber, activeConnections);
   }
 
+  private async getActiveConnections(socket: Socket, subscriberId: string) {
+    const activeSockets = await socket.in(subscriberId).fetchSockets();
+
+    return activeSockets.length;
+  }
+
   private async processConnectionRequest(connection: Socket) {
     const token = this.extractToken(connection);
 
@@ -129,9 +149,29 @@ export class WSGateway implements OnGatewayConnection, OnGatewayDisconnect {
     socket.disconnect();
   }
 
-  private async getActiveConnections(socket: Socket, subscriberId: string) {
-    const activeSockets = await socket.in(subscriberId).fetchSockets();
+  async gracefulShutdown(): Promise<void> {
+    try {
+      if (!this.server) {
+        Logger.error('WS server was not initialized while executing shutdown', LOG_CONTEXT);
 
-    return activeSockets.length;
+        return;
+      }
+
+      Logger.log('Closing WS server for incoming new connections', LOG_CONTEXT);
+      this.server.close();
+
+      Logger.log('Disconnecting active sockets connections', LOG_CONTEXT);
+      this.server.sockets.disconnectSockets();
+    } catch (e) {
+      Logger.error(e, 'Unexpected exception was thrown while graceful shut down was executed', LOG_CONTEXT);
+      throw e;
+    } finally {
+      Logger.log(`Graceful shutdown down has finished`, LOG_CONTEXT);
+    }
+  }
+
+  async onModuleDestroy(): Promise<void> {
+    this.isShutdown = true;
+    await this.gracefulShutdown();
   }
 }

packages/application-generic/src/modules/index.ts

@@ -1,2 +1,3 @@
 export { QueuesModule } from './queues.module';
 export { MetricsModule } from './metrics.module';
+export * from './interfaces';

packages/application-generic/src/modules/interfaces.ts

@@ -0,0 +1,4 @@
+export interface IDestroy {
+  gracefulShutdown?: () => Promise<void>;
+  onModuleDestroy: () => Promise<void>;
+}

packages/application-generic/src/services/readiness/readiness.service.ts

@@ -5,12 +5,11 @@ import { setTimeout } from 'timers/promises';
 import { Worker } from '../bull-mq';
 
 import { IHealthIndicator } from '../../health';
+import { IDestroy } from '../../modules';
 
-export interface INovuWorker {
+export interface INovuWorker extends IDestroy {
   readonly DEFAULT_ATTEMPTS: number;
-  gracefulShutdown: () => Promise<void>;
   readonly topic: string;
-  onModuleDestroy: () => Promise<void>;
   pause: () => Promise<void>;
   resume: () => Promise<void>;
   worker: Worker;

packages/application-generic/src/services/workers/subscriber-process-worker.service.ts

@@ -1,15 +1,11 @@
 import { JobTopicNameEnum } from '@novu/shared';
 
-import { INovuWorker } from '../readiness';
 import { WorkerBaseService } from './worker-base.service';
 import { BullMqService } from '../bull-mq';
 
 const LOG_CONTEXT = 'SubscriberProcessWorkerService';
 
-export class SubscriberProcessWorkerService
-  extends WorkerBaseService
-  implements INovuWorker
-{
+export class SubscriberProcessWorkerService extends WorkerBaseService {
   constructor(private bullMqService: BullMqService) {
     super(JobTopicNameEnum.PROCESS_SUBSCRIBER, bullMqService);
   }

packages/application-generic/src/services/workers/worker-base.service.ts

@@ -1,14 +1,8 @@
-import { IJobData, JobTopicNameEnum } from '@novu/shared';
-import { Inject, Injectable, Logger } from '@nestjs/common';
-
-import {
-  BullMqService,
-  JobsOptions,
-  Processor,
-  Worker,
-  WorkerOptions,
-} from '../bull-mq';
-import { WorkflowInMemoryProviderService } from '../in-memory-provider';
+import { JobTopicNameEnum } from '@novu/shared';
+import { Logger } from '@nestjs/common';
+
+import { BullMqService, Processor, Worker, WorkerOptions } from '../bull-mq';
+import { INovuWorker } from '../readiness';
 
 const LOG_CONTEXT = 'WorkerService';
 
@@ -19,7 +13,7 @@ export type WorkerProcessor =
 
 export { WorkerOptions };
 
-export class WorkerBaseService {
+export class WorkerBaseService implements INovuWorker {
   private instance: BullMqService;
 
   public readonly DEFAULT_ATTEMPTS = 3;