DEV Community

Cover image for Scalable Websockets with AWS API Gateway and AWS Lambda
Leonidas Costas
Leonidas Costas

Posted on

Scalable Websockets with AWS API Gateway and AWS Lambda

Hi Sparta!

In this article I will share with you how to add scalable websockets system in your app with AWS API Gateway and AWS Lambda. Websockets are used to implement any real time system in like a chat or a notification system.

Please note that AWS is not mandatory to implement simple websockets, but it gives us the scalability we are looking for if we are dealing with an app with thousands of users.

On my side, I used this module on top of the React/Node/MySQL starter. This starter has already been presented to you in this article.

What does it bring?

  • Open source code
  • Documentation and "Quick Start"
  • Complete integration of websockets in your React / NestJS / MySQL starter (it can be easily adapted to a node backend)
  • AWS Lambda functions source code
  • 20 hours of work saved :D

Prerequisite

By getting the code here, you'll have the websockets already integrated in the web starter. You'll get a functional project with an authentication and a websocket system in less than 20 minutes :D

Note that using the starter is not mandatory, you can also use the module as standalone. The integration won't be as easy as with the starter, but it should still be simple to integrate it in your already created project :)

Websockets in 3 words

With a standard (REST/SOAP) API, frontend sends information to the server and get a proper answer. This is enough most of the time but it means that the frontend/user need to perform an action to get up to date datas.

Let's imagine a chat where users would have to press a "refresh button" to get new messages displayed... this would be quite annoying. Websockets come to the rescue !

Websocket is a bidirectional connection that you initiate between a server and a client. This connection allows the frontend to speak to the server and vice-versa without any call to an API. If we take back the example of the chat, websockets allow the server to say to the user that he has a new message (without any action from him).

The entire open source code and a step by step integration on the starter is available here.

High level picture of the workflow

The websocket connection will be setup between the client (browser or mobile app) and API Gateway (for sockets). We could have established a websocket connection directly with the backend but this could lead to a shutdown of your API if your server can't scale and if there is too many connections to be maintained. Thanks to API Gateway, the sockets will be handle in a separated server that can scale, independently from your back server.

1) User logs in. He initializes a websocket connection with the API Gateway. The API generates a unique identifier of this connection: connectionId. Once the connection is established, frontend send to the API Gateway (with the socket created) a "connected event". The AWS Lambda that handle the "websocket connection flow" calls our backend endpoint to link this user with this connectionId in our database. The backend saves this infos. For every users connected on our app, we now have one or several connectionId associated.

2) The backend want to send an event to users. For all users, it get their connectionId and ask the API Gateway to send a message to the websocket identified by this connectionId.

3) Users receive the event (without any call to the backend) and adapt their frontend accordingly.

4) A user logs out. This closes the websocket, and notify the API Gateway with a "disconnected event". The AWS Lambda that handle the "websocket disconnection flow" calls our backend endpoint to delete the link between this user and the connectionId in our database. The backend saves this infos.

Configure your AWS API Gateway

1) Create an API Gateway (for sockets) with all default parameters
Add 2 routes:

  • auth: will be called from the frontend when we receive a new connection
  • $disconnect: will be called from the frontend when we receive connection closing event

2) Set your API Gateway credentials in the environments variables of your backend:

    ...
    apiGateway: {
      endpoint:
        'https://xxxxxxx.execute-api.eu-west-3.amazonaws.com/env',
      accessKeyId: 'XXXXXXXXXXXXXXXXXX',
      secretAccessKey: 'XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX',
      region: 'eu-west-3',
    },
Enter fullscreen mode Exit fullscreen mode

3) Set your API Gateway endpoint in the environments variables of your frontend:

socketUrl: "wss://XXXXXXXX.execute-api.eu-west-3.amazonaws.com/env"
Enter fullscreen mode Exit fullscreen mode

Configure your AWS Lambdas

1) Setup the "Connection flow" with AWS Lambda.
Create an AWS Lambda anted websocket-connection and plug it to the auth route of the API Gateway.

In the code of this lambda you should call a backend endpoint we'll create soon. This endpoint will be in charge of saving in database the connectionId of the websocket the user just connected too. Please check here to copy paste the code for this lambda :)

2) Setup the "Disconnection" flow with AWS Lambda.
Create an AWS Lambda anted websocket-disconnection and plug it to the $disconnect route of the API Gateway.

In the code of this lambda you should call a backend endpoint we'll create soon. This endpoint will be in charge of deleting the association between a user and a connectionId in our database. Please check here to copy paste the code for this lambda :)

Setup the websockets in your React frontend

1) Install following package

npm i reconnecting-websocket@4.4.0

2) Init your websocket connection with the API Gateway Socket when the user is connected.

You should establish the connection with the API Gateway thanks to the endpoint stored in your environment variable previously:

        let ws = new ReconnectingWebSocket(
            environment.socketUrl, [], {
            minReconnectionDelay: 500,
            maxReconnectionDelay: 500,
            reconnectionDelayGrowFactor: 1
        });
        ...
Enter fullscreen mode Exit fullscreen mode

You should of course implement:
ws.onopen method: to define what to do when a connection is created.

ws.onmessage method: to define what to do when receiving a new message.

ws.onclose method: to define what to do when a connection is closed.

3) Close the connection when he is loging out: ws.close();

Prepare our NestJS backend for websockets

1) Create a new NestJS module SocketConnectionsModule to manage websockets connections. Do not forget to add it to the import section of your app.module.

Our SocketConnectionEntity will associate a connectionId to a User. A user can have several websockets connections as he may be connected to your app through several browsers or with a mobile application.

@Entity('users')
export class UserEntity implements User {
    ...
    @OneToMany(type => SocketConnectionEntity, socketConnection => socketConnection.user)
    socketConnections: SocketConnectionEntity[];
    ...
}
Enter fullscreen mode Exit fullscreen mode
@Entity('socket_connection')
export class SocketConnectionEntity implements SocketConnection {
    @PrimaryGeneratedColumn()
    id: number;

    @Column({ name: "connection_id" })
    connectionId: string;

    @ManyToOne(() => UserEntity, user => user.socketConnections, { onDelete: 'CASCADE' })
    @JoinColumn({ name: "user_id" })
    user: User;

    @Column({ type: "datetime", default: () => "CURRENT_TIMESTAMP" })
    timestamp: Date;
}
Enter fullscreen mode Exit fullscreen mode

The controller and the service will let us create, get or delete user's connection in database (the SocketConnectionEntity we just created). Those two endpoints are used by the AWS Lambdas we created previously.

@Controller('socket-connections')
export class SocketConnectionsController {

    constructor(private socketConnectionService: SocketConnectionsService) { }

    @Post()
    @Roles('user', 'premium', 'admin')
    async create(@Body() body, @AuthUser() user) {
        return await this.socketConnectionService.create(user, body.connectionId);
    }

    @Delete(':connectionId')
    async delete(@Param() param) {
        return await this.socketConnectionService.deleteConnection(param.connectionId);
    }
}
Enter fullscreen mode Exit fullscreen mode

2) Create a SocketService to ask the API Gateway to send a message to a specific connectionId. Do not forget to import the was sdk import { ApiGatewayManagementApi } from 'aws-sdk'; and create your awsGW object with your API Gateway credentials stored previously in your environment variables.

    async sendMessage(userId, data) {
        const connections = await this.socketConnectionService.getUserConnections(userId);

        for (const connection of connections) {
            console.log("Socket post to : ", connection.connectionId);
            this.awsGW.postToConnection({
                ConnectionId: connection.connectionId,
                Data: JSON.stringify(data)
            }, async (err, success) => {
                if (err) {
                    if (!err.retryable) {
                        // Socket id is disabled
                        await this.socketConnectionService.deleteConnection(connection.connectionId);
                    }
                }
            });
        }
    };
Enter fullscreen mode Exit fullscreen mode

Conclusion

I hope this module will help you saving some time while trying to implement websockets in your project. If you have any question, I'll be present as usual in the comment section !

Links:

  • The platform sharing the starter and it's modules : Fast Modular Project
  • Module "Websocket with API Gateway and AWS Lambda" repository here.

Do not hesitate to pin and like if you appreciated the article ❤️

Top comments (1)

Collapse
 
idevelop profile image
Andrei Gheorghe

All the links out to the code on fast-modular-project.com/ are broken.