SoftClouds

Navigating the cloud, one byte at a time

Working with WebSocket APIs and the Serverless Framework

By AndriesOn

WebSockets are a real-time communication protocol that allow for bi-directional communication between a client and a server. AWS API Gateway provides support for WebSockets, allowing developers to build real-time, interactive applications.

API Gateway makes it easy to manage WebSocket connections by providing a stateless communication layer between the WebSocket clients and your backend. This allows you to focus on the logic of your application, rather than managing the underlying WebSocket infrastructure.

In API Gateway, each WebSocket connection is identified by a unique connection ID, which is generated when a client connects. When your backend needs to send a message to a client, it can make an HTTP call to a URI that contains the connection ID. API Gateway will then forward the message to the correct client over the WebSocket connection.

If we visualize that, it looks as follows:

Working with WebSocket APIs and the Serverless Framework

This approach allows your backend to remain stateless and eliminates the need to manage connection stickyness, as all connections are managed by API Gateway. Additionally, by using API Gateway, you don’t have to worry about the complexities of pubsub, as you can send messages directly to a specific client using the connection ID.

WebSocket routes in API gateway

AWS API Gateway provides the following 4 types of routes for WebSocket APIs:

  1. $connect: This route is used to handle the initial WebSocket connection. You can use this route to perform authentication, authorization, or any other setup that is required when a client connects.
  2. $disconnect: This route is used to handle the termination of a WebSocket connection. You can use this route to clean up resources and perform any necessary logging.
  3. $default: This route acts as a catch-all for incoming WebSocket messages that don’t match any other route. You can use this route to implement default behavior, such as sending a response to the client.
  4. Custom routes: In addition to the $connect, $disconnect, and $default routes, you can also define custom routes for specific WebSocket message types. You can use custom routes to implement specific behavior for different types of messages, such as sending a response to the client or triggering a Lambda function.

By using these routes, you can define the lifecycle behavior of a WebSocket client, from the initial connection to the termination of the connection, and everything in between.

An example Serverless Framework setup

In this example, we will be using the Serverless Framework to build a simple WebSocket application. Our example will handle connecting / disconnecting clients and will demonstrate how you can use a Lambda function to send a message to all connected clients.

I am assuming that you have Webpack and TypeScript configured. This means that you have the necessary dependencies installed, and have set up the appropriate configuration files, such as tsconfig.json and webpack.config.js, to enable building and deploying TypeScript applications with the Serverless Framework. If you don’t, here’s a useful guide that helped me.

Here’s an example serverless.yml file that sets up a simple WebSocket API using the Serverless Framework. It specifies handlers for the $connect and $disconnect routes, and it configures a DynamoDB table that we are going to use to store and query the connection IDs. There’s also an hello function in there, and IAM permissions that allow our functions to interact with DynamoDB.

service: serverless-ws-test

provider:
  name: aws
  runtime: nodejs16.x
  websocketsApiName: custom-websockets-api-name
  websocketsApiRouteSelectionExpression: $request.body.action
  websocketsDescription: Custom Serverless Websockets
  iamRoleStatements:
    - Effect: Allow
      Action:
        - dynamodb:Query
        - dynamodb:PutItem
        - dynamodb:DeleteItem
      Resource:
        - "*"

plugins:
  - serverless-webpack

functions:
  connect:
    handler: handlers/connect.handler
    events:
      - websocket:
          route: $connect

  disconnect:
    handler: handlers/disconnect.handler
    events:
      - websocket:
          route: $disconnect

  hello:
    handler: handlers/hello.handler

resources:
  Resources:
    ConnectionTable:
      Type: AWS::DynamoDB::Table
      Properties:
        TableName: Connection
        AttributeDefinitions:
          - AttributeName: "PK"
            AttributeType: "S"
          - AttributeName: "SK"
            AttributeType: "S"
        BillingMode: PAY_PER_REQUEST
        PointInTimeRecoverySpecification:
          PointInTimeRecoveryEnabled: true
        StreamSpecification:
          StreamViewType: NEW_AND_OLD_IMAGES
        KeySchema:
          - AttributeName: "PK"
            KeyType: "HASH"
          - AttributeName: "SK"
            KeyType: "RANGE"

Handling a connecting client

In the example below we create a handler that deals with a client that just connected to our WebSocket API. We store the connection ID in DynamoDB so we can query it later.

// handlers/connect.handler.ts

import * as AWS from 'aws-sdk';
import { APIGatewayEvent, APIGatewayProxyResult } from 'aws-lambda';

const dynamo = new AWS.DynamoDB.DocumentClient({
    region: 'eu-central-1',
    apiVersion: '2015-03-31',
});

const storeConnection = async (
  connectionId: string
): Promise<void> => {
    await dynamo.putItem({
        TableName: 'Connection',
        Item: {
            PK: 'Connection',
            SK: `Connection-${connectionId}`,
        },
    }).promise();
};

export const handler = async (
    event: APIGatewayEvent,
): Promise<APIGatewayProxyResult> => {
    await storeConnection(event.requestContext.connectionId)
    return {
        statusCode: 200,
        body: '',
    }
};

Disconnecting

Disconnecting works pretty much the same way, except instead of creating a new row in DynamoDB, we will have to remove it.

// handlers/disconnect.handler.ts

import * as AWS from 'aws-sdk';
import { APIGatewayEvent, APIGatewayProxyResult } from 'aws-lambda';

const dynamo = new AWS.DynamoDB.DocumentClient({
    region: 'eu-central-1',
    apiVersion: '2015-03-31',
});

const removeConnection = async (
  connectionId: string
): Promise<void> => {
    await dynamo.deleteItem({
        TableName: 'Connection',
        Key: {
            PK: 'Connection',
            SK: `Connection-${connectionId}`,
            ConnectionId: connectionId,
        },
    }).promise();
};

export const handler = async (
    event: APIGatewayEvent,
): Promise<APIGatewayProxyResult> => {
    await removeConnection(event.requestContext.connectionId)
    return {
        statusCode: 200,
        body: '',
    }
};

Sending a message to all connected clients

In this section, we will create an example of a Lambda function that will post a message to all connected clients. The function will first query a DynamoDB table to retrieve a list of all connected clients, and then send a simple “Hello World” message to each of them.

The first step is to write a generator function that loops through all the connections stored in DynamoDB:

//app/handlers/hello.handler

import { DocumentClient } from 'aws-sdk/clients/dynamodb';

async function* getPaginatedResults(
    params: DocumentClient.QueryInput,
): AsyncGenerator<
  DocumentClient.AttributeMap[] | undefined, 
  void, 
  unknown
> {
    let lastEvaluatedKey;
    do {
        const queryResult = await dynamo.query(params).promise();
        lastEvaluatedKey = queryResult.LastEvaluatedKey;
        params.ExclusiveStartKey = lastEvaluatedKey;

        yield queryResult.Items;
    } while (lastEvaluatedKey);
}

This is useful when the result set is too large to be returned in a single response, as DynamoDB has a limit on the maximum amount of data that can be returned in a single query.

Pagination in DynamoDB is achieved by using the LastEvaluatedKey attribute, which is returned in the response of a query operation. The LastEvaluatedKey attribute indicates the exclusive start key of the next page of results, and can be used as the ExclusiveStartKey parameter in a subsequent query to retrieve the next page of results.

For example, let’s say we have a DynamoDB table with a large number of items, and we want to retrieve all of the items in pages of 100 items each. The first query would be performed without specifying the ExclusiveStartKey parameter. The response from this query would include the first 100 items, and a LastEvaluatedKey attribute that indicates the start key of the next page of results.

Next, we would perform a subsequent query, using the LastEvaluatedKey attribute as the ExclusiveStartKey parameter. This would retrieve the next 100 items, and so on, until all items have been retrieved.

We can now use this function to loop through all our connections and post a message to each.

//app/handlers/hello.handler

import * as AWS from 'aws-sdk';
import { DocumentClient } from 'aws-sdk/clients/dynamodb';

// .... code from getPaginatedResults 

const apiGw = new AWS.ApiGatewayManagementApi({
    endpoint: 'https://abc123.execute-api.eu-central-1.amazonaws.com/env',
    region: 'eu-central-1',
    apiVersion: '2018-11-29',
});

const dynamo = new AWS.DynamoDB.DocumentClient({
    region: 'eu-central-1',
    apiVersion: '2015-03-31',
});

const sendMessage = (connectionId: string): Promise<void> {
    await apiGw.postToConnection({
        ConnectionId: connectionId,
        Data: 'Hello world!',
    }).promise();
}

export const handler = async (): Promise<void> => {
    const generator = getPaginatedResults({
        TableName: 'Connection',
        KeyConditionExpression: 'PK = :PK AND begins_with(SK, :SK)',
        ExpressionAttributeValues: {
            ':PK': 'Connection',
            ':SK': `Connection-`,
        },
        ScanIndexForward: true,
        Select: 'ALL_ATTRIBUTES',
    });

    for await (const items of generator) {
        await Promise.all(
          items.map(item => sendMessage(item.ConnectionId))
        );
    }
};

Summary

In this post, we learned about working with WebSockets in AWS API Gateway and the Serverless Framework. We covered the four types of routes that are related to the lifecycle of a WebSocket client, and how they can be used to manage WebSocket connections and communicate statelessly with a backend. We also explained how DynamoDB pagination works, which is the process of retrieving a portion of data in a DynamoDB table in multiple rounds or “pages”. By using the LastEvaluatedKey attribute.

I hope you enjoyed it!

Please note that the examples provided in this post are for illustration purposes only and are meant to provide a basic understanding of how to work with WebSockets in AWS API Gateway using the Serverless Framework. These examples have minimal error handling and are not meant for production use.