Staying Connected

Maintaining a sustained connection with WebSockets and GraphQL Subscriptions in a serverless environment

Photo by https://unsplash.com/@plhnk

As modern technology and web applications advance, the expectation of immediate feedback and constant communication is apparent. We’re accustomed to this phenomenon of instant gratification where we are fulfilled with very little or hardly any wait time at all. Just think about how we’re able to receive certain packages with same day delivery. The world is at our finger tips. This continuous cause and effect cycle needs to be seamless, otherwise you run the risk of the end user becoming impatient or frustrated.

WebSockets

One way to achieve real time updates is through a persistent connection. This is where WebSockets come into play. A WebSocket is a transport protocol which provides a bi-directional communication channel between client and server via a sustained connection. Typical HTTP/S requests are uni-directional meaning only allowing for requests to be sent from client to server.

WebSockets render polling (aka constant request/response cycles) obsolete. Once a connection is established, connection channels persist and are re-used to exchange messages back-and-forth. Some popular use cases for WebSockets are real time applications such as chat apps and multiplayer games. The server is able to push new updates without the need for frequent requests from the client. Instant gratification!

Serverless

Serverless architecture enables us to build and run our applications without the infrastructure overhead. When it comes to serverless, it’s helpful to keep in mind that it is strictly event based. Events trigger other events. In the case of AWS Lambda, it will run our code on an as-needed basis in response to a given event. Due to this, it is not a long running process and is very transactional in nature.

This brings me to the question at hand, what if we wanted to take advantage of WebSockets but with all the advantages of serverless architecture? Existing WebSocket implementations are inherently stateful while serverless is purely event driven. What a conundrum. If only there was a way to maintain the connection so new Lambda invocations can reference it…🤔

Persistence Layer

Amazon API Gateway provides a WebSocket API which is essentially a collection of routes. Based on the message content it receives from the client, WebSocket API will determine which route to use (we’ll dive into this more in just a second). Lambdas are short lived (max 15 minute duration), so we need a way to persist the information so we can continue to listen for new events within the socket we’ve established.

When the client initiates a new WebSocket connection, API Gateway issues a connection ID for reference. Amazon DynamoDB will be our persistence layer to store that ID and keep track of any subscriptions (stay tuned), headers, request context, etc. that we will need to send data back to the client. Similarly, on disconnect, we will need the connection ID to remove the item from the table since we are no longer concerned with it.

Configuring Our Routes

In our serverless config, we will set up the three default routes that WebSocket API predefines for us in our Lambda function. We can create custom ones as well, but for this demonstration we’ll use the predefined ones.

• $connect — client initializes a connection to the API
• $disconnect — either client or server disconnects from the API
• $default — handles all other messages if no matching route is found

It’s important to remember Lambdas are not responsible for routing. They’re essentially just handlers. API Gateway forwards these requests to our Lambda handler function:

We must return a status code of 200 to inform API Gateway the event was handled successfully. Else, a status code of 500 will be returned to the client.

GraphQL Subscriptions

Before we go further, let’s discuss how we can incorporate subscriptions as a layer on top of WebSockets and use Apollo to instantiate our client on the frontend.

Subscriptions are a GraphQL feature that allow the server to push data to clients using a sustained connection such as a WebSocket. It’s important to remember that GraphQL is a query language. It is agnostic to any transport layer. Meaning we can send GraphQL queries via HTTP, WebSockets, paper planes — whatever your heart desires.

So how do we set up this WebSocket client then? We need to specify the subscription endpoint. This will be the WebSocket URL located in API Gateway. You can also set up custom domain names. You’ll notice the URL will look something like wss://. This is an encrypted protocol just like https:// but for WebSockets.

With Apollo, we can compose hybrid links. Queries and mutations will use an HTTP Link, while subscriptions will use a WebSocket Link. We’ll be creating this Link with graphql-ws — a GraphQL over WebSocket library.

Using the split method from Apollo, we can route a request to a specific middleware link. The first argument is a function that returns a boolean. The second and third arguments are Links. If the first argument returns true, the request will be forwarded to the second link. If false, it will forward to the third link. In our example above, the function is checking if the operation is a subscription and routing the requests accordingly. We can now pass the client we’ve instantiated with the proper link to our Apollo Provider.

Communication

The WebSocket protocol we’ll adhere to is graphql-transport-ws. Within this specification, messages are structured in JSON and stringified before being sent over the network. Every message contains a type field which is simply the action of the given message. Both the server and client must comply with this message structure.

Message Types

• connection_init — the first message server gets from the client asking to establish a connection. As we discussed earlier, we will persist the contents of this message in Dynamo for reference later on.

{
  type: 'connection_init',
  payload: {
    headers: {...}
  }
}

• connection_ack — server acknowledging a successful connection with the client after it’s been requested with connection_init.

{
  type: 'connection_ack',
  payload: {}
}

• subscribe (start) — client requests the start of a subscription. We also want to persist the contents of the subscribe message so we can keep track of all subscribers and pass the arguments to resolvers when there is an event. All future communication will be associated with this unique ID.

{
  id: 'subscription-id',
  type: 'subscribe',
  payload: {
    operationName: 'NewMessage',
    query: 'subscription NewMessage($chatId: ID!) {
      newMessage(chatId: $chatId) {
        id
        content
        createdAt
        senderName
      }
    }',
    variables: { chatId: '123' }
  }
}

To demonstrate, think back to the chat app mentioned in the beginning. We want to receive real time updates when a new message is sent in our chat. Apollo Client will execute the newMessage subscription on the frontend by establishing a connection to the GraphQL server and listen for response data to update the corresponding query (let’s say we call it getNewMessages). The server only publishes data to the client when a new message event is triggered on the backend.

• next—server sends operation execution result from the subscribe message back to the client.

{
  id: 'subscription-id',
  type: 'next',
  payload: { data: { newMessage: [] } }
}

• complete (stop)— client requests to stop subscribing. We will then remove/clean up the subscription from Dynamo.

{
  id: 'subscription-id',
  type: 'complete'
}

Sending Messages To Client

Now we understand the different message types, but how does the server communicate with the client? When the server wants to send a message back to the client, we need to instantiate the ApiGatewayManagementApi class.

Here, we can use the postToConnection method to send the message with the corresponding connectionId.

Bringing It Home

Pulling everything together we’ve just discussed, let’s go back to our Lambda handler function and fill in the gaps of how to write the connection and subscriptions to persistence.

There are a couple of different ways we can write to Dynamo. First, we’ll create the tables by adding to our serverless config (if you want to see an example check this out). The Amazon DynamoDB DataMapper library provides a simple way for us to define the relationship between items in our DynamoDB table and object instances. This object-relational mapping (ORM) allows us to perform basic put/get/update/delete (CRUD) operations and load items using the vocabulary we set in our application. Alongside DataMapper, DynamoDB’s DataMapper Annotations defines the model using Typescript property decorators. These annotations map to properties of records in our table — attributes, hash keys, tables, etc.

We’ll create two models:

Connection Model

Subscription Model

With our models defined and associated tables created, we can perform CRUD operations on objects from the table. We’ll do this by first instantiating DataMapper.

Using our mapper instance, notice how we’re able to create a get using the Connection model defined previously.

Now we’re ready to start using our persistence layer ✨. In our handler function, let’s go ahead and add our WebSocket protocol to the connection headers.

Next, we will properly handle writing/removing from persistence based on the message type from the client.

Almost done! Lastly, we want to clean up our table if the client or server decides to close the connection.

Final Thoughts 💭

Congrats! We’ve made it to the end. Phew. You’ve just learned what WebSockets are and have the tools to successfully incorporate them with serverless architecture. As an added cherry on top, we also know how to layer GraphQL Subscriptions to get the most out of our connection.

Gotchas

It goes without saying, not every solution or implementation is perfect. There are a couple of caveats or gotchas here. The first is in regards to the $disconnect event. In clean closures, the client intentionally sends a message to the WebSocket server asking to disconnect, resulting in the $disconnect handler being called immediately. In the case of unclean closures, where the user either closes the tab/browser or disconnects from the network, there is no message. Meaning, the server may not know the client has “disconnected” until idle timeout in API Gateway is reached (10 minutes). This is because API Gateway currently doesn’t support the WebSocket spec for ping/pong which detects idle connections. On the same note of idleness, this means even if both parties are connected and no message is sent over the socket, API Gateway will close it.

Alternatives

Another solution in the serverless world would be to use AWS Fargate over WebSocket API Gateway with Lambdas. Fargate runs on docker containers and uses Application Load Balancers (ALB) that support content-based routing and standard protocols such as WebSocket/HTTP. Clients maintain a connection with the load balancer and it forwards requests to a target group. It works like a traditional server as far as WebSockets are concerned since load balancers support sticky sessions. Having said that, keep in mind that unlike with Lambdas, you will need to manage auto-scaling yourself.

Wrapping Up 🎁

Handling stateful WebSockets in a stateless serverless environment can seem a bit daunting. After all, Lambdas are here for a good time, not a long time 🕺. However, with API Gateway and DynamoDB at our disposal, we’re able to maintain a persistent connection and effectively communicate between client and server. Lambdas, although short lived, can be super efficient in large scale production environments. Combined with how powerful WebSockets are for real time updates in applications, it’s certainly worth taking it for a test drive. I hope this inspires you to introduce WebSockets in your applications — serverless or not. Good luck friends!

Restaurant Brands International is Hiring 🍔 🍗 ☕️

See our open roles in the United States, Canada, and Europe.

Still in school? Check out our internship opportunities!

All trademarks and products are the respective property of their owners. No affiliation or endorsement is intended or implied.