This post will teach you how to create a WebSocket server to attend to millions of simultaneous connections. We will build the WebSocket server with a serverless architecture using the AWS API Gateway. This will require the following services from AWS: API Gateway, Lambda, and DynamoDB.
NEW RESEARCH: LEARN HOW DECISION-MAKERS ARE PRIORITIZING DIGITAL INITIATIVES IN 2024.
To simplify building serverless apps, we will use Serverless Framework, a multi-provider (AWS, Microsoft Azure, Google Cloud Platform, Apache OpenWhisk, Cloudflare Workers, or a Kubernetes-based solution like Kubeless) and a mature framework for serverless architecture.
WebSocket
WebSocket resolves a communication need between the server and the client, where clients need to make successive requests to the server or keep a connection open until it responds.
Before WebSockets, several approaches were used for this purpose, such as long polling and short polling. But these strategies relied on intensive resource usage.
WebSocket was created to allow bidirectional communication between client and server when a TCP connection is open.
Today, WebSocket is present in numerous applications, such as chat apps, games, and other real-time applications.
Preparing for the Project
As we’ll use AWS and Serverless Framework, we’ll need to go through a few steps before kicking off the project.
Let’s start by creating credentials for programmatic access to AWS through IAM, where we will specify which permissions our users should have. After getting the credentials, we must configure the Serverless Framework to use them when interacting with AWS. You can learn more about configuring credentials here.
Next, install the Serverless Framework. You can install it directly through the standalone binaries or, if you have npm installed locally, through “npm -g install serverless”.
With everything configured, it’s time to kick off our project.
Starting the project
To start an SLS project, type “sls” or “serverless”, and the prompt command will guide you to creating a new serverless project.
Then, install the “serverless-iam-roles-per-function” plugin, which will help us define the AWS Role for each AWS Lambda. To install, you just need to type “sls plugin install –name serverless-iam-roles-per-function”.
After that, our workspace will have the following structure:
Now our project is ready to start defining the resources for AWS and implementing our lambda functions, which will handle the client’s connection, communication, and disconnection.
Serverless.yml File
To use AWS Services, we can create the services through the AWS Console, AWS CLI, or through a framework that helps us build serverless apps, such as the Serverless Framework.
The serverless.yml contains all the project configuration, such as the info about runtime, environment variables, our functions and the definitions of how we want to package them, plugins we are using, and the required resources (it’s using the CloudFormation template).
service: aws-ws-api-gateway
frameworkVersion: '2 || 3'
provider:
name: aws
runtime: nodejs12.x
lambdaHashingVersion: 20201221
stage: ${opt:stage, "dev"}
environment:
${file(./config/${opt:stage, 'dev'}.yaml)}
functions:
websocketConnect: ${file('./deploy/functions/WebSocketConnect.yaml')}
websocketDisconnect: ${file('./deploy/functions/WebSocketDisconnect.yaml')}
websocketOnMessage: ${file('./deploy/functions/WebSocketOnMessage.yaml')}
package:
individually: true
patterns:
- '!deploy/**'
- '!config/**'
- '!src/**'
plugins:
- serverless-iam-roles-per-function
resources:
Resources:
ConnectionsWebsocketTable: ${file(./deploy/resources/ConnectionsWsDynamoDb.yaml)}
serverless.yml
The Serverless Framework has a complete system of variables that can be used to substitute the values inside the “serverless.yml”. Here we use the ${file} expression that accepts a file path and reads the file’s content, and ${opt}, which allows us to specify a default value when a value doesn’t pass through the CLI operations.
When working with Serverless Framework in a single repository, it’s essential to package functions in a way that reduces the package size and, consequently, the start time of the lambda. To help with this, the Serverless Framework has the package option, where we can specify different output behaviors such as packaging the files individually, removing files from the final artifact, or adding extra files.
Another major advantage of using the Serverless Framework is the vast number of plugins created by the community to facilitate common tasks, such as the serverless-iam-roles-per-function, which helps define specific AWS Roles for each function.
Below is the complete structure of the project and the files containing the DynamoDb table, the environment variables, and the definitions of the functions.
Let’s start with the files in the config folder, which will have the environment variables referenced in provider.environment, according to the deployment stage.
ENV: "dev" CONNECTIONS_WEBSOCKET_TABLE: "ConnectionsWebSocketdev"
dev.yaml
ENV: "test" CONNECTIONS_WEBSOCKET_TABLE: "ConnectionsWebSockettest"
test.yaml
To organize the functions and resources definitions, separate them into “functions” and “resources” folders.
Under the “functions” folder, we’ll have the following files.
handler: src/onconnect.handle
events:
- websocket:
route: $connect
package:
patterns:
- ./src/onconnect.js
iamRoleStatements:
- Effect: Allow
Action:
- dynamodb:PutItem
Resource:
- Fn::GetAtt: [ConnectionsWebsocketTable, Arn]
WebSocketConnect.yaml
handler: src/ondisconnect.handle
events:
- websocket:
route: $disconnect
package:
patterns:
- ./src/ondisconnect.js
iamRoleStatements:
- Effect: Allow
Action:
- dynamodb:DeleteItem
Resource:
- Fn::GetAtt: [ConnectionsWebsocketTable, Arn]
WebSocketDisconnect.yaml
handler: src/onmessage.handle
events:
- websocket:
route: onMessage
package:
patterns:
- ./src/onmessage.js
iamRoleStatements:
- Effect: Allow
Action:
- dynamodb:Scan
Resource:
- Fn::GetAtt: [ConnectionsWebsocketTable, Arn]
- Effect: Allow
Action:
- execute-api:Invoke
- execute-api:ManageConnections
Resource:
Fn::Sub:
- "arn:aws:execute-api:${Region}:${AccountId}:${WebSocketId}/*"
- { Region: !Ref AWS::Region, AccountId: !Ref AWS::AccountId, WebSocketId: !Ref WebsocketsApi }
WebSocketOnMessage.yaml
Lastly, the resources folder in our project will contain only one file to define the AWS DynamoDb table, which will contain the information about connected clients into our WebSocket server.
Type: AWS::DynamoDB::Table
Properties:
AttributeDefinitions:
- AttributeName: connectionId
AttributeType: S
KeySchema:
- AttributeName: connectionId
KeyType: HASH
BillingMode: PAY_PER_REQUEST
TableName: ${self:provider.environment.CONNECTIONS_WEBSOCKET_TABLE}
ConnectionsWsDynamoDb.yaml
AWS API Gateway WebSocket APIs
A WebSocket API in API Gateway is a collection of WebSocket routes integrated with backend HTTP endpoints, Lambda functions, or other AWS services. Incoming JSON messages are directed to backend integrations based on your configured routes.
You can use some predefined routes: $connect, $disconnect, and $default. You can also create custom routes.
To reach our goal, we are going to use $connect and $disconnect and create a custom route called onMessage.
For JSON messages, the routing is done using a route key, which is a defined property when creating the Websocket API. The default route key will be action. Therefore, the messages sent to the Websocket server must contain an action property. If it’s not sent, the $default route will be selected by the route selection expression.
Non-JSON messages are directed to a $default route if it exists.
Coding Our Functions
The $connect route will be responsible for storing the connection ID into a DynamoDb table, which will later be used to deliver the messages to the connected clients.
const AWS = require('aws-sdk');
class OnConnect {
constructor({ repository }) {
this.repository = repository
}
async handle(event) {
const putParams = {
TableName: process.env.CONNECTIONS_WEBSOCKET_TABLE,
Item: {
connectionId: event.requestContext.connectionId
}
};
try {
await this.repository.put(putParams).promise();
} catch (err) {
return {
statusCode: 500,
body: 'Failed to connect: ' + JSON.stringify(err)
};
}
return {
statusCode: 200,
body: JSON.stringify({ connectionId: event.requestContext.connectionId })
};
}
}
const ddb = new AWS.DynamoDB.DocumentClient();
const onConnect = new OnConnect({ repository: ddb });
module.exports.handle = onConnect.handle.bind(onConnect);
onconnect.js
The onMessage will be responsible for receiving and rebroadcasting the messages for all connected clients. To do this, we will use the @connections API from AWS.
const AWS = require('aws-sdk');
const { CONNECTIONS_WEBSOCKET_TABLE } = process.env;
class OnMessage {
constructor({ repository }) {
this.repository = repository;
}
async handle(event) {
try {
let connectionData = await this.findConnections();
const postCalls = this.postMessages(event, connectionData.Items);
await Promise.all(postCalls);
} catch (err) {
return {
statusCode: 500,
body: err.stack
};
}
return { statusCode: 200 };
}
async findConnections() {
return this.repository.scan({ TableName: CONNECTIONS_WEBSOCKET_TABLE, ProjectionExpression: 'connectionId' }).promise();
}
postMessages(event, items) {
const apigwManagementApi = new AWS.ApiGatewayManagementApi({
endpoint: event.requestContext.domainName + '/' + event.requestContext.stage
});
const message = JSON.stringify(JSON.parse(event.body).data);
return items
.map(async ({ connectionId }) => {
try {
await apigwManagementApi.postToConnection({ ConnectionId: connectionId, Data: message }).promise();
} catch (err) {
if (err.statusCode === 410) {
console.log(`Found stale connection, deleting ${connectionId}`);
await this.repository.delete({ TableName: CONNECTIONS_WEBSOCKET_TABLE, Key: { connectionId } }).promise();
} else {
throw err;
}
}
});
}
}
const ddb = new AWS.DynamoDB.DocumentClient();
const onMessage = new OnMessage({ repository: ddb });
module.exports.handle = onMessage.handle.bind(onMessage);
onmessage.js
Lastly, the $disconnect route will remove the connection ID from DynamoDb, which shouldn’t receive anything else.
const AWS = require('aws-sdk');
class OnDisconnect {
constructor({ repository }) {
this.repository = repository;
}
async handle(event) {
const deleteParams = {
TableName: process.env.CONNECTIONS_WEBSOCKET_TABLE,
Key: {
connectionId: event.requestContext.connectionId
}
};
try {
await this.repository.delete(deleteParams).promise();
} catch (err) {
return {
statusCode: 500,
body: 'Failed to disconnect: ' + JSON.stringify(err)
};
}
return { statusCode: 200, body: 'Disconnected.' };
}
}
const ddb = new AWS.DynamoDB.DocumentClient();
const onDisconnect = new OnDisconnect({ repository: ddb });
module.exports.handle = onDisconnect.handle.bind(onDisconnect);
ondisconnect.js
To complete our journey, we need to deploy our functions to AWS, making this service available for testing. For deployment, simply run sls deploy. By default, the Serverless Framework will publish our services at us-east-1 region on AWS and dev stage, but we can change this by using the –region and –stage flags. After the deployment process finishes, the Serverless Framework will list the URL we should use to interact with our WebSocket server.
Output of SLS deployment
Time to Test
To test our function, let’s use wscat, a convenient tool for testing a WebSocket API.
As we can see, our server is rebroadcasting the messages to all connected clients.
Finally, to avoid any costs with the provisioned resources, like our DynamoDB table, we can remove our stack from AWS by running sls remove and passing the region and stage if it wasn’t deployed with the default values.
Conclusion
That’s it. Now you know how easy it is to develop and publish a fully serverless WebSocket server with the help of Serverless Framework.
A major advantage of building a WebSocket server using the AWS API Gateway with the DynamoDb is that it will handle the auto scale for us independently of how many clients the application has in the future. This solves many problems developers face when using a non-serverless architecture.
You can view the source code of this post here.
