Feature Request: Powertools Web Sockets Support

[msailes]

Key information

• RFC PR:
• Related issue(s), if known: n/a
• Area: Web Sockets
• Meet tenets: Yes
• Approved by: ''
• Reviewed by: ''

Summary

Customers would like to use AWS Lambda and Amazon API Gateway Web Sockets but find it hard to build a suitable solution to handle managing the connections. Powertools could provide a best practice approach to managing connection information similar to how the Idempotency module works.

Motivation

To provide a best practice solution to working with API Gateway Web Sockets with AWS Lambda

Proposal

This is the bulk of the RFC.

Explain the design in enough detail for somebody familiar with Powertools to understand it, and for somebody familiar with the implementation to implement it.

If this feature should be available in other runtimes (e.g. Java), how would this look like to ensure consistency?

User Experience

How would customers use it?

Any configuration or corner cases you'd expect?

Demonstration of before and after on how the experience will be better

Drawbacks

Why should we not do this?

Do we need additional dependencies? Impact performance/package size?

Rationale and alternatives

• What other designs have been considered? Why not them?
• What is the impact of not doing this?

Unresolved questions

Optional, stash area for topics that need further development e.g. TBD

[rr-on-gh]

Proposal
Here is a proposal for the Powertools that uses annotations, similar to the Idempotency PowerTool. The code samples and much of the thoughts are based on Java as the target programming language.

TL;DR; Create annotations as part of the proposed PowerTools which users can import into their project. Users can annotate the Lambda functions that handle WebSocket events like $connect, $disconnect and other custom routes. The AspectJ code will intercept invocations to these functions and manage the co-relation between WebSocket connectionIDs and the customer entity in a DynamoDB table.

Initial Setup
The solution would use a DynamoDB table to store connection information which is created with a predefined schema. The table details are then initialized in the Lambda constructor similar to how the Idempotency PoweTools handles DDB creation:

    public Connect() {
        // we need to initialize connection store before the handleRequest method is called
        WebSocket.config().withPersistenceStore(
           DynamoDBPersistenceStore.builder()
                .withTableName(System.getenv("TABLE_NAME"))
                .build()
        ).configure();
    }

Scenario: Initial connection

Users would initiate the WebSocket connection to the API Gateway invoking the $connect route and if required, present the authentication token like JWT and API GW authenticates the request. API Gateway then invokes the Lambda function that handles connect events. The lambda would look something like this:

public class Connect implements RequestHandler<APIGatewayV2WebSocketEvent, APIGatewayProxyResponseEvent> {
    ... 
    @WebSocketConnect
    public APIGatewayV2WebSocketResponse handleRequest(APIGatewayV2WebSocketEvent input, Context context) {
        // Any business logic if needed
        return new APIGatewayV2WebSocketResponse();
    }

    @ExternalId
    private List<String> getExternalIDs(APIGatewayV2WebSocketEvent input, Context context) {
        //Eg: Extract external id from JWT token here.
        return externalIDs;
    }
}

The annotation WebSocketConnect uses AspectJ (similar to Idempotency powertool) and intercepts the Lambda’s handleRequest method invocation. The advice would then invoke the method that is annotated with @ExternalId. In the exampke above, it is the getExternalIDs method. This method would return a list of externalIds. An externalID is the identifier that connects the message that you want to send, to a WebSocket connection your end user has established. Let’s take an example of an Investment Management application. If the WebSocket connection’s business use case is to keep your end customers updated about their total portfolio value, then the externalID could be a ‘customer identifier’ that uniquely identifies them. However if the business use case of the WebSocket is to keep them informed on updates to a certain ticker symbol, then externalID could be the ticker symbol. If you want to do both on the same WebSocket connection, you can associate a WebSocket to multiple externalIDs. The advice invokes this method and then stores this as [connectionID, externalID] in DynamoDB and then the rest of the handleRequest is executed and the reply is sent back to customer via API Gateway

Scenario: Backend needs to push a message to WebSockets
Building up the investment management application, lets say every 5 minutes or so, the application needs to send the updated value of a certian ticker symbol to the connected clients. But clients should only get updates to the ticker symbols that they are subscribed to. An event (SQS, EventBridge, SNS etc) will trigger a backend lambda with AMZN in its payload indicating all customers interested in this ticker AMZN should be notified. The implementation of this Lambda would look something like this:

public class TickerUpdate implements RequestHandler<SQSEvent, String> {
    @WebSocketBroadcast
    public String handleRequest(SQSEvent input, Context context) {
        return "message to send";
    }

    @ExternalId
    private List<String> getExternalIDs(SQSEvent input, Context context) {
        return externalIDs;
    }
}

The advice backing @WebSocketBroadcast would intercept the call, get the externalIDs by invoking the getExternalIDs, fetch the list of corresponding websocket connectionids from DynamoDB and broadcast the updates to them.
The externalID should be designed taking this into consideration. For this example, storing AMZN as the externalID would result in a single call to the persistent store to fetch all WebSocket connections that are interested in AMZN.

Scenario: Disconnect from client and server side
This would work similar to connect workflow, and a @WebSocketDisconnect would delete the corresponding items from the DynamoDB table.

Scenario: Custom WebSocket route is invoked
This is essentially any call to a custom route on the WebSocket connection. In this case customer invokes the subscribeToTicker route to subscribe to a certain ticker symbol, for eg. MSFT. The lambda that handles this route is invoked by the API Gateway and the needed logic is executed in the Lambda. This results in a change to the business entity, i.e. the connections list maintained by the PowerTool must update to add an entry to the DynamoDB table for MSFT for this connectionID. This is done as follows:

public class SubscribeToTicker implements RequestHandler<APIGatewayV2WebSocketEvent, APIGatewayProxyResponseEvent> {

    @WebSocketModelUpdate
    public APIGatewayV2WebSocketResponse handleRequest(APIGatewayV2WebSocketEvent input, Context context) {
        // Any business logic if needed
        return new APIGatewayV2WebSocketResponse();
    }

    @ExternalId
    private List<String> getExternalIDs(APIGatewayV2WebSocketEvent input, Context context) {
        //Eg: Return the updated list of external IDs (??)
        return externalIDs;
    }
}

A Lambda that is annotated with @WebSocketModelUpdate will be intercepted and the business logic will be executed. At the end of the execution the advice will call the getExternalIDs method to get the updated list of externalIDs for this connection. The connectionId information in DynamoDB is then updated with the new set of externalIDs.

[willfarrell]

Within Middy (for NodeJS runtime), we have ws-routter and ws-response as part of our collection of middlwares people can use. Maybe these can be a starting point if this RFC moves forward.