Realtime communication using API Gateway (Websockets), DynamoDB, and AWS Lambda

I’ve always worked with REST APIs, and it’s been quite a learning experience with every project that I’ve worked on. Recently, I took up a task to create a simple application that uses WebSockets to communicate. The application should update the client’s side autonomously when it receives a message over the communication channel without polling the backend at regular intervals. Regular polling increases the load on the backend should the number of users increase.

Hence, the first constraint was to use WebSockets as a mode of the communication channel with the backend. Secondly, I wanted to keep the technology stack within the AWS environment. Services like Pusher provide an easy way out by using their API. However, employing third-party services might not be feasible all the time.

The post's objective is to comprehend all the moving pieces and develop a small little POC (Proof of Concept) for ourselves. Moreover, I want to work with many managed services to keep my setup and management costs to the minimum.

You can find the code for the Lambda functions here: https://github.com/mychewcents/real-time-app

Let’s dive right into it. Shall we?

Let’s understand the role of each of the technologies that I plan to use:

1. API Gateway: To manage the WebSocket connections and route messages
2. DynamoDB: To store the active WebSocket connections
3. AWS Lambda: To process the requests and interact between API Gateway and DynamoDB

The overall architecture diagram of the article is shown below.

The full architecture of the system developed in the article

First, let’s set up a simple DynamoDB table that stores the connection attributes. We’ll name the key as “connectionId.” I’ve used the default settings to keep the setup phase to a minimum.

DynamoDB table to store active connections

Second, let’s create a WebSockets-based API Gateway using the AWS Console. I’ll make a basic API with defaults for now. I’ve used “Mock” as my endpoint integration to keep it simple. We have the following endpoints defined:

1. $connect: Provided by API Gateway
2. $disconnect: Provided by API Gateway
3. subscribe: A custom route to store the message type/topics that the user wants to listen

Finally, one would need to pass action in the BODY of the WebSocket’s message to trigger the subscribe endpoint.

Here’s how my final API Gateway looks like:

API Gateway review page for WebSockets API — 1

API Gateway review page for WebSockets API — 2

Finally, let’s create the AWS Lambda functions that would be responsible for handling our requests. We’ll need four Lambda functions:

1. Connect: Simple endpoint to return an empty response. (One can do a lot in the connect phase, but we’ll skip that for this article)
2. Subscribe: Save the connection details in DynamoDB for sending messages
3. Publish: Read from the DynamoDB table for active connections and send the messages via the API Gateway
4. Disconnect: Remove the link and delete any DynamoDB table entries for the connection

Let’s start with the “onConnect” lambda function. The function returns a status code 200 and a dummy message. We don’t really do much here and keep it simple. Moreover, I’ve used the default settings for the AWS Lambda and created a new role for the execution.

AWS Lambda information for the “onConnect” function

Second, the “onSubscribe” function. Following are the tasks we need to perform in the function:

1. Read the connectionId for the WebSocket connection.
2. Read the message type/topic that the user wants to subscribe to.
3. Store the above details in a DynamoDB table.

For step 3, we need to provide our Lambda function with DynamoDB access. Moreover, the function needs the table name to store the details. We provide the table name using Lambda’s ENV variables.

AWS Lambda information for the “onSubscribe” function — 1

AWS Lambda information for the “onSubscribe” function — 2

Now, head over to “Configuration > Environment Variables” to add the DynamoDB table name. The usage of the ENV variable removes any need to edit the code should the table name change.

Environment variables for the “onSubscribe” function

Now, let’s move to the “onDisconnect” function. Once the user closes the connection, we should remove the corresponding entries from the DynamoDB table. Hence, the procedure is very similar to “onSubscribe” regarding the accesses required and the ENV variables.

AWS Lambda information for the “onDisconnect” function

We’ll skip the ENV screenshot here as it is the same as the one above for “onSubscribe.”

Finally, we create the “onPublish” function to send out the messages to the active connections. Following are the tasks that the function needs to perform:

1. Read the message type/topic and the body.
2. Get all the active connections from the DynamoDB table.
3. Send the messages to the active connections using API Gateway.

For the function, apart from DynamoDB access, it should have access to API Gateway as well.

New role details for the “onPublish” function.

The above role might not be fulfilling the “Least Privilege Principle” and hence, should not be used in production.

AWS Lambda information on the “onPublish” function

Now, let’s add the required ENV variables. Apart from the table name, we also need to provide the API Gateway’s URL. It will be available in the “Stages” section of the API Gateway.

Environment variables for the “onPublish” function

Finally, we have all the pieces ready. Let’s integrate the API Gateway with the Lambda functions. We’ll be using the trigger as Lambda for each of the endpoint’s integration. The below screenshot showcases the “$connect” endpoint only. However, you can repeat the same steps for the other two endpoints too.

API Gateway settings for “$connect” endpoint

One thing to note here is that “onPublish” doesn’t have an endpoint. For now, we’re assuming that the backend application can send messages to the “onPublish” lambda function directly or via any other third party or AWS Services.

Now, deploy the API and use the “wss://…” to connect to from your frontend.