WebSocket is an effective way for full-duplex, real-time communication between a web server and a client. It is widely used for building real-time web applications along with helper libraries that offer better features. Implementing WebSockets requires a persistent connection between two parties. Serverless functions are known for short execution time and non-persistent behavior. However, with the API Gateway support for WebSocket endpoints, it is possible to implement a Serverless service built on AWS Lambda, API Gateway, and DynamoDB.

Prerequisites

A basic understanding of real-time web applications will help with this implementation. Throughout this article, we will be using Serverless Framework for developing and deploying the WebSocket service. Also, Node.js is used to write the business logic. 

Behind the scenes, Serverless uses Cloudformation to create various required resources, like API Gateway APIs, AWS Lambda functions, IAM roles and policies, etc.

Why Serverless?

Serverless Framework abstracts the complex syntax needed for creating the Cloudformation stacks and helps us focus on the business logic of the services. Along with that, there are a variety of plugins available that help developing serverless applications easier.

Why DynamoDB?

We need persistent storage for WebSocket connection data, along with AWS Lambda. DynamoDB, a serverless key-value database from AWS, offers low latency, making it a great fit for storing and retrieving WebSocket connection details.

Overview

In this application, we’ll be creating an AWS Lambda service that accepts the WebSocket connections coming via API Gateway. The connections and subscriptions to topics are persisted using DynamoDB. We will be using ws for implementing basic WebSocket clients for the demonstration. The implementation has a Lambda consuming WebSocket that receives the connections and handles the communication. 

Base Setup

We will be using the default Node.js boilerplate offered by Serverless as a starting point.

CODE: https://gist.github.com/velotiotech/d7b99ac4f55333226ec88cdf8306f127.js 

A few of the Serverless plugins are installed and used to speed up the development and deployment of the Serverless stack. We also add the webpack config given here to support the latest JS syntax.

Adding Lambda role and policies:

The lambda function requires a role attached to it that has enough permissions to access DynamoDB and Execute API. These are the links for the configuration files:

Link to dynamoDB.yaml

Link to lambdaRole.yaml

Adding custom config for plugins:

The plugins used for local development must have the custom config added in the yaml file.

This is how our serverless.yaml file should look like after the base serverless configuration:

CODE: https://gist.github.com/velotiotech/e34ebddcbe980032c13a9967410f4f04.js

Add WebSocket Lambda:

We need to create a lambda function that accepts WebSocket events from API Gateway. As you can see, we’ve defined 3 WebSocket events for the lambda function.

• $connect
• $disconnect
• $default

These 3 events stand for the default routes that come with WebSocket API Gateway offering. $connect and $disconnect are used for initialization and termination of the socket connection, where $default route is for data transfer.

CODE: https://gist.github.com/velotiotech/d2dcb5b7fbe63f681a2de2b8c7960dac.js

We can go ahead and update how data is sent and add custom WebSocket routes to the application.

The lambda needs to establish a connection with the client and handle the subscriptions. The logic for updating the DynamoDB is written in a utility class client. Whenever a connection is received, we create a record in the topics table.

CODE: https://gist.github.com/velotiotech/9740a04578742b4603476ef94b2e5793.js

The Client utility class internally creates a record for the given connectionId in the DynamoDB topics table.

CODE: https://gist.github.com/velotiotech/a05239bf9e98ae5cd7b078eae7a67930.js

Similarly, for the $disconnect route, we remove the INITIAL_CONNECTION topic record when a client disconnects.

CODE: https://gist.github.com/velotiotech/2267009b57d6284403e131dad8caa9a6.js

The client.unsubscribe method internally removes the connection entry from the DynamoDB table. Here, the getTopics method fetches all the topics the particular client has subscribed to.

CODE:  https://gist.github.com/velotiotech/e86fa012ead2f3370eb859e11cd01022.js

Now comes the default route part of the lambda where we customize message handling. In this implementation, we’re relaying our message handling based on the event.body.type, which indicates what kind of message is received from the client to server. The subscribe type here is used to subscribe to new topics. Similarly, the message type is used to receive the message from one client and then publish it to other clients who have subscribed to the same topic as the sender.

CODE: https://gist.github.com/velotiotech/6b1acefc9c45a06097ffef0889dddca7.js

Similar to $connect, the subscribe type of payload, when received, creates a new subscription for the mentioned topic.

Publishing the messages

Here is the interesting part of this lambda. When a client sends a payload with type message, the lambda calls the publishMessage method with the data received. The method gets the subscribers active for the topic and publishes messages using another utility TopicSubscriber.sendMessage

CODE: https://gist.github.com/velotiotech/c214c301df4c6dff0ccf3957c099e8e2.js

The sendMessage executes the API endpoint, which is the API Gateway URL after deployment. As we’re using serverless-offline for the local development, the IS_OFFLINE env variable is automatically set.

CODE: https://gist.github.com/velotiotech/21d9869277e05c57044ee4c55884f080.js

Instead of manually invoking the API endpoint, we can also use DynamoDB streams to trigger a lambda asynchronously and publish messages to topics.

Implementing the client

For testing the socket implementation, we will be using a node.js script ws-client.js. This creates two nodejs ws clients: one that sends the data and another that receives it.

CODE: https://gist.github.com/velotiotech/8e3bee1f4732eda2b3ad8ce1cb74aab4.js

The first client on connect sends the data at an interval of one second to a topic named general. The count increments each send.

CODE: https://gist.github.com/velotiotech/94dcb281674f48033f6893571f51ac53.js

The second client on connect will first subscribe to the general topic and then attach a handler for receiving data.

CODE: https://gist.github.com/velotiotech/5fc6b2b55a2cd82886acf2a99c323467.js

Once the service is running, we should be able to see the following output, where the two clients successfully sharing and receiving the messages with our socket server.

Conclusion

With API Gateway WebSocket support and DynamoDB, we’re able to implement persistent socket connections using serverless functions. The implementation can be improved and can be as complex as needed.

WebSocket is an effective way for full-duplex, real-time communication between a web server and a client. It is widely used for building real-time web applications along with helper libraries that offer better features. Implementing WebSockets requires a persistent connection between two parties. Serverless functions are known for short execution time and non-persistent behavior. However, with the API Gateway support for WebSocket endpoints, it is possible to implement a Serverless service built on AWS Lambda, API Gateway, and DynamoDB.