04 August 2020
How I implemented WebSocket with AWS Lambda and API Gateway
AWS Lambda is known for its scalability and ease of use. However, it is also known for its limitation – the maximum duration of 15 minutes. That’s why, most of the time, it is connected with HTTP API. What if I told you that you can build WebSocket-related systems with AWS Lambda? Heresy! Let me tell you a story of a simple system that we migrated from DigitalOcean droplet to AWS Lambda. A system where WebSocket is a necessity. So here’s how I implemented WebSocket with AWS Lambda and API Gateway.
At The Software House, we use Slack… a lot. It became a kind of foundation for our company and so did all the emojis. 😀 Obviously, at some point we started to think about which emoji is used the most often. A really simple question but difficult to answer without coding. In order to get the necessary data, I decided to build a very simple app.
- Node.js backend app with MongoDB consisting of two parts:
- Express.js API used as a subscriber for Slack reacting added/removed events and GET reaction endpoint.
- WebSocket server allowing us to send information about changes in reactions count.
- React SPA used for displaying the current count.
Everything hosted on a Digital Oceans droplet.
Obviously, even the smallest droplet was overkill for such a small app. However, because of the stateful nature of WebSocket, I couldn’t use a Serverless approach… or could I?
AWS Lambda = HTTP?
When we think about AWS Lambda, the first thing that comes to our mind is the REST API. The second one is probably something related to cron-like jobs or event/queue handlers. But why is it so?
The answer is simple. All of those are stateless and have short execution time. That’s why we never hit Lambda maximum duration (15 minutes). On the other hand, WebSocket is stateful, we need to store existing connections somewhere to be able to send messages to them.
However, there is one thing I forgot. AWS Lambda is not responsible for routing or even connecting a user to specific Lambda code. In fact Lambda, it is just a handler that is called with two parameters – the event (caller related set of data for example HTTP event) and the context.
What is responsible for routing then you may ask. Well, it’s Amazon API Gateway.
Initially, API Gateway supported only RESTful APIs, however at the end of 2018, AWS added support for WebSocket communication. At the same time, they made it possible to use AWS Lambda for real-time apps.
When it comes to WebSocket and API Gateway, there are three predefined routes we need to handle:
- $connect – the route responsible for connection initialization
- $disconnect – the route called when the client or the server disconnects from API
- $default – the route used as the message default handler.
You can configure them using Serverless:
Of course, we can also configure custom handlers. First, we need to define websocketApiRouteSelectionExpression. This will allow us to point which field on a request will be used for route matching. For example, let’s define the action field to be used as a route source:
Now we only need to configure a custom handler for it, with a route defined for myCustomAction….
…and then send a request with a proper body.
Here we face the challenge. How to send a message? Lambdas are stateless, but we want to send a message to a user connected to API Gateway. So how can we do that?
WebSocket event gives us an access to a few important fields:
- routeKey – this one is used to get to know what operation we want to handle,
- connectionId – this one is a connectionId we can use for sending messages back to the user.
API Gateway has a special API called execute-api. By calling it, we can easily send messages to a specific connectionId.
This API is available on:
Of course, instead of building all of the requests manually, we can use a specialized API library. In this case, it is ApiGatewayManagementApi.
We are using Serverless Framework and so we can define API Gateway endpoint using it:
The last thing that is left is to use one of the methods of ApiGatewayManagementApi to send messages to the specific connectionId. For this, we can use the postToConnection method.
So… Does it work?
Yes, it does! Of course, there are a few challenges:
- we need to store connectionId somewhere – you can use RDS or DynamoDB for it,
- it requires a different architectural approach – I recommend making good use of SNS (for example, we can build a Lambda that will be responsible for sending messages to specific connections based on SNS message payload).
But it means nothing if we realize how scalable this approach is. Do you remember how problematic WebSocket scaling is?
With API Gateway and AWS Lambda – you get an architecture that allows you to scale infinitely! 🚀