How I implemented WebSocket API Gateway with AWS Lambda

App for… couting emojis

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?

DigitalOcean Droplets are Linux-based virtual machines (VMs) running on virtualized hardware. Each new server is a droplet. It can be either standalone or part of a larger infrastructure.

Before I go any further, I’d like to give you a brief refresher of Lambda and WebSocket. If you have more experience with those, skip to the next section (AWS Lambda = HTTP?).

Overview of AWS Lambda (and its limitations in terms of maximum duration)

AWS Lambda is a serverless computing service that lets you run code without provisioning or managing servers. AWS Lambda offers several benefits, such as cost savings, scalability, and ease of use, but it has some limitations that we need aware of.

The first (and most significant) limitation of AWS Lambda is the maximum duration. It allows running for a maximum of 15 minutes only. It should be enough for many use cases. However, if our application process a large amount of data, it can exceed that limit. To overcome these limitations, we can split our task into smaller functions or offload part of the logic to another service. 

There are some other AWS Lambda limitations, including concurrent executions, memory allocation, or even environment variables size. You can read more about those directly on the AWS site.

Introduction to WebSocket API Gateway (and its advantages over HTTP API Gateway)

The WebSocket API is a powerful tool that improves the real-time interactivity of applications. Traditional HTTP API relies on the request-response communication model. In contrast to that, WebSocket API has a different communication model. It creates long-term, bi-directional WebSocket connections between the client and server to communicate over it. The WebSocket API communication model allows the server to notify all of the connected client applications immediately when something changed. That helps in the development of chat applications, collaboration tools, real-time charts, or even gaming platforms.

We as developers don’t need to think about long-pooling techniques anymore and client applications don’t need to call the server every few seconds to show the latest data. And because WebSocket keeps a single, persistent connection, it eliminates latency problems that arise with the HTTP request-response model. Important concepts include WebSocket routes, WebSocket messages.

Despite those advantages, there are some scenarios where the WebSocket connection is not preferred. In case our, application benefits from caching (because our data changes rarely) WebSocket might be a bad choice. It’ doesn’t support any edge caching.

AWS Lambda = HTTP?

When we think about AWS Lambda and Lambda functions, the first thing that comes to our mind is the REST API, definitely not WebSocket APIs. 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 and Lambda function. 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.

Configuration

When it comes to WebSocket and AWS API Gateway, there are three predefined routes we need to handle:

• $connect – the route key responsible for connection initialization.
• $disconnect – the custom route called when the connected clients or server disconnect 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 the route selection expression, that is 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.

Sending messages

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:

https://<api-id>.execute-api.<region>.amazonaws.com/<stage>

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.

WebSocket API Gateway with Lambda… 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; in most cases DynamoDB is preferable because it has lower cost and development/infrastructure required. Any lambda can connect to DynamoDB Table without any additional setup, while in the case of the RDS, your lambda has to be in VPC. Thus, RDS Proxy and NAT Gateway may be required, which means additional costs.
  
• 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 these issues are peanuts in light of 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! 🚀

WebSocket API Gateway with Lambda – Frequently Asked Questions

Not that the big issue is out of the way, let’s ask a couple of other questions you might just have on your mind.

Can I use this approach successfully in a major commercial project?

It all comes down to the definition of “big”. The main thing you have to be aware of with any AWS services are quotas. There are limits for concurrent invocations, limits for new connections per seconds, per account limit, connection duration etc.

As long as you are aware of those and build your app with them in mind, you should be able to use the same idea for any commercial project.

Did we use it in any of our projects?

Yes, we did! We do have an experience in Lambda-based, WebSocket-based projects. One of them was a system that made it possible to manage active sessions on multiple devices, restricting access to content based on a variety of criteria. Unfortunately, I can’t share any more details at this moment. There a lot more projects that used Lambda among the AWS projects we developed or consulted.