feat: add websocket server to the backend to publish vote count updates

Author: lucaazalim

backend/package.json

@@ -5,8 +5,6 @@
     "deploy": "wrangler deploy --minify",
     "cf-typegen": "wrangler types --env-interface CloudflareBindings",
     "db:generate": "drizzle-kit generate",
-    "db:migrate": "wrangler d1 migrations apply --env=production poll",
-    "db:migrate:local": "wrangler d1 migrations apply --local poll",
     "format": "prettier --write .",
     "format:check": "prettier --check ."
   },

backend/src/index.ts

@@ -130,6 +130,30 @@ const createContext = (env: Env, request: Request): Context => ({
 // Export default function to handle requests
 export default {
   async fetch(request: Request, env: Env): Promise<Response> {
+    const url = new URL(request.url);
+
+    // Handle WebSocket upgrade requests
+    if (url.pathname.endsWith("/websocket")) {
+      const upgradeHeader = request.headers.get("Upgrade");
+
+      if (!upgradeHeader || upgradeHeader !== "websocket") {
+        return new Response("Durable Object expected Upgrade: websocket", {
+          status: 426,
+        });
+      }
+
+      const pollId = url.searchParams.get("pollId");
+
+      if (!pollId) {
+        return new Response("Poll-Id header is required", { status: 400 });
+      }
+
+      const durableObjectId = env.POLL_DURABLE_OBJECT.idFromName(pollId);
+      const stub = env.POLL_DURABLE_OBJECT.get(durableObjectId);
+
+      return stub.fetch(request);
+    }
+
     // Handle CORS preflight requests
     if (request.method === "OPTIONS") {
       return new Response(null, {

backend/src/poll-durable-object.ts

@@ -6,17 +6,26 @@ import { Env } from ".";
 import migrations from "../drizzle/migrations";
 import { Result } from "./result";
 import { poll, pollOptions, votes } from "./schema";
-import { PollWithOptions, PollWithResults, Vote, VoteInsert } from "./types";
+import {
+  PollWithOptions,
+  PollWithResults,
+  Vote,
+  VoteInsert,
+  WebSocketMessage,
+} from "./types";
 
 export class PollDurableObject extends DurableObject {
   storage: DurableObjectStorage;
   db: DrizzleSqliteDODatabase<any>;
 
+  currentlyConnectedWebSockets: WebSocket[];
+
   constructor(ctx: DurableObjectState, env: Env) {
     super(ctx, env);
 
     this.storage = this.ctx.storage;
     this.db = drizzle(this.storage, { logger: false });
+    this.currentlyConnectedWebSockets = [];
 
     ctx.blockConcurrencyWhile(async () => {
       await this._migrate();
@@ -146,12 +155,47 @@ export class PollDurableObject extends DurableObject {
       .values(voteInserts)
       .returning();
 
+    this.currentlyConnectedWebSockets.forEach((ws) => {
+      if (ws.readyState === WebSocket.OPEN) {
+        ws.send(
+          JSON.stringify({
+            type: "newVote",
+            data: insertedVotes,
+          } satisfies WebSocketMessage<Vote[]>)
+        );
+      }
+    });
+
     return {
       success: true,
       data: insertedVotes,
     };
   }
 
+  async fetch(request: Request): Promise<Response> {
+    // Creates two ends of a WebSocket connection.
+    const webSocketPair = new WebSocketPair();
+    const [client, server] = Object.values(webSocketPair);
+
+    // Calling `accept()` tells the runtime that this WebSocket is to begin terminating
+    // request within the Durable Object. It has the effect of "accepting" the connection,
+    // and allowing the WebSocket to send and receive messages.
+    this.ctx.acceptWebSocket(server);
+    this.currentlyConnectedWebSockets.push(client);
+
+    // If the client closes the connection, the runtime will close the connection too.
+    server.addEventListener("close", (cls: CloseEvent) => {
+      this.currentlyConnectedWebSockets =
+        this.currentlyConnectedWebSockets.filter((ws) => ws !== client);
+      server.close(cls.code, "Durable Object is closing WebSocket");
+    });
+
+    return new Response(null, {
+      status: 101,
+      webSocket: client,
+    });
+  }
+
   async _migrate() {
     migrate(this.db, migrations);
   }

backend/src/types.ts

@@ -1,5 +1,10 @@
 import { poll, pollOptions, votes } from "./schema";
 
+export type WebSocketMessage<T> = {
+  type: "newVote";
+  data: T;
+};
+
 export type Poll = typeof poll.$inferSelect;
 export type PollInsert = typeof poll.$inferInsert;
 export type PollOption = typeof pollOptions.$inferSelect;