In Real-Time notification system using WebSockets on AWS
⚡ TL:DR
We will cover how to efficiently create a WebSockets API Gateway and send real-time messages from our Lambda functions to our WebSockets client to make our applications more interactive.
📋 Table of content
· 🌍 Introduction
· 💪 Goals and Objectives
· 🛠️ Requirements
· 🤔 Solution Architecting
· 🛠️ Implementing our solution
· 👀 Room for improvements
· 🏁 Conclusion
🌍 Introduction
In today’s fast-paced digital landscape, the demand for instant communication and real-time updates has become more and more important. Whether it’s receiving breaking news alerts, tracking deliveries, or staying connected with friends and colleagues, the need for timely information has never been greater. In such a dynamic environment, traditional methods of communication often fall short, unable to provide the immediacy and responsiveness required.
This is where real-time notification systems step in to bridge the gap. By leveraging technologies like WebSockets, these systems enable seamless, bidirectional communication between clients and servers, allowing for instantaneous updates and alerts.
In the past year, I’ve completed a project for a client that was using a fully event-driven architecture, and for me, it was missing a crucial feature, real-time notification.
💪 Goals and Objectives
In this article, we will dive into the creation of a simple system on AWS by providing an architecture diagram, a few principles, the limits of performing this architecture on AWS, and a little bit of code.
We will be using the following AWS services:
- AWS API Gateway (Rest API and WebSockets API)
- AWS Lambda
- AWS Simple Queue Service
- AWS DynamoDB
🛠️ Requirements
In this article, we will mostly be using AWS so make sure you have an account ready to play with.
🤔 Solution Architecting
We have a simple web application that performs Rest API calls, some methods will create new AWS EC2 instances and take some time, we want to be notified through the whole creation process, at every step.
Here we are, with a blank page, tons of ideas in mind but can’t get them out. We know what we are looking for, a simple notification system. While some AWS services out there perform notification such as SNS (Simple Notification Service), it might not be our best fit for real-time notification.
When thinking about web and Rest API, an idea would be to store notifications in a database and make queries to our API that fetch stored notifications. But how should we fetch these notifications? Should we add a button to our web application that fetches these notifications? Should we create code that automatically makes the request to the right API endpoint every X seconds and collects the notifications? This is a solution, but not the one we are looking for.
What we want to do is real-time, a protocol that comes into my mind is WebSockets. WebSocket is a communication protocol that provides full-duplex communication channels over a single, long-lived connection between a client and a server. Unlike traditional HTTP connections, which are stateless and involve a request-response cycle meaning that the connection is terminated when the response is received, WebSockets enable persistent, low-latency communication.
With that definition in mind, WebSockets seems to be a solid candidate for our problem. On top of that, the AWS API Gateway service allows us to create WebSockets API, just what we need.
We now have to think about how we correlate our Rest API requests and WebSockets communication. If we take a look at how WebSockets API Gateway works on AWS, we might find one solution. If we refer to the AWS documentation, to send a message to a WebSocket client, we need to perform an HTTP request to the following URL:
POST https://{api-id}.execute-api.us-east-1.amazonaws.com/{stage}/@connections/{connection_id}Here, the interesting part of the URL is “{connection_id}”. That means that each time a user connects to the WebSockets API Gateway a connection ID is generated. Don’t we have a similar feature on Rest API Gateway? When requesting Rest API Gateway, a Request ID is generated, awesome! We now have a way to correlate WebSocket Connection and Rest API requests. For each request made to the Rest API Gateway, we will have to provide the WS connection ID. We can use a database solution to store these IDs and use them in our system to send messages later.
We have to find a way to send messages from AWS to our web application. Since we are dealing with an application that creates instances, we may want to send messages once instance creation starts, when instances are successfully created, and when instances are usable. Wouldn’t it be convenient to have a centralized way of delivering messages?
First thing first, we saw earlier that sending a message to one WebSocket client can be done by making a POST HTTP request to a specific URL. We can create a Lambda Function that will perform this HTTP request. In order to invoke that lambda function, we can either invoke the Lambda Function directly from other lambda functions, we can also create an SQS Queue that will receive every message and then send them to our Lambda Function, this is the solution that we will use because it decouples our architecture, it is more resilient to failure and it gives us more visibility about the number of messages to process.
We have one last topic to address, we want to send messages when the instance is created and ready to be used, but how? We can leverage Amazon EventBridge and catch those events.
At the end, we should have the following architecture:
🛠️ Implementing our solution
Before actually starting to build our solution, let’s give some context here and define some boundaries:
- In the provided architecture diagram above, we have mentioned the use of an SQS queue. We won’t be implementing it here.
- We’ve also suggested the use of DynamoDB to store Requests IDs and remove them on disconnection, we won’t be implementing them on here.
- We won’t deploy a Rest API Gateway as well, we will simply deploy a WebSockets API Gateway and create a Lambda that we will trigger manually to send back messages to our WebSocket client.
We have a simple web application that allows users to create VDI instances (Virtual Desktop Interface) on the go. The application is fairly simple, users are authenticated and can request their instance from the web interface.
Users would like to get in real-time notifications about their order of instance.
Here, users want to create new instances. Creation of EC2 instances can take some time due to many factors:
- EC2 creation time
- Instance setup (using any automation solution)
- Compliance check
The user requested to get an update at each step of the instance creation process. Let’s create our notification system, starting on AWS!
First, let’s log in to AWS and head to the API Gateway service. Create a new WebSocket API Gateway.
When communicating to a WebSocket API Gateway, you send a payload, in this payload you will need to specify the action you are trying to perform. In this input, you set which key inside of your payload is the action you are trying to perform. We will come back to this later.
Add the $connect and “$disconnect” predefined routes to your API Gateway.
- $connect: This route is the one called when your client connects to the WebSocket API Gateway
- $disconnect: This route is the one called when your client disconnects to the WebSocket API Gateway
In our use-case, we don’t need to add any custom routes yet.
Then you will want to create some Lambda Functions, one for the $connect route, and another one for the $disconnect route. This is the Lambda function that will be called when the client connect/disconnect to our API Gateway. Remember earlier when we said that we will have to match WebSocket connection ID and Rest API Request ID? Well, the $connect lambda function is a good place to write our logic that store the connection ID somewhere.
In the $disconnect lambda, we can write the logic that removes our stored connection ID.
We will write the lambda code later.
The final step is the stage. Here we are testing things, so we will simply create a first stage named “development”. Click on “Next” and deploy your API Gateway.
Once our WebSocket API Gateway is deployed, let’s try it. We will use Postman which supports WebSockets.
Grab your API Gateway WebSocket URL and open Postman.
Click on Connect, and you shall see that the connection is successful. Great!
We will add some code to our $connect Lambda function to analyze what’s being passed to that lambda.
import json
def lambda_handler(event, context):
print(json.dumps(event))
return {
'statusCode': 200,
'body': json.dumps('Hello from Lambda!')
}This code will simply dump every passed parameter to that Lambda. Let’s disconnect and then reconnect to our WebSocket API and see those parameters.
{
"headers": {
"Host": "c9if5l9rgd.execute-api.eu-west-1.amazonaws.com",
"Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits",
"Sec-WebSocket-Key": "q2ueME10A76weGCzBqgq2Q==",
"Sec-WebSocket-Version": "13",
"X-Amzn-Trace-Id": "Root=1-65feee73-002d1ca83f839f353a4c89b4",
"X-Forwarded-For": "xxx",
"X-Forwarded-Port": "443",
"X-Forwarded-Proto": "https"
},
"multiValueHeaders": {
...
},
"requestContext": {
"routeKey": "$connect",
"eventType": "CONNECT",
"extendedRequestId": "VFoyGG1vDoEF2dg=",
"requestTime": "23/Mar/2024:15:00:03 +0000",
"messageDirection": "IN",
"stage": "development",
"connectedAt": 1711206003575,
"requestTimeEpoch": 1711206003576,
"identity": {
"sourceIp": "xxx"
},
"requestId": "VFoyGG1vDoEF2dg=",
"domainName": "c9if5l9rgd.execute-api.eu-west-1.amazonaws.com",
"connectionId": "VFoyGcaIjoECIjw=",
"apiId": "c9if5l9rgd"
},
"isBase64Encoded": false
}Did you notice the important parameter here? “connectionId” that has a value of “VFoyGcaIjoECIjw=”. Note that this number is unique and generated on the client connection.
We will now create a new Lambda Function that will send a message to a given WebSocket connectionId, to ensure that we are able to reach this connectionId.
import json
import boto3
api_gw_id = "c9if5l9rgd" # Change with your WebSocket API Gateway ID
region = "eu-west-1" # Change with the region where your WebSocket API Gateway ID is
def send_message(connectionId: str, message: str, stage: str) -> None:
client = boto3.client(
"apigatewaymanagementapi",
endpoint_url=f"https://{api_gw_id}.execute-api.{region}.amazonaws.com/{stage}",
)
client.post_to_connection(
ConnectionId=connectionId, Data=json.dumps(message).encode("utf-8")
)
def lambda_handler(event, context):
send_message(
connectionId=event["connectionId"],
message=event["message"],
stage=event["stage"]
)
return {
'statusCode': 200,
'body': json.dumps('Hello from Lambda!')
}Create a new test, and set these values:
{
"connectionId": "VFoyGcaIjoECIjw=", # The Connection ID caught in the $connect dump
"stage": "development",
"message": "Hello Medium" # The message you want to send
}
We are not done yet! We need to give additional permission to our Lambda Function otherwise we will face a permission denial while invoking our test!
{
"Sid": "Statement1",
"Effect": "Allow",
"Action": [
"execute-api:ManageConnections"
],
"Resource": [
"arn:aws:execute-api:YOUR_REGION:YOUR_ACCOUNT_ID:API_GATEWAY_ID/*/*"
]
}Now, execute your test. You should now see your message on Postman!
We’re not quite done yet. There is another thing to do, include theses messages in our web application.
There are tons of ways to perform WebSocket. Here we will use Vanilla JavaScript.
We will simply add the following code to our web application:
<script>
const socket = new WebSocket('ws://your-websocket-server-address');
// Event listener for when the connection is established
socket.addEventListener('open', function (event) {
console.log('WebSocket connected');
});
// Event listener for incoming messages
socket.addEventListener('message', function (event) {
// Display an alert with the received message
alert('Message received: ' + event.data);
});
// Function to send a message to the WebSocket server
function sendMessage(message) {
if (socket.readyState === WebSocket.OPEN) {
socket.send(message);
console.log('Message sent: ' + message);
} else {
console.error('WebSocket connection is not open.');
}
}
</script>Once this code is added, we can refresh the page. Our web application will then send a connect request to our WebSocket API Gateway. On the AWS side, just like before, we will fetch the connection ID, and manually run our Lambda function.
And guess what? Our Web Application displays our message! We have made it!
👀 Room for improvements
Well, as you are probably wondering, I have mentioned earlier the use of a REST API Gateway, an SQS queue but not above, why?
The general idea was to provide you with a simple use case of how to use WebSockets within the AWS environment.
On our Web Application, we have a small interface that allows users to create virtual desktop instances. When clicking on the “Request New Instance” button, an HTTP request is sent to our REST API Gateway. When sending this request, we will need to send the Connection ID into our initial payload, that way we will be able to send notifications later, during our instance initialization process.
Each Lambda Function has to send a message to our SQS queue, to ensure message ordering and delivery. We can then either configure a Lambda that’s triggered on message pushed that will send messages to our WebSocket client or having a scheduler that will invoke our Lambda function and fetch messages from our SQS Queue.
🏁 Conclusion
We have learned how to use WebSockets on AWS, and it wasn’t that hard, right? I hope that you have enough resources to dig on your own, and maybe implement some of my ideas in one of your projects! Feel free to comment on this post with a link to your project, I will happily take a look at it.
