Serverless Websocket with AWS: Chat Application

-

Serverless and Websockets seemed an unnatural combination to me so when I found out it was possible, I had to try it. Serverless is event-driven in it’s foundation so instead of unnatural it appears to be a match made in heaven. In this blog, I show how to set up a serverless websocket with AWS without much code. Easy as a breeze!

This may appear like a lot at first but I promise, we only need 120 (easy to read) lines of code to get the basics running: a chat application. Add 75 extra lines and three commands in the terminal and you have yourself automated deployment via CDK too.

chat serverless aws

The basics

We need to do a few things, so let’s spell them out:

  1. An index.html file for the UI
  2. An API Gateway with Websocket configuration
  3. Three lambda’s to connect, send messages and disconnect
  4. A DynamoDB table to store connection ID’s

The code

We only need an HTML form and Lambda’s to connect, disconnect and send messages. The API Gateway and DynamoDB table are configuration only.

HTML

Lets write the first 26 lines of code, a simple form. Note that the Stop button isn’t needed, if you close the browser window it will terminate automatically and remove the database ID.

<div id="chat-box"></div>
<input type="text" id="message" placeholder="Enter message"/>
<button onclick="connect()">start</button>
<button onclick="sendMessage()">Send</button>
<button onclick="webSocket.close();">Stop</button>

<script>
let webSocket;
const chatBox = document.getElementById('chat-box');
const messageInput = document.getElementById('message');

function connect() {
    const webSocketURL = 'WEBSOCKET API GATEWAY URL HERE';
    webSocket = new WebSocket(webSocketURL);

    webSocket.onmessage = message => 
      chatBox.innerHTML += `<div> ${message.data} </div>`;
    webSocket.onopen = () => 
      chatBox.innerHTML += '<div>Connected!</div>';
    webSocket.onclose = () => 
      chatBox.innerHTML += '<div>Disconnected!</div>';
}

function sendMessage() {
    const message = messageInput.value;
    webSocket.send(JSON.stringify({message}));
    messageInput.value = '';
}
</script>

 Lambda’s

We need 3 lambda’s: one to connect, one to send messages to other connections and one to disconnect. I wrote them in Typescript. You need to remove the types to deploy them to a Lambda as javascript, or follow the guide using CDK (recommended!).

In this blog, I won’t show the imports, but you can always find them on GitHub (link at the bottom of this post).
Connecting

export const handler = async (event) => {
    const addConnectionParameters = {
        TableName: process.env.TABLE_NAME!,
        Item: {
            connectionId: event.requestContext.connectionId,
            timestamp: new Date().toISOString(),
        }
    };

    try {
        await dynamo.put(addConnectionParameters);
        return { statusCode: 200, body: 'Connected.' };
    } catch (err) {
        console.error('Error during onConnect:', err);
        return { statusCode: 500, body: 'Failed.' };
    }
};

Disconnecting

export const handler = async (event) => {
    const deleteConnectionsParameters = {
        TableName: process.env.TABLE_NAME!,
        Key: {
            connectionId: event.requestContext.connectionId
        }
    };

    try {
        await dynamo.delete(deleteConnectionsParameters);
        return { statusCode: 200, body: 'Disconnected.' };
    } catch (err) {
        console.error('Error during onDisconnect:', err);
        return { statusCode: 500, body: 'Failed.' };
    }
};

Sending messages
We read the connection ID’s from the database and tell the API Gateway to send them to all active connections that appeared in the database. You could add metadata to group connections.

export const handler = async ({ body, requestContext }) => {
    const messageData = JSON.parse(body!);

    const sendMessage = (connectionId: string) => {
        const apiGatewayClient = new ApiGatewayManagementApiClient({
            endpoint: `https://${requestContext.domainName}/${requestContext.stage}`
        });
        const postCommand = new PostToConnectionCommand({
            ConnectionId: connectionId,
            Data: Buffer.from(JSON.stringify(messageData))
        });
        return apiGatewayClient.send(postCommand);
    };

    const connectionTableName = {
        TableName: process.env.TABLE_NAME!,
    };
    const connections = await dynamoClient.scan(connectionTableName);

    const sendMessagesToAllConnections = connections.Items?.map((item) =>
        sendMessage(item.connectionId)
    );

    try {
        await Promise.all(sendMessagesToAllConnections!);
        return { statusCode: 200, body: 'Message sent.' };
    } catch (err) {
        console.error('Error during sendMessage:', err);
        return { statusCode: 500, body: 'Failed.' };
    }
};

CDK

CDK is not a requirement. This is for automated deployment and makes my life a lot easier, yours too probably! If you don’t want to use it you can read the configuration I used and apply that via the user interface AWS provides.

To begin, install NPM and the AWS CLI and log in

$ aws sso login

Next, install the CDK CLI using the command

$ npm i -g aws-cdk

Finally, create an empty directory for your project and run this command:

$ cdk init app --language=typescript

We need a database, 3 lambda’s with the code shown above and an API Gateaway (plus some settings to make them communicate).

First, let’s build a table using CDK:

const table = new dynamodb.Table(this, 'MessagesTable', {
  tableName: "ChatConnections",
  partitionKey: {
    name: 'connectionId',
    type: dynamodb.AttributeType.STRING
  },
  billingMode: dynamodb.BillingMode.PAY_PER_REQUEST
});

That table is going to be used to store the connection ID. In the data we will also store a timestamp to see when users connected, but that’s fully optional.

Next are the Lambda’s. They will be invoked by the API Gateway and store, read and delete the connection ID in the database. I stored the code for these lambda’s in a folder called ‘lambda’ in the root directory of the project.

const onConnectLambda = new lambda.Function(this, 'OnConnectLambda', {
  functionName: 'ConnectLambda',
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: 'connect.handler',
  code: lambda.Code.fromAsset(path.join(__dirname, '../lambda')),
  environment: {
    TABLE_NAME: table.tableName
  }
});

const onDisconnectLambda = new lambda.Function(this, 'OnDisconnectLambda', {
  functionName: 'DisconnectLambda',
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: 'disconnect.handler',
  code: lambda.Code.fromAsset(path.join(__dirname, '../lambda')),
  environment: {
    TABLE_NAME: table.tableName
  }
});

const sendMessageLambda = new lambda.Function(this, 'SendMessageLambda', {
  functionName: 'MessageLambda',
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: 'message.handler',
  code: lambda.Code.fromAsset(path.join(__dirname, '../lambda')),
  environment: {
    TABLE_NAME: table.tableName
  }
});

Now, we need an API Gateway with configuration to support websockets.

const webSocketApi = new apigateway.WebSocketApi(this, 'WebSocketApi', {
  connectRouteOptions: { integration: new integrations.WebSocketLambdaIntegration( 'connect', onConnectLambda ) },
  disconnectRouteOptions: { integration: new integrations.WebSocketLambdaIntegration('disconnect', onDisconnectLambda) },
  defaultRouteOptions: { integration: new integrations.WebSocketLambdaIntegration('message', sendMessageLambda ) }
});

const deploymentStage = new apigateway.WebSocketStage(this, 'DevelopmentStage', {
  webSocketApi,
  stageName: 'dev',
  autoDeploy: true
});

// Output the WebSocket URL
new CfnOutput(this, 'WebSocketUrl', { value: deploymentStage.url });

Make sure to copy the lambda code of the previous chapter into three files:

lambda/connect.ts
lambda/disconnect.ts
lambda/message.ts

Put the lambda folder in the root directory (not the /lib directory).

That’s it! We now have the code required to create everything. Lets do this thing. Only a few more commands:

$ npm run build
$ cdk deploy

See the results

The console should output the URL of the websocket connection. Make sure to replace the right string in the index.html file.

After that, you should be able to run the index.html file in any browser (no need for a webserver). If everything was alright, it should work!

If you see any errors and think you might have skipped something, clone my repository and run it from there to make sure there are no typo’s.

Conclusion

This websocket setup is easy! We hardly wrote any code and got this cool serverless setup that is super cheap and reliable. It feels awesome and makes reactive applications truly reactive. Note that there is a maximum message size of 128kb so for media or large queries you may want to refer to another type of API call and only send notifications that data has been updated through websockets. If you are building an application with websockets I would recommend reading up on event-driven architectures as that is a great fit. Check out the repository in the link below for the full code and CDK configurations.

https://github.com/vroegop/websockets-serverless-aws-blog