Forem: Arnav Gosain

Implementing Live Visitor Count using Serverless WebSockets

Arnav Gosain — Mon, 05 Dec 2022 10:30:41 +0000

Originally posted on: https://arnavgosain.com/serverless-websockets-visitor-count

As a fun experiment, I wanted to add a live visitor count to my blog. I've seen this done in a few different ways, most commonly in a server-full way, but I wanted to try implementing it as a serverless function. Unlike other serverless solutions, Cloudflare Workers support WebSockets & provide a neat interface for consistent storage via Durable Objects, so I decided to give it a shot.

PS: If you're familiar with Cloudflare Workers, you can skip to Implementing the live visitors counter.

The Basics

Wrangler

Cloudflare provides a CLI tool called wrangler to control/deploy to any of their serverless offerings (Workers, D1, etc). Wranger can be installed using npm.

npm i -g wrangler

Scaffolding a Worker

To scaffold a new Worker, run wrangler generate and follow the prompts.

wrangler init my-worker

Wrangler will prompt you to create a package.json file, as well as whether you want to use TypeScript or not and if you want it to setup Jest/Vitest for testing. I recommend saying yes to all of these, as it will make your life easier later.

Wrangler will create the following files:

package.json - Self-explanatory.
wrangler.toml - Configuration file for the Wrangler
src/index.ts - The main entry point for the Worker
src/index.test.ts - Vitest test file for the Worker

Creating a basic WebSocket server with Cloudflare Workers

async function handleRequest(request: Request, env: Env) {
    const upgradeHeader = request.headers.get('Upgrade');

    // If the upgrade header is not set, or it's not set to "websocket", return 426
    if (!upgradeHeader || upgradeHeader !== 'websocket') {
        return new Response('Expected Upgrade: websocket', { status: 426 });
    }

    const [client, server] = new WebSocketPair();

    server.addEventListener('message', event => {
        // Handle messages from the client
    });

    server.addEventListener('close', () => {
        // Handle the client disconnecting
    });

    server.accept();

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

export default {
    fetch: handleRequest
};

On the client side, you can connect to the Worker using the WebSocket constructor.

const socket = new WebSocket('wss://my-worker.example.com');

socket.addEventListener('open', () => {
    console.log('Connected to the Worker');
});

socket.addEventListener('message', () => {
    console.log('Connected to the Worker');
});

socket.addEventListener('close', () => {
    console.log('Disconnected from the Worker');
});

Making a Counter using Durable Objects

On the surface level, a Durable Object is just a class. But an instance of a Durable Object class cannot be created manually, as it needs to be persisted and used in multiple requests.

That's where the Env object comes in. The Env object is passed to the request & the Durable Object constructor by the Workers runtime. It contains the special properties that let you access your Durable Object in the request handler.

For the runtime to know which Durable Object class to use, you need to register it in the wrangler.toml file.

[durable_objects]
bindings = [{ name = "counter", class_name = "Counter" }]

[[migrations]]
tag = "v1" # Should be unique for each entry
new_classes = ["Counter"]

After which it's accessible in the Env object as a DurableObjectNamespace.

interface Env {
    counter: DurableObjectNamespace;
}

export class Counter {
    constructor(private state: DurableObjectState, private env: Env) {}

    async decrement() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count - 1;
        await this.state.storage.put('count', newCount);
        return newCount;
    }

    async increment() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count + 1;
        await this.state.storage.put('count', newCount);
        return newCount;
    }
}

Hooking up our `Counter` Durable Object with the request handler

A Durable Object class's methods can't be called directly as we don't get access to the instance of the class itself. Instead, Workers runtime gives us to a DurableObjectStub, through which we can pass our request to the Durable Object by calling its fetch method.

So instead of calling counter.increment() directly, we can call await counter.fetch(server). After this, our persisted Durable Object instance will then handle our WebSocket connection.

interface Env {
    counter: DurableObjectNamespace;
}

export class Counter {
    // Code from the previous section
    // ...

    async fetch(request: Request) {
        const [client, server] = new WebSocketPair();

        server.addEventListener('message', async event => {
            // Messages are received/sent as strings, so we need to parse it into JSON
            // to use it as an object
            const action = JSON.parse(event.data);

            if (action.type === 'increment') {
                const newCount = await this.increment();
                server.send(JSON.stringify({ type: 'update/count', count: newCount }));
            } else if (action.type === 'decrement') {
                const newCount = await this.decrement();

                server.send(JSON.stringify({ type: 'update/count', count: newCount }));
            }
        });

        server.addEventListener('close', async () => {
            // When the client disconnects, we can delete all the data in Durable Object
            // Deleting all data automatically discards the Durable Object instance
            await this.state.storage.deleteAll();
        });

        server.accept();

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

async function handleRequest(request: Request, env: Env) {
    // Cloudflare provides the client IP in the 'CF-Connecting-IP' header
    const ip = request.headers.get('CF-Connecting-IP') as string;

    // We'll use the IP as the Durable Object's id
    const counterId = env.counter.idFromName(ip);
    // .get() will fetch the existing Durable Object instance, or create a new one if it doesn't exist
    const counter = env.counter.get(counterId);

    return await counter.fetch(request);
}

export default {
    fetch: handleRequest
};

Implementing the live visitors counter

Now that we have the basics out of the way, let's implement the live visitors counter.

So here's the rundown of how it works:

Each client connects to the Worker via a WebSocket
When the connection is established, we'll increment the count in our DurableObject by 1.
When the connection is closed, we'll decrement the count in our DurableObject by 1.
Whenever the count changes, the new count will be broadcasted to all connected clients.
When all clients disconnect, we'll delete all the data in the DurableObject, which will discard the DurableObject instance.

Keeping track of connections

Since our DurableObject's fetch method is called for every WebSocket connection, we can store the WebSocket connection in a Set.

export class Counter {
    // Refer to the previous section for the full code
    // ...

    private conns = new Set<WebSocket>();

    async fetch(request: Request) {
        const upgradeHeader = request.headers.get('Upgrade');

        // If the upgrade header is not set, or it's not set to "websocket", return 426
        if (!upgradeHeader || upgradeHeader !== 'websocket') {
            return new Response('Expected Upgrade: websocket', { status: 426 });
        }

        const [client, server] = new WebSocketPair();

        server.addEventListener('message', async event => {
            // ...
        });

        server.addEventListener('close', async () => {
            // Remove the session from the Set
            this.conns.delete(server);
        });

        server.accept();

        // Add the session to the Set
        this.conns.add(server);

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

Broadcasting a message to all connected clients

export class Counter {
    // Refer to the previous section for the full code
    // ...

    private conns = new Set<WebSocket>();

    async fetch(request: Request) {
        // ...
    }

    private broadcast(message: string) {
        for (const session of this.conns) {
            session.send(message);
        }
    }
}

Now we'll modify the `increment` and `decrement` methods to broadcast the new count to all connected clients

export class Counter {
    // Refer to the previous section for the full code
    // ...

    private conns = new Set<WebSocket>();

    private broadcast(message: string) {
        for (const conn of this.conns) {
            conn.send(message);
        }
    }

    async increment() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count + 1;
        await this.state.storage.put('count', newCount);

        // Broadcast the new count to all connected clients
        this.broadcast(JSON.stringify({ type: 'update/count', count: newCount }));

        return newCount;
    }

    async decrement() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count - 1;
        await this.state.storage.put('count', newCount);

        // Broadcast the new count to all connected clients
        this.broadcast(JSON.stringify({ type: 'update/count', count: newCount }));

        return newCount;
    }

    async fetch(request: Request) {
        /* ...*/
    }
}

Discarding closed connections

export class Counter {
    // Refer to the previous section for the full code
    // ...

    private conns = new Set<WebSocket>();

    private broadcast(message: string) {
        for (const conn of this.conns) {
            // Check if the connection is still alive
            try {
                conn.send(message);
            } catch {
                // If the connection is closed, remove it from the Set
                this.conns.delete(session);
            }
        }
    }

    async increment() {
        /* ... */
    }
    async decrement() {
        /* ... */
    }
    async fetch(request: Request) {
        /* ... */
    }
}

Full code

interface Env {
    counter: DurableObjectNamespace;
}

export class Counter {
    private conns = new Set<WebSocket>();

    constructor(private state: DurableObjectState, private env: Env) {}

    private broadcast(message: string) {
        for (const conn of this.conns) {
            // Check if the connection is still alive
            try {
                conn.send(message);
            } catch {
                // If the connection is closed, remove it from the Set
                this.conns.delete(session);
            }
        }
    }

    async increment() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count + 1;
        await this.state.storage.put('count', newCount);

        // Broadcast the new count to all connected clients
        this.broadcast(JSON.stringify({ type: 'update/count', count: newCount }));

        return newCount;
    }

    async decrement() {
        const count = (await this.state.storage.get('count')) ?? 0;
        const newCount = count - 1;
        await this.state.storage.put('count', newCount);

        // Broadcast the new count to all connected clients
        this.broadcast(JSON.stringify({ type: 'update/count', count: newCount }));

        return newCount;
    }

    async fetch(request: Request) {
        const [client, server] = new WebSocketPair();

        server.addEventListener('message', async event => {
            // Messages are received/sent as strings, so we need to parse it into JSON
            // to use it as an object
            const action = JSON.parse(event.data);

            if (action.type === 'increment') {
                const newCount = await this.increment();
                server.send(JSON.stringify({ type: 'update/count', count: newCount }));
            } else if (action.type === 'decrement') {
                const newCount = await this.decrement();

                server.send(JSON.stringify({ type: 'update/count', count: newCount }));
            }
        });

        server.addEventListener('close', async () => {
            // Remove the session from the Set
            this.conns.delete(server);

            if (this.conns.size === 0) {
                // When the client disconnects, we can delete all the data in Durable Object
                // Deleting all data automatically discards the Durable Object instance
                await this.state.storage.deleteAll();
            }
        });

        server.accept();

        // Add the session to the Set
        this.conns.add(server);

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

async function handleRequest(request: Request, env: Env) {
    const upgradeHeader = request.headers.get('Upgrade');

    // If the upgrade header is not set, or it's not set to "websocket", return 426
    if (!upgradeHeader || upgradeHeader !== 'websocket') {
        return new Response('Expected Upgrade: websocket', { status: 426 });
    }

    // Cloudflare provides the client IP in the 'CF-Connecting-IP' header
    const ip = request.headers.get('CF-Connecting-IP') as string;

    // We'll use the IP as the Durable Object's id
    const counterId = env.counter.idFromName(ip);
    // .get() will fetch the existing Durable Object instance, or create a new one if it doesn't exist
    const counter = env.counter.get(counterId);

    return await counter.fetch(request);
}

export default {
    fetch: handleRequest
};

How to use Litestream with SQLite & Prisma

Arnav Gosain — Tue, 31 May 2022 14:52:47 +0000

Originally posted on: https://arnavgosain.com/sqlite-prisma-litestream

One of my side projects is Rune, an alternative email client built for reading and annotating newsletter content.

Rune is built with Remix¹ and uses Prisma² as the ORM layer. Prisma supports quite a few databases, of which I prefer Postgres. But I wanted to keep things simple, and didn't want to set up a full-blown Postgres instance, so I decided to use SQLite.

Since since my preferred way to deploy is via Docker on Railway.app, I needed a way to persist SQLite since Docker containers are refreshed on every deployment. Litestream³ was the perfect tool for the job.

But since Litestream only works with SQLite's WAL journaling mode⁴, it needs to be enabled on the SQLite database before running the Remix app.

`enable-wal.js`

const { PrismaClient } = require("@prisma/client");

let client = new PrismaClient();

client.$queryRaw`PRAGMA journal_mode = WAL;`
  .then(() => {
    console.log("ENABLED WAL MODE FOR DATABASE");
  })
  .catch((err) => {
    console.log("DB SETUP FAILED", err);
    process.exit(1);
  });

`run.sh`

#!/bin/bash

set -ex

## Restores the database from S3
litestream restore -v -config /etc/litestream.yml -if-replica-exists -o /data/db /data/db

## Runs migrations on the restored database
npx prisma migrate deploy

## Enables WAL Mode
node ./scripts/enable-wal.js

npx concurrently "litestream replicate -config /etc/litestream.yml" "npm run start"

`Dockerfile`

CMD ["bash", "./scripts/run.sh"]

Explanation

Dockerfile runs the run.sh script.
The script first restores the latest snapshot of the database from S3 using Litestream.
It runs npx prisma migrate deploy to deploy any pending migrations.
Then, enable-wal.js is run to enable WAL journaling mode using PrismaClient's $queryRaw
Finally, the Remix app & Litestream are started concurrently using the concurrently⁵ CLI tool.

https://remix.run ↩
https://prisma.io ↩
https://litestream.io ↩
https://sqlite.org/wal.html ↩
https://www.npmjs.com/package/concurrently ↩

Improving API Documentation using React Query and TypeScript

Arnav Gosain — Sat, 02 Oct 2021 16:43:30 +0000

As your codebase grows, there is only one way to increase developer productivity: Documentation. One of many reasons I prefer TypeScript to JavaScript is that overtime, as your codebase grows, developer productivity increases because TypeScript (and typed languages in general) offer something that dynamically typed languages cannot, automatic documentation in your IDE.

This article assumes you're familiar with React Query. If you're not, I highly recommend you read the official docs and this intro guide by Sai Kranthi.

Why React Query

Imagine an simple app that does two things based on the PokeAPI:

Renders a list of Pokemons that link to their own dedicated page
Have dedicated pages for all pokemons

To fetch the list of pokemons, with Redux (before RTK Query) you would have to:

Create a global store
Create a reducer with an action to update the list in the store
Write a thunk action to fetch the data.
Write a useEffect hook inside to dispatch the thunk action.
Render the list.

And then you'd have to write invalidation logic, loading status logic and much more.

But with React Query, fetching your list of pokemons is as easy as wrapping your App in a QueryClientProvider and then making use of the useQuery and useMutation hooks.

Example of basic React Query usage:

This approach works for simple apps like a Pokemon List, but it quickly becomes unmanageable as you add more endpoints to your API. In which case, you would have to create many such custom hooks.

This is the problem I ran into as I hopped on my first project after joining TartanHQ. While it's a fairly simple CRUD app, it makes use of many endpoints and making custom hooks for each endpoint simply isn't an option.

One Hook For All Queries

To counteract this problem, we created a layer of abstraction over React Query's useQuery hook, a hook that makes use of TypeScript to improve discoverability of endpoints across the entire application.

import * as React from "react";
import {
  useQuery as useReactQuery,
  UseQueryOptions,
  UseQueryResult,
} from "react-query";
import { queryFetchers, QueryKeys } from "~/lib/api/queries";

type Await<T>  = T extends Promise<infer U> ? U : T;

export function useQuery<
  Key extends QueryKeys,
  Params = Parameters<typeof queryFetchers[Key]>,
  Data = Await<ReturnType<typeof queryFetchers[Key]>>
>(key: Key, options?: UseQueryOptions<Data>): UseQueryResult<Data>;

export function useQuery<
  Key extends QueryKeys,
  Params = Parameters<typeof queryFetchers[Key]>,
  Data = Await<ReturnType<typeof queryFetchers[Key]>>
>(
  key: Key,
  params: Params,
  options?: UseQueryOptions<Data>
): UseQueryResult<Data>;

export function useQuery<
  Key extends QueryKeys,
  Params = Parameters<typeof queryFetchers[Key]>,
  Data = Await<ReturnType<typeof queryFetchers[Key]>>
>(
  key: Key,
  arg2?: Params | UseQueryOptions<Data>,
  arg3?: UseQueryOptions<Data, unknown, Data>
) {
  const params = Array.isArray(arg2) ? arg2 : [];
  const options = !!arg3 && Array.isArray(arg2) ? arg3 : arg2;

  return useReactQuery(
    key,
    () => queryFetchers[key].apply(null, params),
    options
  );
}

/**
 * Legend:
 *
 * QKEY = Query Key
 * QData = Query Data
 */

const GET_ALL_POKEMONS_QKEY = "pokemons/all" as const;
type GetAllPokemonsQData = {
  count: number;
  next: string;
  previous: string;
  results: { name: string; url: string }[];
};
const getAllPokemons = (): Promise<GetAllPokemonsQData> => {
  return fetch("https://pokeapi.co/api/v2/pokemon?limit=151").then(
    (response) => response.json() as GetAllPokemonsQData
  );
};

const POKEMON_BY_ID_QKEY = "pokemons/byId" as const;
type GetPokemonByIdQData = Record<string, unknown>;
const getPokemonById = (id: string) => {
  return fetch(`https://pokeapi.co/api/v2/pokemon/${id}/`).then(
    (res) => res.json() as GetPokemonByIdQData
  );
};

export type QueryKeys = typeof GET_ALL_POKEMONS_KEY | typeof POKEMON_BY_ID_QKEY;
export const queryFetchers = {
  [GET_ALL_POKEMONS_QKEY]: getAllPokemons,
  [POKEMON_BY_ID_QKEY]: getPokemonById,
} as const;

Example:

Now that you're all done, you can take full advantage of VSCode autocomplete.

If you have an alternative idea or found this useful: I'd love to connect with you on Twitter!