Repository: novuhq/novu
File: apps/ws/src/socket/ws.gateway.ts (TypeScript)

Code before:

  private async processConnectionRequest(connection: Socket) {
    const token = this.extractToken(connection);

    if (!token || token === 'null') {
      Logger.warn(`No token was found during counnection process for ${connection.id}`, LOG_CONTEXT);

      return this.disconnect(connection);
    }

    const subscriber = await this.getSubscriber(token);
    if (!subscriber) {
      Logger.warn(`No subscriber was found for specified token ${connection.id}`, LOG_CONTEXT);

      return this.disconnect(connection);
    }

    Logger.log(
      `Connection request received from ${subscriber._id} external id: ${subscriber.subscriberId} organization id: ${subscriber.organizationId}`,
      LOG_CONTEXT
    );

    await connection.join(subscriber._id);

    Logger.log(`Connection request accepted for ${subscriber._id}`, LOG_CONTEXT);

    await this.subscriberOnlineService.handleConnection(subscriber);
  }

  async sendMessage(userId: string, event: string, data: any) {
    if (!this.server) {
      Logger.error('No sw server available to send message', LOG_CONTEXT);

      return;
    }

    Logger.log(`Sending event ${event} message to ${userId}`, LOG_CONTEXT);

    this.server.to(userId).emit(event, data);
  }

  private disconnect(socket: Socket) {
    socket.disconnect();
  }

  async gracefulShutdown(): Promise<void> {
    try {
      if (!this.server) {
        Logger.error('WS server was not initialized while executing shutdown', LOG_CONTEXT);

        return;

Review comment (refactor):
I also left the original subscriberId room for backwards compatibility, so after this change is rolled-out, the number of open rooms will double (until we remove this line), but that shouldn't be a problem, its just a server side grouping mechanism not a new WS connection. 

Code after:
  }

  private async processConnectionRequest(connection: Socket) {
    const token = this.extractToken(connection);

    if (!token || token === 'null') {
      Logger.warn(`No token was found during counnection process for ${connection.id}`, LOG_CONTEXT);

      return this.disconnect(connection);
    }

    const subscriber = await this.getSubscriber(token);
    if (!subscriber) {
      Logger.warn(`No subscriber was found for specified token ${connection.id}`, LOG_CONTEXT);

      return this.disconnect(connection);
    }

    const contextKeys = subscriber.contextKeys || [];

    Logger.log(
      `Connection request received from ${subscriber._id} external id: ${subscriber.subscriberId} organization id: ${subscriber.organizationId}`,
      LOG_CONTEXT
    );

    await connection.join(subscriber._id);

    // Join context-specific rooms for efficient message routing
    if (contextKeys.length === 0) {
      // Join default room for connections without context
      await connection.join(`${subscriber._id}:${DEFAULT_CONTEXT_ROOM}`);
      Logger.log(`Connection ${connection.id} joined default room for ${subscriber._id}`, LOG_CONTEXT);
    } else {
      // Join a room for each context key
      for (const contextKey of contextKeys) {
        await connection.join(`${subscriber._id}:${contextKey}`);
      }
      Logger.log(
        `Connection ${connection.id} joined context rooms for ${subscriber._id}: ${contextKeys.join(', ')}`,
        LOG_CONTEXT
      );
    }

    Logger.log(`Connection request accepted for ${subscriber._id}`, LOG_CONTEXT);

    await this.subscriberOnlineService.handleConnection(subscriber);
  }

  async sendMessage(userId: string, event: string, data: any, contextKeys?: string[]) {
    if (!this.server) {
      Logger.error('No sw server available to send message', LOG_CONTEXT);