Adding Web Sockets to Workers

Web Sockets allow for applications to communicate continuously without requiring HTTP requests to be explicitly executed to receive data from a server. What if you want to just be told when you need new data instead of keep asking? Parents would rejoice worldwide if they could keep their kids from asking “Are we there yet” every 30 seconds, and instead live in a world where the kids would be content with the parents telling them “We’re here”. That’s the same problem web sockets help the internet solve – lucky.

A more practical example on how it would be used is let us pretend we’re developing a chat application where two or more users can communicate. When using the HTTP method each user would be required to send a network request at some polling interval (let’s say every 5 seconds) to see if any new messages exist for that chat. Most of the time those HTTP requests come back with nothing new to report which just feels clunky.

What web sockets provide is the ability for any number of clients to “subscribe” to the server and listen for events. A client would tell the server “hey, I’m here” for each user. When a user sends a message to the server, everyone who has told the server they are there would get the message pushed to them immediately instead of them having to keep asking if a new message exists. How ideal.

Why in StarbaseDB?

TLDR; At the moment, convenience. In the future, optimizations.

Communicating between a client and server via web sockets can reduce latency because there is an already established tunnel of communication that can be used. Each subsequent message sent or received can be header-less and reduce the amount of excessive data a typical HTTP request/response might retain.

In the future we can provide a number of optimizations including CDC (Change Data Capture) where multiple clients can listen for events to be triggered on a database, such as a new entry into a table being written to it, and all clients can receive notification of that event. Streaming data back as the cursor iterates over the data SQL responds back with instead of waiting for it to be packaged together and sent back as a single object. How about two-way data syncing between a local data source and your database when internet connectivity exists.

There are a number of future optimizations and benefits we can develop where having a baseline web socket implementation opens up those future opportunities.

Worker Implementation

Time to get to work. I’ll show you how easy it is for you to support web socket connections in workers yourself, by essentially showing you the path we took.

Below what you will see is a scaffolding of our Worker code. Really when you boil it down we’re handling incoming HTTP requests to establish the connection with fetch(request: Request) and the other three functions will do the heavier lifting of solidifying the client/server connection, and sending and receiving messages.

export class DatabaseDurableObject extends DurableObject {
    // Map of WebSocket connections to their corresponding session IDs
    private connections = new Map<string, WebSocket>();

    constructor(ctx: DurableObjectState, env: Env) {
        super(ctx, env);
    }

    async fetch(request: Request): Promise<Response> {
        const url = new URL(request.url);

        // TBD...
    }

    // Defined by us
    clientConnected() {
        // TBD...
    }

    // Defined by Cloudflare
    async webSocketMessage(ws: WebSocket, message: any) {
        // TBD...
    }

    // Defined by Cloudflare
    async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean) {
        // TBD...
    }
}

Taking this one function at a time to keep it in chunks we can easily digest, let’s start with our fetch function. Our fetch functions primary objective is to receive an incoming HTTP request from our user, and establish the client/server relationship handshake for future web socket communications. Simply all we’re going to do here is pass the logic into our clientConnected() function to keep our fetch clean.

async fetch(request: Request): Promise<Response> {
    const url = new URL(request.url);

    if (url.pathname === '/socket') {
        return this.clientConnected();
    }
}

Makes sense for us to move onto the clientConnected() function details now. To establish the client/server relationship we need to create that pairing so the server knows who the client is, and then we’ll return the client back to the calling application so they can start using that communication tunnel.

clientConnected() {
    const webSocketPair = new WebSocketPair();
    const [client, server] = Object.values(webSocketPair);
    const wsSessionId = crypto.randomUUID();

    this.ctx.acceptWebSocket(server, [wsSessionId]);
    this.connections.set(wsSessionId, client);

    return new Response(null, { status: 101, webSocket: client });
}

Above we create a web socket pair, assign the pair with a randomUUID to uniquely identify it for future calls, and then tell the context of our application to go ahead and acceptWebSocket. You notice we have this connections Map object as well, we’re just storing the connections in a Map in case we want to implement checks in the future or handle logic such as rate limiting. I’ve added it just to show you how you can do it but if it doesn’t pertain to you then feel free to omit it. Lastly, we return back to the calling application the client object in the response which is something they will hold onto for future communication.

🎉 Hooray – connections are established! Now let’s get them talking. The next two functions we implement are defined and handled mostly by the Cloudflare Worker itself. When they are called is largely out of our control, but what logic they execute is fully within our control and that’s what we’ll focus on.

The primary function here is webSocketMessage(ws: WebSocket, message: any) as it handles receiving messages, and typically sending a message back. For StarbaseDB we use it to listen for when a user wants to complete a query action, execute their SQL statement and respond back with the results. Take a look at how our implementation roughly looks like.

async webSocketMessage(ws: WebSocket, message: any) {
    const { sql, params, action } = JSON.parse(message);

    if (action === 'query') {
        const response = this.executeSQL(sql, params)
        ws.send(JSON.stringify(response.result));
    }
}

It’s simple enough. We piece apart the components we want to utilize from our message, in this case is a key/value object, and hold reference to them. If our action value is query then we proceed to execute some custom code for our user. Once that statement has provided a response we can call ws.send(…) to respond back to the requesting client with any message we’d like. Actually quite a simple implementation.

Last but not least it’s good practice for us to handle closing sockets. Cloudflare Workers gives us another helper method in webSocketClose(…) for us to observe when a socket connection has closed for a client so we can handle custom logic.

async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean) {
    // If the client closes the connection, the runtime will invoke the webSocketClose() handler.
    ws.close(code, "StarbaseDB is closing WebSocket connection");

    // Remove the WebSocket connection from the map
    const tags = this.ctx.getTags(ws);
    if (tags.length) {
        const wsSessionId = tags[0];
        this.connections.delete(wsSessionId);
    }
}

Really all that is mandatory is the ws.close(code, ““) call to properly end the connection on the server. The code below is how we continue on to remove the connection from our Map object. If you recall earlier when we assigned our web socket connection a randomUUID() value and added it to our map, well that UUID value was assigning it as one of the sockets tags (which is an array of values). Here we just get the tags array from our web socket, capture the ID we assigned to it and remove it from our connections Map.

Frontend Implementation

This section we’re not going to dissect as much as we did our Workers implementation but I want to include code here so it’s a comprehensive example of how you get the entire stack talking harmoniously.

Below is an entire HTML document you can save as index.html on your computer and open up. To connect to your socket implementation you must replace the URL in the wss:// URL used in the code below, otherwise it won’t work.

The key component is calling our Worker to establish the connection, and then the world becomes our oyster to listen to events such as onopen, onmessage, onerror and onclose.

let socket = new WebSocket('wss://starbasedb.{YOUR-IDENTIFIER}.workers.dev/socket');

// ...

socket.onmessage = function(event) {
    logMessage("Received: " + event.data);
};

// ...

socket.send(JSON.stringify({
    sql: 'SELECT 1 + 1;',
    params: [],
    action: 'query'
}));

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Test</title>
    <style>
        #log {
            white-space: pre-wrap;
            max-height: 200px;
            overflow-y: scroll;
        }
    </style>
</head>
<body>
    <h1>WebSocket Test</h1>
    <div>
        <input type="text" id="messageInput" placeholder="Type a message..." />
        <button onclick="sendMessage()">Send Message</button>
    </div>
    <div id="log"></div>

    <script>
        let socket;
        const log = document.getElementById('log');

        function logMessage(message) {
            log.textContent += message + '\n';
            log.scrollTop = log.scrollHeight;
        }

        function connectWebSocket() {
            logMessage("Connecting to WebSocket...");

            socket = new WebSocket('wss://starbasedb.{YOUR-IDENTIFIER}.workers.dev/socket');

            socket.onopen = function() {
                logMessage("WebSocket connection opened.");
            };

            socket.onmessage = function(event) {
                logMessage("Received: " + event.data);
            };

            socket.onclose = function(event) {
                logMessage(`WebSocket closed with code: ${event.code}, reason: ${event.reason}, clean: ${event.wasClean}`);
            };

            socket.onerror = function(error) {
                logMessage("WebSocket error: " + error.message);
            };
        }

        function sendMessage() {
            const message = document.getElementById('messageInput').value;
            if (socket && socket.readyState === WebSocket.OPEN) {
                logMessage("Sending: " + message);

                socket.send(JSON.stringify({
                    sql: message,
                    params: [],
                    action: 'query'
                }));
            } else {
                logMessage("WebSocket is not open.");
            }
        }

        window.onload = connectWebSocket;
    </script>
</body>
</html>

Considerations

Where do web sockets REALLY shine? When you need to broadcast messages to many users when a particular event happens. For databases maybe you want everyone to know when someone has just signed up as a new user, or made a purchase on a website, or sent a new chat message to the chat.

A problem you will begin to experience as your usage scales is how do you handle so many concurrent connections at a time? That’s a story for a different day, but a consideration to make nonetheless.

When adding support for web sockets it’s generally a good idea to implement some idea of IP rate limiting so that clients can’t unnecessarily bombard your server with requests – because you can end up getting charged additionally for the number of messages that are passed back and forth. Thankfully with Cloudflare the limits to reach are pretty high before you get charged additionally.

Pull Requests

Want to see the code that was contributed alongside this article? Check the contributions out below!

https://github.com/Brayden/starbasedb/pull/15 (Support added)

Join the Adventure

We’re working on building an open-source database offering described with the basic building blocks we talked about above. When we get to join compute with storage the way we can with durable objects, we think there’s a lot more that can be offered to users than what some other providers are capable of.

Twitter: https://twitter.com/BraydenWilmoth

Github Repo: github.com/Brayden/starbasedb

Outerbase: https://outerbase.com