Scalable Websockets with AWS API Gateway and AWS Lambda
Tags: webdev, react, node, aws
Published: 2021-04-30 15:52:22

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](https://dev.to/fastmodularproject/web-starter-react-nestjs-mysql-3o9i).

# 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](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda), 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**](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda).

# 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)](https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-overview-developer-experience.html#api-gateway-overview-websocket) 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',
    },
```

3) Set your API Gateway endpoint in the environments variables of your frontend:
```
socketUrl: "wss://XXXXXXXX.execute-api.eu-west-3.amazonaws.com/env"
```

# 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](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda) 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](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda) 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
        });
        ...
```

You should of course implement:
`ws.onopen` method: to define [what to do when a connection is created](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda).

`ws.onmessage` method: to define [what to do when receiving a new message](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda).

`ws.onclose` method: to define [what to do when a connection is closed](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda).


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[];
    ...
}
```
```
@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;
}
```

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);
    }
}
```

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);
                    }
                }
            });
        }
    };
```

# 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](https://fast-modular-project.com)
* Module "Websocket with API Gateway and AWS Lambda" repository [here](https://fast-modular-project.com/modules/websocket-apigateway-aws-lambda).

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