How to Create WebSocket Backend with API Gateway and Micronaut

EDIT: After evaluating similar API in production I need to warn you that implementing the similar WebSocket event bus using API Gateway may cost you four digits number so use with caution.

Since December 2018 you have a choice to create either REST or WebSocket API when you create new API using AWS API Gateway.

Once a new WebSocket API is created you need to declare new route handlers:

There are three predefined routes keys $connect, $disconnect , $default. You can declare more routes if you need to. The route key is derived using the Route Selection Expression defined when the API was created but it can be changed later.

Obviously, we need to implement the Lambda function first. Micronaut provides great support for creating AWS Lambda functions. Make sure you have the latest Micronaut CLI installed and run following command:

mn create-function mn-ws-echo

First of all, you need to open build.gradle file and add a new dependency on AgoraPulse Micronaut Libraries:

dependencies {
    compile 'com.agorapulse:micronaut-aws-sdk:1.0.4.4'

    // .. rest of the default dependencies
}

Once you have the dependency in place you can update the generated function bean class MnWsEchoFunction:

@FunctionBean("mn-ws-echo")
public class MnWsEchoFunction 
        extends FunctionInitializer
        implements Function<WebSocketRequest, WebSocketResponse> {

    @Inject MessageSenderFactory factory;

    @Override
    public WebSocketResponse apply(WebSocketRequest event) {
        RequestContext ctx = event.getRequestContext();
        MessageSender sender = factory.create(ctx);
        String connectionId = ctx.getConnectionId();

        switch (ctx.getEventType()) {
            case CONNECT:
                // e.g. register new connection in cache
                break;
            case MESSAGE:
                sender.send(connectionId, event.getBody());
                break;
            case DISCONNECT:
                // e.g unregister connection from cache
                break;
        }

        return WebSocketResponse.OK;
    }

}

There is a couple of changes which needs to be done

1. The function bean has to implement Function<WebSocketRequest, WebSocketResponse>
2. The handler method has to return WebSocketResponse.OK
3. There is MessageSenderFactory being injected which helps to create MessageSender instance which can send messages back to the connected client
4. If the new event is MESSAGE we send the message back to the current client

Before we deploy the function we need to have the Lambda role ready which allows us to call the connection endpoint and reply to the messages.

First, we need to create a policy to execute the WebSocket API @connections endpoint:

The policy can allow all execution API actions and should be restricted to the created API. Use your own Api id when creating the policy.

Once the policy is created we can create the role for the Lambda function:

The role is using the policy created in the previous step and also AWSLambdaBasicExecutionRole policy. Update the role in the build.gradle file and you are ready for deployment. You may also need to increase the memory a bit.

task deploy(/* ... */) {
    // ..
    role = "arn:aws:iam::794392443626:role/mn-ws-echo-role
    // ...
    memorySize = 512
    // ...
}

Be sure you supply your AWS credentials using the credentials file. Then you can run Gradle deploytask to deploy the function:

./gradlew deploy

Now we can finally set up the route handlers for $connect, $disconnect and $default routes to mn-ws-echo:

Once set we can deploy the API using the dropdown Actions button.

Select the new stage and call it test. When the stage is deployed you can see the WebSocket URL and Connection URL in the detail page.

You can use the WebSocket URL to test if everything is working as expected using this simple WebSocket tester:

WebSocket Test

Paste the WebSocket URL into endpoint input and click Connect. As soon as you see the text CONNECTED which may take a while you can write some text into the message text area and click Send. If everything was set up properly you should see the text starting with SEND followed by the same text starting with MESSAGE.

Please, read WebSockets for API Gateway part of the documentation for further reference and use cases.

We’re Hiring