Sign in Subscribe

Part 1/3: How to Scale a Chat App to Millions of Users in NestJS

Rares Tarabega

• 6 min read
Part 1/3: How to Scale a Chat App to Millions of Users in NestJS

Hello, in this guide, I will teach you how to scale a NestJS app to millions of users.

Prerequisites:

  • Experience with JavaScript.
  • Experience with NestJS.
  • Experience as a backend engineer.

Access the GitHub repository for this project here.

To scale your app, you will need to scale horizontally by creating multiple instances of your NestJS app.

Each instance of your app should have the following features:

  1. WebSocket integrations that keep track of connected clients through websockets in a singleton memory array. (this is done by default when using nestjs, reference: Injection Scopes)
  2. A Redis database publisher that announces new requests received by all other instances and a subscriber that listents to these requests. This is necessary because a client may send a request to one instance (e.g., instance number 3) and be registered as a websocket client there. However, the client may then send a request to another instance because you will have a load balancer installed (e.g., instance number 32), triggering a websocket chat notification. Unfortunately, this instance does not have the connected client in its singleton memory array. To handle this, the instance needs to announce the request to all other instances through Redis. The other instances can then look up their array of connected clients and notify them if a match is found.

Here's how you can achieve this:

  1. Install websockets in your NestJS app. Follow the official NestJS guide. If you'd like to quickly jump out and see the GitHub repo for the full example, you can find it here.

  2. Create the WebSocket client manager service (explained below).

After setting up websockets in your app, we need to define an entrypoint gateway for websockets.

@WebSocketGateway({ path: '/entrypoint' })
export class EntrypointGateway {}

Then, create a new lifecycle gateway that extends this entrypoint. This lifecycle gateway contains the logic to insert clients into the WebSocket client manager.

@Injectable()
export class LifecycleGateway
  extends EntrypointGateway
  implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect
{
  private readonly logger = new Logger(LifecycleGateway.name);

  constructor(
    private readonly jwtService: JwtService,
    private readonly wsClientManager: WsClientManager,
  ) {
    super();
  }

  afterInit() {
    this.logger.debug('Websockets initialized ' + this.constructor.name);
  }

  handleConnection(client: any) {
    const authUserTokenData = this.getDecodedAuthToken(client);

    if (!authUserTokenData) {
      client.close();
      return;
    }

    this.wsClientManager.addConnection(client, authUserTokenData);
  }

  handleDisconnect(client: any) {
    this.wsClientManager.removeConnection(client);
  }

  getDecodedAuthToken(client: any) {
    let decodedJwt: DecodedAuthToken | null = null;

    try {
      if (client.protocol) {
        decodedJwt = this.jwtService.verify(client.protocol, {
          secret: jwtConstants.secret,
        }) as DecodedAuthToken;
      }
    } catch (e) {}

    return decodedJwt;
  }
}

To inject the Redis connection into our app, we utilize @liaoliaots/nestjs-redis package with the following yarn command:

 yarn add @liaoliaots/nestjs-redis

Moving forward, we will add the Redis module configuration for subscriber and publisher channels.

RedisModule.forRootAsync({
  inject: [ConfigService],
  useFactory: async (config: ConfigService) => ({
    config: [
      {
        namespace: 'subscriber',
        host: config.get('REDIS_HOST'),
        port: config.get('REDIS_PORT'),
      },
      {
        namespace: 'publisher',
        host: config.get('REDIS_HOST'),
        port: config.get('REDIS_PORT'),
      },
    ],
  }),
}),

In the WebSocket client manager, you will find the following:

  • We generate a random Redis client using the crypto module's randomUUID from the official Node.js package.

Each NestJS app instance subscribes to the three Redis channels we defined:

  • SendWsMessageToAllClients
  • SendWsMessageToSomeClients
  • SendWsMessageToOneClient

When a NestJS instance receives a message, it first checks that the publisher is not the same as the subscriber. Then, it proceeds to send the message. However, this time, we set the shouldPublishToRedis parameter to false to avoid an infinite loop. The received message will be sent by each instance if it finds the clients in its singleton memory array.

This approach also handles cases where a single user uses the app through multiple devices simultaneously (e.g., PC, mobile app). The user will receive chat notifications on all devices because, upon inspecting the code, you will notice that we are utilizing a map object. For every user ID, we store all associated WebSocket connected clients.

@Injectable()
export class WsClientManager {
  private readonly connectedClients = new Map<string, any[]>();
  private readonly redisClientId = `ws_socket_client-${crypto.randomUUID()}`;

  constructor(
    @InjectRedis('subscriber') private readonly subscriberRedis: Redis,
    @InjectRedis('publisher') private readonly publisherRedis: Redis,
  ) {
    this.subscriberRedis.subscribe(
      RedisSubscribeChannel.SendWsMessageToAllClients,
    );
    this.subscriberRedis.subscribe(
      RedisSubscribeChannel.SendWsMessageToSomeClients,
    );
    this.subscriberRedis.subscribe(
      RedisSubscribeChannel.SendWsMessageToOneClient,
    );

    this.subscriberRedis.on('message', (channel, message) => {
      const data = JSON.parse(message) as RedisPubSubMessage;
      if (data.from !== this.redisClientId) {
        switch (channel) {
          case RedisSubscribeChannel.SendWsMessageToAllClients:
            this.sendMessageToAllClients(data.message, false);
            break;
          case RedisSubscribeChannel.SendWsMessageToSomeClients:
            this.sendMessageToClients(
              (data as RedisPubSubMessageWithClientIds).clientIds,
              data.message,
              false,
            );
            break;
          case RedisSubscribeChannel.SendWsMessageToOneClient:
            this.sendMessageToClient(
              (data as RedisPubSubMessageWithClientId).clientId,
              data.message,
              false,
            );
            break;
        }
      }
    });
  }

  addConnection(client: any, authUserTokenData: DecodedAuthToken) {
    const userId = authUserTokenData.sub;

    this.setUserIdOnClient(client, userId);
    const clientsPool = this.getClientsPool(client);
    this.connectedClients.set(
      userId,
      clientsPool ? [...clientsPool, client] : [client],
    );

    setTimeout(() => {
      client.close(); // This will trigger removeConnection from the LifecycleGateway's handleDisconnect
    }, this.getConnectionLimit(authUserTokenData));
  }

  removeConnection(client: any) {
    const clientsPool = this.getClientsPool(client);
    const newPool = clientsPool!.filter((c) => c !== client);

    if (!newPool.length) {
      this.connectedClients.delete(client.userId);
    } else {
      this.connectedClients.set(client.userId, newPool);
    }
  }

  private setUserIdOnClient(client: any, userId: string) {
    client.userId = userId;
  }

  private getClientsPool(client: any) {
    return this.connectedClients.get(client.userId);
  }

  private getConnectionLimit(tokenData: DecodedAuthToken) {
    return tokenData.exp * 1000 - Date.now();
  }

  getConnectedClientIds() {
    const clientIds: string[] = [];

    const iterator = this.connectedClients.keys();
    let current = iterator.next();
    while (!current.done)

 {
      clientIds.push(current.value);
      current = iterator.next();
    }

    return clientIds;
  }

  sendMessageToClient(
    clientId: string,
    message: string,
    shouldPublishToRedis = true,
  ) {
    if (shouldPublishToRedis) {
      this.publisherRedis.publish(
        RedisSubscribeChannel.SendWsMessageToOneClient,
        JSON.stringify({
          message,
          clientId,
          from: this.redisClientId,
        }),
      );
    }

    const clientPool = this.connectedClients.get(clientId);

    if (clientPool) {
      clientPool.forEach((client) => {
        client.send(message);
      });
    }
  }

  sendMessageToClients(
    clientIds: string[],
    message: string,
    shouldPublishToRedis = true,
  ) {
    if (shouldPublishToRedis) {
      this.publisherRedis.publish(
        RedisSubscribeChannel.SendWsMessageToSomeClients,
        JSON.stringify({
          message,
          clientIds,
          from: this.redisClientId,
        }),
      );
    }

    this.connectedClients.forEach((clientPool, clientId) => {
      if (clientIds.includes(clientId)) {
        clientPool.forEach((client) => {
          client.send(message);
        });
      }
    });
  }

  sendMessageToAllClients(message: string, shouldPublishToRedis = true) {
    if (shouldPublishToRedis) {
      this.publisherRedis.publish(
        RedisSubscribeChannel.SendWsMessageToAllClients,
        JSON.stringify({
          message,
          from: this.redisClientId,
        }),
      );
    }

    this.connectedClients.forEach((clientPool) => {
      clientPool.forEach((client) => {
        client.send(message);
      });
    });
  }
}

Now it's time to test our app. First, clone the project. Make sure you have Docker installed.

This is the Dockerfile for our NestJS app (represents a single instance):

Dockerfile:

FROM node:18.13.0-alpine AS development

WORKDIR /usr/src/app
COPY package.json ./
COPY yarn.lock ./

RUN yarn install --frozen-lockfile \
&& yarn cache clean

COPY . .

RUN yarn build

CMD yarn install; yarn start:debug;

In the docker-compose file, we use our backend image generated from the Dockerfile above. We also have an Nginx server to act as a load balancer between app instances, Postgres as the database, and Redis as the centralized database used for communication between app instances. In this example, we simulate the presence of five NestJS app instances.

docker-compose.yml:

version: '3.7'

services:
  backend:
    image: scalable-chat-app-example-backend
    build:
      context: ./../
      dockerfile: ./docker/Dockerfile
      target: development
    volumes:
      - ./../:/usr/src/app
      - scalable-chat-app-example-backend-node-modules:/usr/src/app/node_modules
      - scalable-chat-app-example-backend-dist:/usr/src/app/dist
    ports:
      - '3000'
    networks:
      - mainnet
    depends_on:
      - postgres
      - redis
    restart: unless-stopped
    scale: 5

  nginx:
    container_name: scalable-chat-app-example-nginx-load-balancer
    image: nginx:latest
    volumes:
      - ./../nginx.conf:/etc/nginx/nginx.conf
    depends_on:
      - backend
    networks:
      - mainnet
    ports:
      - '3000:3000'

  postgres:
    container_name: scalable-chat-app-example-postgres-db
    image: postgres:15.0
    networks:
      - mainnet
    environment:
      TZ: ${DB_TIMEZONE}
      POSTGRES_USER: ${DB_USERNAME}
      POSTGRES_PASSWORD: ${DB_PASSWORD}
      POSTGRES_DB: ${DB_NAME}
      PG_DATA: /var/lib/postgresql/data
    ports:
      - '5432:5432'
    volumes:
      - scalable-chat-app-example-pgdata:/var/lib/postgresql/data

  redis:
    container_name: scalable-chat-app-example-redis-db
    image: redis:7.0.7
    networks:
      - mainnet
    expose:
      - '6379'
    volumes:
      - scalable-chat-app-example-redisdata:/data

networks:
  mainnet:

volumes:
  scalable-chat-app-example-pgdata:
  scalable-chat-app-example-backend-node-modules:
  scalable-chat-app-example-backend-dist:
  scalable-chat-app-example-redisdata:

We also need the following Nginx configuration to handle WebSocket connections:

nginx.conf:

user nginx;
events {

	worker_connections 1000;
}
http {

	upstream app {

		server scalable-chat-app-example-backend-1:3000;
		server scalable-chat-app-example-backend-2:3000;
		server scalable-chat-app-example-backend-3:3000;
		server scalable-chat-app-example-backend-4:3000;
		server scalable-chat-app-example-backend-5:3000;
	}

	server {

		listen 3000;
		location / {

			proxy_pass http://app;
			# WebSocket support
			proxy_http_version 1.1;
			proxy_set_header Upgrade $http_upgrade;
			proxy_set_header Connection $http_connection;

		}
		client_max_body_size 1000M;
	}
}

To set up the example, navigate to the repository directory and execute the following three commands:

Copy sample.env:

cp sample.env .env

Run containers:

yarn dc up

Initiate docker database:

yarn dc-db-init

If you want to clean the project you can run the following command:

yarn dc-clean

Your database has 5 users initiated:

Each one of them has 1234 as password.

Open swagger at http://localhost:3000/docs.

Login multiple users so you obtain the bearer token needed for ws authentication.

Open as many browser tabs as you want and write the following code in each of them:

const ws = new WebSocket('ws://localhost:3000/entrypoint', 'introduceHereBearerTokenFromAuthLogin')

Check out your users table from docker database and extract them as a json so you can use the array of users in the browser.

Now, type the following in any browser window:

 users.forEach(user => ws.send(JSON.stringify({event: 'send_chat_message_to_participant', data: {message: 'test', participantId: user.id}})));

This will send individual chat message to all participants. You can also check docker logs to see it in action.

Regardless of which NestJS instance grabs the WebSocket connection, the messages are consistently sent to participants. This is how you scale a NestJS chat app to millions of users.

If you are interested in learning how to deploy this stack on Kubernetes, leave a comment and i will do the tutorial.

If you want to collaborate on potential start-up projects, you can reach me at: rares.tarabega27@gmail.com

If you want to check out an improved and more efficient model for broadcasting messages between instances, as well as how to deploy this stack on Kubernetes, you can find Part 2 of this series here.