Merge pull request #37 from jcreek/35-replace-custom-websockets-with-partykit

35 replace custom websockets with partykit

README.md

@@ -4,11 +4,18 @@ A web app to allow scrum teams to easily perform estimations quickly and without
 There are two components to this:
 
 1. Frontend, built using SvelteKit
-2. WebSockets server, built using Node.js
+2. Realtime server, built using PartyKit (WebSocket-compatible)
 
 ## Frontend
 
-For full functionality, ensure there is a copy of the WebSockets server running. 
+For full functionality, ensure there is a PartyKit server running.
+
+### Environment variables
+
+The frontend uses PartyKit defaults unless overridden:
+
+- `PUBLIC_PARTYKIT_HOST`: PartyKit host (defaults to `localhost:1999` in dev, `websocket.jcreek.co.uk` in production)
+- `PUBLIC_PARTYKIT_PARTY`: Party name (defaults to `main`)
 
 ### Developing
 
@@ -44,12 +51,14 @@ This favicon was [generated](https://favicon.io/favicon-generator/) using the fo
 
 ### Developing
 
-Once you've installed dependencies with `npm install`, start a development server:
+Install dependencies in `src/server`, then run the PartyKit dev server:
 
 ```bash
-npm run dev
+cd src/server
+npm install
+npx partykit dev
 ```
 
 ### Deploying
 
-See the README in `src/server` for details.
+See the README in `src/server` for PartyKit deployment details.

package.json

@@ -29,6 +29,7 @@
 	"dependencies": {
 		"@fireworks-js/svelte": "^2.10.7",
 		"@joeattardi/emoji-button": "^4.6.4",
+		"partysocket": "1.0.2",
 		"uuid": "^9.0.0",
 		"ws": "^8.13.0"
 	}

src/routes/+page.svelte

@@ -1,10 +1,8 @@
 <script>
-	import { onMount } from 'svelte';
 	import { connectToWebSocket, sendMessage, generateId } from './estimation/estimation.js';
 	import { goto } from '$app/navigation';
 
 	let roomId = '';
-	let socket;
 	const cardSets = [
 		{ name: 'Fibonacci', values: [1, 2, 3, 5, 8, 13, 21, '?'], example: '(1, 2, 3, 5)' },
 		{
@@ -18,9 +16,39 @@
 
 	let selectedCardSet = cardSets[0];
 
-	function startARoom() {
+	function waitForSocketOpen(socket) {
+		if (socket.readyState === WebSocket.OPEN) {
+			return Promise.resolve();
+		}
+
+		return new Promise((resolve, reject) => {
+			const handleOpen = () => {
+				socket.removeEventListener('open', handleOpen);
+				socket.removeEventListener('error', handleError);
+				resolve();
+			};
+			const handleError = (event) => {
+				socket.removeEventListener('open', handleOpen);
+				socket.removeEventListener('error', handleError);
+				reject(event);
+			};
+
+			socket.addEventListener('open', handleOpen);
+			socket.addEventListener('error', handleError);
+		});
+	}
+
+	async function startARoom() {
 		const roomId = generateId();
-		sendMessage(socket, { roomId: roomId, type: 'create-room', cardSetName: selectedCardSet.name });
+		const socket = connectToWebSocket(roomId, onMessageReceived);
+		try {
+			await waitForSocketOpen(socket);
+			sendMessage(socket, { roomId: roomId, type: 'create-room', cardSetName: selectedCardSet.name });
+		} catch (error) {
+			console.error('Failed to connect to PartyKit:', error);
+			alert('Unable to create the room right now. Please try again.');
+			return;
+		}
 		goto(`/estimation/${roomId}`);
 	}
 
@@ -33,16 +61,6 @@
 	}
 
 	function onMessageReceived(message) {}
-
-	onMount(() => {
-		socket = connectToWebSocket(null, onMessageReceived);
-
-		setInterval(() => {
-			if (socket.readyState === WebSocket.OPEN) {
-				socket.send(JSON.stringify({ type: 'ping' }));
-			}
-		}, 45000); // send a ping every 45 seconds
-	});
 </script>
 
 <div class="container">

src/routes/estimation/estimation.js

@@ -1,18 +1,29 @@
 import { v4 as uuidv4 } from 'uuid';
+import PartySocket from 'partysocket';
 
 function generateId() {
 	const id = uuidv4();
 	return id;
 }
 
-function connectToWebSocket(roomId, onMessageReceived) {
-	let socket;
-    const isDevelopment = import.meta.env.DEV;
-	if (isDevelopment) {
-		socket = new WebSocket(`ws://localhost:8080`);
-	} else {
-		socket = new WebSocket(`wss://websocket.jcreek.co.uk`);
+const partyHost = import.meta.env.PUBLIC_PARTYKIT_HOST;
+const partyName = import.meta.env.PUBLIC_PARTYKIT_PARTY || 'main';
+
+function getPartyKitHost() {
+	if (partyHost) {
+		return partyHost;
 	}
+	return import.meta.env.DEV ? 'localhost:1999' : 'websocket.jcreek.co.uk';
+}
+
+function connectToWebSocket(roomId, onMessageReceived) {
+	const host = getPartyKitHost();
+	const socket = new PartySocket({
+		host,
+		room: roomId ?? 'lobby',
+		party: partyName,
+		...(import.meta.env.DEV ? { protocol: 'ws' } : {})
+	});
 
 	socket.addEventListener('open', () => {
 		console.log(`Connected to WebSocket server from roomId: ${roomId}`);

src/server/classes/Room.ts

@@ -1,5 +1,5 @@
-import { WebSocket } from 'ws';
-import { User } from './User';
+import type * as Party from 'partykit/server';
+import { User } from './User.js';
 
 class CardSet {
 	name: string;
@@ -27,7 +27,7 @@ export const cardSets: Array<CardSet> = [
 export class Room {
 	id: string;
 	cardSet: CardSet;
-	users: Map<WebSocket, User>;
+	users: Map<Party.Connection, User>;
 
 	constructor(id: string, cardSetName: string) {
 		this.id = id;
@@ -37,49 +37,48 @@ export class Room {
 		} else {
 			this.cardSet = cardSets[0];
 		}
-		
+
 		this.users = new Map();
 	}
 
-	addUser(ws: WebSocket, user: User) {
-		this.users.set(ws, user);
+	addUser(connection: Party.Connection, user: User) {
+		this.users.set(connection, user);
 	}
 
-	removeUser(ws: WebSocket) {
-		this.users.delete(ws);
+	removeUser(connection: Party.Connection) {
+		this.users.delete(connection);
 	}
 
 	getUsers() {
 		return Array.from(this.users.values());
 	}
 
-	getUserByWebSocket(ws: WebSocket) {
-		return this.users.get(ws);
+	getUserByConnection(connection: Party.Connection): User | undefined {
+		return this.users.get(connection);
 	}
 
-	getWebSocketByUser(user: User): WebSocket | null {
-		for (const [ws, u] of this.users) {
+	getConnectionByUser(user: User): Party.Connection | null {
+		for (const [connection, u] of this.users) {
 			if (u === user) {
-				return ws;
+				return connection;
 			}
 		}
 		return null;
 	}
 
-	getWebSocketByUserId(userId: String): WebSocket | null {
-		for (const [ws, u] of this.users) {
+	getConnectionByUserId(userId: string): Party.Connection | null {
+		for (const [connection, u] of this.users) {
 			if (u.userId === userId) {
-				return ws;
+				return connection;
 			}
 		}
 		return null;
 	}
 
 	broadcast(message: any) {
-		this.users.forEach((user, client) => {
-			if (client.readyState === WebSocket.OPEN) {
-				client.send(JSON.stringify(message));
-			}
+		const messageString = JSON.stringify(message);
+		this.users.forEach((user, connection) => {
+			connection.send(messageString);
 		});
 	}
 

src/server/package.json

@@ -23,6 +23,10 @@
     "map-stream": "0.0.7",
     "tsc-watch": "^5.0.3",
     "typescript": "^4.6.3",
-    "ws": "^8.12.1"
+    "ws": "^8.12.1",
+    "partysocket": "1.0.2"
+  },
+  "devDependencies": {
+    "partykit": "0.0.111"
   }
 }

src/server/party/index.ts

@@ -0,0 +1,31 @@
+import type * as Party from "partykit/server";
+
+export default class Server implements Party.Server {
+  constructor(readonly room: Party.Room) {}
+
+  onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    // A websocket just connected!
+    console.log(
+      `Connected:
+  id: ${conn.id}
+  room: ${this.room.id}
+  url: ${new URL(ctx.request.url).pathname}`
+    );
+
+    // let's send a message to the connection
+    conn.send("hello from server");
+  }
+
+  onMessage(message: string, sender: Party.Connection) {
+    // let's log the message
+    console.log(`connection ${sender.id} sent message: ${message}`);
+    // as well as broadcast it to all the other connections in the room...
+    this.room.broadcast(
+      `${sender.id}: ${message}`,
+      // ...except for the connection it came from
+      [sender.id]
+    );
+  }
+}
+
+Server satisfies Party.Worker;

src/server/partykit.json

@@ -0,0 +1,6 @@
+{
+  "$schema": "https://www.partykit.io/schema.json",
+  "name": "websocket-server-party",
+  "main": "websocket.ts",
+  "compatibilityDate": "2024-11-07"
+}

src/server/tsconfig.json

@@ -1,10 +1,11 @@
 {
-    "compilerOptions": {
-      "target": "ES2022",
-      "module": "ES2022",
-      "esModuleInterop": true,
-      "forceConsistentCasingInFileNames": true,
-      "strict": true,
-      "skipLibCheck": true
-    }
-  }
+	"compilerOptions": {
+		"target": "ES2022",
+		"module": "nodenext",
+		"moduleResolution": "nodenext",
+		"esModuleInterop": true,
+		"forceConsistentCasingInFileNames": true,
+		"strict": true,
+		"skipLibCheck": true
+	}
+}

src/server/websocket.ts

@@ -1,233 +1,239 @@
-import express from 'express';
-import { WebSocketServer, WebSocket } from 'ws';
-import { Room, cardSets } from './classes/Room.js';
+import type * as Party from 'partykit/server';
+import { Room } from './classes/Room.js';
 import { User } from './classes/User.js';
 import { JoinRoomMessage } from './classes/messages/JoinRoomMessage.js';
-import { ChangeEstimateMessage } from './classes/messages/ChangeEstimateMessage.js';
 import { SelectEstimateMessage } from './classes/messages/SelectEstimateMessage.js';
 
-const app = express();
-const port = process.env.PORT || 8080;
+type StoredRoom = {
+	id: string;
+	cardSetName: string;
+};
 
-function onSocketPreError(e: Error) {
-	console.log(e);
-}
+export default class EstimationParty implements Party.Server {
+	room: Room | null = null;
+	private nudgeTimeout: ReturnType<typeof setTimeout> | null = null;
 
-function onSocketPostError(e: Error) {
-	console.log(e);
-}
+	constructor(public party: Party.Room) {}
 
-function groupEstimates(users: Array<User>) {
-	let estimateGroups: { [key: number | string]: string[] } = {};
-	users.forEach((user) => {
-		if (user.estimate !== null) {
-			if (estimateGroups[user.estimate]) {
-				estimateGroups[user.estimate].push(user.name);
-			} else {
-				estimateGroups[user.estimate] = [user.name];
-			}
+	async onStart() {
+		const savedRoom = (await this.party.storage.get('room')) as StoredRoom | undefined;
+		if (savedRoom?.cardSetName) {
+			this.room = new Room(savedRoom.id ?? this.party.id, savedRoom.cardSetName);
 		}
-	});
-
-	return estimateGroups;
-}
-
-console.log(`Attempting to run server on port ${port}`);
-
-const server = app.listen(port, () => {
-	console.log(`Listening on port ${port}`);
-});
-
-const wss = new WebSocketServer({ noServer: true });
-const rooms = new Set<Room>();
-
-server.on('upgrade', (req, socket, head) => {
-	socket.on('error', onSocketPreError);
-
-	// perform auth
-	if (!!req.headers['BadAuth']) {
-		socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
-		socket.destroy();
-		return;
 	}
 
-	wss.handleUpgrade(req, socket, head, (ws) => {
-		socket.removeListener('error', onSocketPreError);
-		wss.emit('connection', ws, req);
-	});
-});
+	async onConnect(_connection: Party.Connection) {}
 
-wss.on('connection', (ws: WebSocket, req) => {
-	let room: Room | null = null;
+	async onClose(connection: Party.Connection) {
+		if (!this.room) {
+			return;
+		}
 
-	ws.on('error', onSocketPostError);
+		const user = this.room.getUserByConnection(connection);
+		if (!user) {
+			return;
+		}
 
-	function createRoom(roomId: string, cardSetName: string) {
-		room = new Room(roomId, cardSetName);
-		rooms.add(room);
+		this.room.removeUser(connection);
+		await this.updateRoomState();
+		this.party.broadcast(JSON.stringify({ type: 'user-left', userId: user.userId }));
+
+		if (this.room.getUsers().length === 0) {
+			this.clearNudgeTimeout();
+			await this.party.storage.delete('room');
+			this.room = null;
+		}
 	}
 
-	ws.on('message', (message) => {
-		const data = JSON.parse(message.toString());
+	async onMessage(message: string | ArrayBuffer | ArrayBufferView, connection: Party.Connection) {
+		if (typeof message !== 'string') {
+			console.log('Unsupported message payload type');
+			return;
+		}
+		const data = JSON.parse(message);
 
-		if (data.type === 'create-room') {
-			const roomId = data.roomId;
-			const cardSetName = data.cardSetName;
-			createRoom(roomId, cardSetName);
-		} else if (data.type === 'join-room') {
-			const joinRoomMessage: JoinRoomMessage = data as JoinRoomMessage;
-			const { roomId, userId, name } = joinRoomMessage;
-			room = getRoomById(roomId);
+		switch (data.type) {
+			case 'create-room':
+				this.room = new Room(this.party.id, data.cardSetName ?? '');
+				await this.updateRoomState();
+				break;
 
-			// Create room if it doesn't already exist
-			if (!room) {
-				createRoom(roomId, data.cardSetName);
-			}
+			case 'join-room':
+				{
+					this.room = await this.getOrCreateRoom(data.cardSetName);
+					if (!this.room) {
+						break;
+					}
+					const joinRoomMessage: JoinRoomMessage = data as JoinRoomMessage;
+					const { userId, name } = joinRoomMessage;
+					const user = new User(userId, name);
+					this.room.addUser(connection, user);
+					await this.updateRoomState();
+					this.party.broadcast(
+						JSON.stringify({ type: 'user-joined', message: `${name} has joined the room` })
+					);
+				}
+				break;
 
-			// Create the user in the room
-			const user = new User(userId, name);
-			room!.addUser(ws, user);
-
-			broadcastToRoom(roomId, { type: 'user-joined', message: `${name} has joined the room` });
-		} else if (data.type === 'select-estimate') {
-			const selectEstimateMessage: SelectEstimateMessage = data as SelectEstimateMessage;
-			const { userId, estimate } = selectEstimateMessage;
-
-			if (room) {
-				const user = room.getUserByWebSocket(ws);
-				if (user) {
+			case 'select-estimate':
+				{
+					this.room = await this.getOrCreateRoom();
+					if (!this.room) {
+						break;
+					}
+					const selectEstimateMessage: SelectEstimateMessage = data as SelectEstimateMessage;
+					const { userId, estimate } = selectEstimateMessage;
+					const user = this.room.getUserByConnection(connection);
+					if (!user) {
+						break;
+					}
 					user.estimate = estimate;
-					broadcastToRoom(room.id, {
-						type: 'estimate-selected',
-						userId,
-						name: user.name,
-						estimate
-					});
+					await this.updateRoomState();
+					this.party.broadcast(
+						JSON.stringify({
+							type: 'estimate-selected',
+							userId,
+							name: user.name,
+							estimate
+						})
+					);
 
-					const users = room.getUsers();
-					let timeout = undefined;
-
-					const allEstimates = room.getAllEstimates();
+					const users = this.room.getUsers();
+					const allEstimates = this.room.getAllEstimates();
 					if (allEstimates.size === users.length) {
-						// If all users have estimated, clear the timeout
-						clearTimeout(timeout);
-						timeout = undefined;
-
-						broadcastToRoom(room.id, {
-							type: 'estimation-closed',
-							groupedEstimates: groupEstimates(users)
-						});
+						this.clearNudgeTimeout();
+						this.party.broadcast(
+							JSON.stringify({
+								type: 'estimation-closed',
+								groupedEstimates: this.groupEstimates()
+							})
+						);
 					} else if (allEstimates.size === users.length - 1) {
-						// The last user has 30 seconds to add their estimation, otherwise send the nudge message
-						let notEstimatedUser: User | undefined = undefined;
-						for (let i = 0; i < users.length; i++) {
-							if (!allEstimates.has(users[i].name)) {
-								notEstimatedUser = users[i];
-								break;
-							}
-						}
-						if (notEstimatedUser !== undefined) {
-							const ws = room.getWebSocketByUser(notEstimatedUser);
-							if (ws) {
-								// Cancel the old timeout if it's still running
-								if (timeout) {
-									clearTimeout(timeout);
-								}
-								// Set a new timeout
-								timeout = setTimeout(() => {
-									if (!room!.getAllEstimates().has(notEstimatedUser!.name)) {
-										ws.send(JSON.stringify({ type: 'nudge' }));
-									}
-								}, 30000);
-							}
+						const notEstimatedUser = users.find((candidate) => !allEstimates.has(candidate.name));
+						if (notEstimatedUser) {
+							this.scheduleNudge(notEstimatedUser);
 						}
 					}
 				}
-			}
-		} else if (data.type === 'restart-estimation') {
-			if (room) {
-				room.clearEstimates();
-				broadcastToRoom(room.id, { type: 'estimation-restarted' });
-			}
-		} else if (data.type === 'get-user-estimates') {
-			if (room) {
-				const users = room.getUsers();
-				broadcastToRoom(room.id, { type: 'user-estimates', users, selectedCardSet: room.cardSet });
-			}
-		} else if (data.type === 'trigger-emoji') {
-			const { cardId, emoji } = data;
-			if (room) {
-				broadcastToRoom(room.id, {
-					type: 'trigger-emoji',
-					cardId,
-					emoji
-				});
-			}
-		} else if (data.type === 'ping') {
-			if (room) {
-				broadcastToRoom(room.id, 'pong');
-			}
-		} else if (data.type === 'kick-user') {
-			console.log('kick user message received');
-			if (room) {
-				const userIdToKick = data.userId;
+				break;
 
-				const ws = room.getWebSocketByUserId(userIdToKick);
-				console.log(userIdToKick, ws);
-				if (ws) {
-					console.log('closing ws');
-					ws.close();
+			case 'restart-estimation':
+				if (this.room) {
+					this.clearNudgeTimeout();
+					this.room.clearEstimates();
+					await this.updateRoomState();
+					this.party.broadcast(JSON.stringify({ type: 'estimation-restarted' }));
 				}
+				break;
 
-				// Remove them from the room even if the active websocket can't be found or can't be closed
-				room.removeUser(userIdToKick);
-			}
-		}
-	});
-
-	ws.on('close', () => {
-		if (room) {
-			const user = room.getUserByWebSocket(ws);
-			if (user) {
-				room.removeUser(ws);
-
-				if (room.getUsers().length === 0) {
-					rooms.delete(room);
-					broadcastToRoom(room.id, {
-						type: 'room-deleted',
-						message: 'Room deleted'
-					});
+			case 'get-user-estimates':
+				if (this.room) {
+					const users = this.room.getUsers();
+					connection.send(
+						JSON.stringify({ type: 'user-estimates', users, selectedCardSet: this.room.cardSet })
+					);
 				}
+				break;
 
-				broadcastToRoom(room.id, {
-					type: 'user-left',
-					userId: user.userId
-				});
+			case 'trigger-emoji':
+				if (this.room) {
+					const { cardId, emoji } = data;
+					this.party.broadcast(JSON.stringify({ type: 'trigger-emoji', cardId, emoji }));
+				}
+				break;
+
+			case 'ping':
+				connection.send(JSON.stringify({ type: 'pong' }));
+				break;
+
+			case 'kick-user':
+				if (this.room) {
+					const userIdToKick = data.userId;
+					const connectionToKick = this.room.getConnectionByUserId(userIdToKick);
+					if (connectionToKick) {
+						connectionToKick.close();
+						this.room.removeUser(connectionToKick);
+					}
+					await this.updateRoomState();
+					this.party.broadcast(JSON.stringify({ type: 'user-left', userId: userIdToKick }));
+				}
+				break;
+
+			default:
+				console.log('Unknown message type:', data.type);
+		}
+	}
+
+	// Helper to broadcast the current room state to all users in the room
+	private async updateRoomState() {
+		if (this.room) {
+			await this.party.storage.put('room', {
+				id: this.room.id,
+				cardSetName: this.room.cardSet.name
+			});
+		}
+	}
+
+	// Group user estimates into a map of estimates -> list of user names
+	private groupEstimates() {
+		const users = this.room ? this.room.getUsers() : [];
+		const estimateGroups: { [key: number | string]: string[] } = {};
+		users.forEach((user) => {
+			if (user.estimate !== null) {
+				if (estimateGroups[user.estimate]) {
+					estimateGroups[user.estimate].push(user.name);
+				} else {
+					estimateGroups[user.estimate] = [user.name];
+				}
 			}
+		});
+		return estimateGroups;
+	}
+
+	private async getOrCreateRoom(cardSetName?: string): Promise<Room | null> {
+		if (this.room) {
+			return this.room;
 		}
-	});
-});
 
-function broadcastToRoom(roomId: string, message: any) {
-	const room = getRoomById(roomId);
-	if (room) {
-		room.broadcast(message);
+		const savedRoom = (await this.party.storage.get('room')) as StoredRoom | undefined;
+		if (savedRoom?.cardSetName) {
+			this.room = new Room(savedRoom.id ?? this.party.id, savedRoom.cardSetName);
+			return this.room;
+		}
+
+		this.room = new Room(this.party.id, cardSetName ?? '');
+		await this.updateRoomState();
+		return this.room;
 	}
-}
 
-function getUsersInRoom(roomId: string) {
-	const room = getRoomById(roomId);
-	if (room) {
-		return room.getUsers();
+	private scheduleNudge(user: User) {
+		this.clearNudgeTimeout();
+
+		if (!this.room) {
+			return;
+		}
+
+		const connection = this.room.getConnectionByUser(user);
+		if (!connection) {
+			return;
+		}
+
+		this.nudgeTimeout = setTimeout(() => {
+			if (!this.room) {
+				return;
+			}
+
+			const allEstimates = this.room.getAllEstimates();
+			if (!allEstimates.has(user.name)) {
+				connection.send(JSON.stringify({ type: 'nudge' }));
+			}
+		}, 30000);
 	}
-	return [];
-}
 
-function getRoomById(roomId: string) {
-	for (const room of rooms) {
-		if (room.id === roomId) {
-			return room;
+	private clearNudgeTimeout() {
+		if (this.nudgeTimeout) {
+			clearTimeout(this.nudgeTimeout);
+			this.nudgeTimeout = null;
 		}
 	}
-	return null;
 }