WebSocket cluster with NestJs and Redis

Scaling is an inevitable part of back-end app life, once you decided to scale your application to multiple instances, you are going to face a problem with handling users with multiple clients ( phone, laptop, … ) each connecting to a random instance of your cluster.

In this article, we are going to define the problem and make a solution around it using NestJs and Redis.

Requirements:

• Experience working with Nodejs and NestJs
• Having Nodejs installed
• Having NestJs CLI installed
• Having Redis installed

The Problem:

• Messages emitted on WebSocket need to be sent on every device of the recipient which is connected to one of our instances.

The Solution:

• We will handle messages to multiple instances using Redis PubSub streams, To accomplish this on NestJs we will create a module called socket module and we will put a gateway for handling socket clients and a service for doing discovery, connecting to Redis and distrobuting messages.
  we will cover this with step by step guide.

Install NestJs CLI as a global package

# run with sudo if you are on ubuntu
npm i -g @nestjs/cli

Create new NestJS project, also dependencies will be installed through a wizard by this command

nest g socket-cluster-app

Generate the Socket module we were talking about

# go into project folder
cd socket-cluster-app/
# generate socket module
nest g module socket
# generate socket service
nest g service socket
# generate socket gateway
nest g gateway socket/socket

Using “nest g” command will be automatically adding your services and sockets to their relative modules

Install WebSocket adapter

npm i @nestjs/platform-ws
npm i @nestjs/websockets

Register adapter in “main.ts” file

Then we will identify each socket on handleConnection calls, and we will put a “userId” property to each client. in this example we will set userId by token cookie sent from the client, in a real-world example, you need to validate the token and assign userId to the client by querying your database or some authentication service.

src/socket.gateway.ts

Now we need to implement socket service, we will require a Redis package for distributing messages between instances.

npm i redis
npm i --save-dev @types/redis

Socket service is going to have multiple functions

• constructor, Step zero is to assign a random id to our service in the constructor method and inject the “SocketGateWay” which we have implemented in the last step.

src/main.ts

Also, we are implementing onModuleInit function in socket service which will create and connect to 3 Redis clients.
- redisClient for updating service key by channel discovery
- subscriberClient to get distributed messages
- publisherClient to distribute messages to other instances

src/socket/socket.service.ts

createClient is imported from “redis” package

• channelDiscovery will save its serviceId on Redis with the expiration of 3 seconds. it will also start self-repeating timeout to re-execute every 2 seconds. this way all instances will have access to an updated list of socket services ready for distribution of messages.
  clearing discovery interval timeout would be a good idea to prevent open handler's problem when testing this service.

src/socket/socket.service.ts

• sendMessage final step would be sending messages to every connected client of a specific user.
  We are sending the message to our connected clients and also distributing this message to other instances.
  “if(!fromRedisChannel)” will prevent distributing if the message is already distributed by another instance.

src/socket/socket.service.ts

Okay, we are done, now we can set up our test scenario.
First, we are going to create a simple test script that will connect to one of our instances and print the received messages.

install ws package by running “npm i ws”

const ws = require('ws');
const port = 3001;
const socket = new ws(`ws://localhost:${port}`, {
   headers: { Cookie: 'token=user1' },
});
socket.on('message', data => {
  console.log(`Received message`, data);
});
socket.on('open', data => {
  console.log(`Connected to port ${port}`);
});
socket.on('close', data => {
  console.log(`Disconnected from port ${port}`);
});

Then we add a simple interval to our socket service for sending time to user1.

Finally, run the following commands in order

PORT=3001 npm start
PORT=3002 npm start
node test-script.js

test script should log a message from both instances every 3 seconds.

# output
Received message 8:21:55 AM | from server on port 3001
Received message 8:21:57 AM | from server on port 3002

This shows us that now our service is capable to distribute WebSocket messages from different instances to a specific client.

A fully working example of what we stepped trough in this article is available at https://github.com/m-esm/socket-cluster-app