Refactor socket tracker

server/routerlicious/packages/lambdas/src/nexus/connect.ts

@@ -270,16 +270,16 @@ async function connectOrderer(
 	};
 }
 
-function trackSocket(
+async function trackSocket(
 	socket: IWebSocket,
 	tenantId: string,
 	documentId: string,
 	claims: ITokenClaims,
 	{ socketTracker }: INexusLambdaDependencies,
-): void {
+): Promise<void> {
 	// Track socket and tokens for this connection
 	if (socketTracker && claims.jti) {
-		socketTracker.addSocketForToken(
+		return socketTracker.addSocketForToken(
 			createCompositeTokenId(tenantId, documentId, claims.jti),
 			socket,
 		);
@@ -712,7 +712,7 @@ export async function connectDocument(
 		(connectedMessage as any).timestamp = connectedTimestamp;
 		connectionTrace.stampStage(ConnectDocumentStage.MessageClientConnected);
 
-		trackSocket(socket, tenantId, documentId, claims, lambdaDependencies);
+		await trackSocket(socket, tenantId, documentId, claims, lambdaDependencies);
 		connectionTrace.stampStage(ConnectDocumentStage.SocketTrackerAppended);
 
 		trackCollaborationSession(

server/routerlicious/packages/routerlicious-base/package.json

@@ -82,6 +82,7 @@
 		"serialize-error": "^8.1.0",
 		"sha.js": "^2.4.12",
 		"sillyname": "^0.1.0",
+		"socket.io": "^4.8.0",
 		"uuid": "^11.1.0",
 		"winston": "^3.6.0",
 		"ws": "^7.5.10"

server/routerlicious/packages/routerlicious-base/src/nexus/customizations.ts

@@ -15,6 +15,7 @@ import type {
 } from "@fluidframework/server-services-core";
 import type { SocketIoAdapterCreator } from "@fluidframework/server-services-shared";
 import type { IRedisClientConnectionManager } from "@fluidframework/server-services-utils";
+import type { Server } from "socket.io";
 
 /**
  * @internal
@@ -36,4 +37,5 @@ export interface INexusResourcesCustomizations {
 	redisClientConnectionManagerForInvalidTokenCache?: IRedisClientConnectionManager;
 	customCreateSocketIoAdapter?: SocketIoAdapterCreator;
 	readinessCheck?: IReadinessCheck;
+	customSocketIoSetupFunction?: (io: Server) => void;
 }

server/routerlicious/packages/routerlicious-base/src/nexus/runnerFactory.ts

@@ -594,6 +594,7 @@ export class NexusResourcesFactory implements core.IResourcesFactory<NexusResour
 					httpServerConfig,
 					socketIoConfig,
 					customizations?.customCreateSocketIoAdapter,
+					customizations?.customSocketIoSetupFunction,
 			  );
 
 		const startupCheck = new StartupCheck();

server/routerlicious/packages/services-core/src/tokenRevocationManager.ts

@@ -179,11 +179,6 @@ export interface ITokenRevocationManager {
 		jwtId: string,
 		options?: IRevokeTokenOptions,
 	): Promise<ITokenRevocationResponse>;
-
-	/**
-	 * @deprecated move this function to IRevokedTokenChecker
-	 */
-	isTokenRevoked?(tenantId: string, documentId: string, jwtId: string): Promise<boolean>;
 }
 
 /**

server/routerlicious/packages/services-core/src/webSocketTracker.ts

@@ -12,7 +12,7 @@ import type { IWebSocket } from "./http";
  */
 export interface IWebSocketTracker {
 	// Add a token to socket mapping
-	addSocketForToken(compositeTokenId: string, webSocket: IWebSocket);
+	addSocketForToken(compositeTokenId: string, webSocket: IWebSocket): Promise<void>;
 
 	// Get the socket objects with the given token
 	getSocketsForToken(compositeTokenId: string): IWebSocket[];

server/routerlicious/packages/services-shared/src/webServer.ts

@@ -13,6 +13,7 @@ import * as util from "util";
 import type * as core from "@fluidframework/server-services-core";
 import { Lumberjack } from "@fluidframework/server-services-telemetry";
 import type { IRedisClientConnectionManager } from "@fluidframework/server-services-utils";
+import { Server } from "socket.io";
 import { setupMaster, setupWorker } from "@socket.io/sticky";
 
 import * as socketIo from "./socketIoServer";
@@ -102,6 +103,7 @@ export class SocketIoWebServerFactory implements core.IWebServerFactory {
 		private readonly httpServerConfig?: IHttpServerConfig,
 		private readonly socketIoConfig?: socketIo.ISocketIoServerConfig,
 		private readonly customCreateAdapter?: socketIo.SocketIoAdapterCreator,
+		private readonly customIoSetupFunction?: (io: Server) => void,
 	) {}
 
 	public create(requestListener: RequestListener): core.IWebServer {
@@ -115,7 +117,7 @@ export class SocketIoWebServerFactory implements core.IWebServerFactory {
 			server,
 			this.socketIoAdapterConfig,
 			this.socketIoConfig,
-			undefined /* ioSetup */,
+			this.customIoSetupFunction,
 			this.customCreateAdapter,
 		);
 

server/routerlicious/packages/services-utils/src/webSocketTracker.ts

@@ -16,31 +16,38 @@ export class WebSocketTracker implements IWebSocketTracker {
 	// Map of socketId to token ids. It assumes one socket could be used for connections with multiple tokens
 	private readonly socketIdToTokenIdMap: Map<string, Set<string>>;
 
-	constructor() {
+	constructor(private readonly useSocketIoRoomFeature: boolean = false) {
 		this.socketIdToSocketMap = new Map();
 		this.tokenIdToSocketIdMap = new Map();
 		this.socketIdToTokenIdMap = new Map();
 	}
 
-	public addSocketForToken(compositeTokenId: string, webSocket: IWebSocket) {
-		const socketIds = this.tokenIdToSocketIdMap.get(compositeTokenId);
-		if (socketIds) {
-			socketIds.add(webSocket.id);
+	public async addSocketForToken(compositeTokenId: string, webSocket: IWebSocket): Promise<void> {
+		if (this.useSocketIoRoomFeature) {
+			await webSocket.join(compositeTokenId);
 		} else {
-			this.tokenIdToSocketIdMap.set(compositeTokenId, new Set([webSocket.id]));
-		}
+			const socketIds = this.tokenIdToSocketIdMap.get(compositeTokenId);
+			if (socketIds) {
+				socketIds.add(webSocket.id);
+			} else {
+				this.tokenIdToSocketIdMap.set(compositeTokenId, new Set([webSocket.id]));
+			}
 
-		const tokenIds = this.socketIdToTokenIdMap.get(webSocket.id);
-		if (tokenIds) {
-			tokenIds.add(compositeTokenId);
-		} else {
-			this.socketIdToTokenIdMap.set(webSocket.id, new Set([compositeTokenId]));
-		}
+			const tokenIds = this.socketIdToTokenIdMap.get(webSocket.id);
+			if (tokenIds) {
+				tokenIds.add(compositeTokenId);
+			} else {
+				this.socketIdToTokenIdMap.set(webSocket.id, new Set([compositeTokenId]));
+			}
 
-		this.socketIdToSocketMap.set(webSocket.id, webSocket);
+			this.socketIdToSocketMap.set(webSocket.id, webSocket);
+		}
 	}
 
 	public getSocketsForToken(compositeTokenId: string): IWebSocket[] {
+		if (this.useSocketIoRoomFeature) {
+			throw new Error("Method not supported when socket room feature is enabled.");
+		}
 		const socketIds = this.tokenIdToSocketIdMap.get(compositeTokenId);
 
 		if (!socketIds) {
@@ -58,10 +65,17 @@ export class WebSocketTracker implements IWebSocketTracker {
 	}
 
 	public addSocket(webSocket: IWebSocket) {
+		if (this.useSocketIoRoomFeature) {
+			throw new Error("Method not supported when socket room feature is enabled.");
+		}
 		this.socketIdToSocketMap.set(webSocket.id, webSocket);
 	}
 
 	public removeSocket(socketId: string) {
+		if (this.useSocketIoRoomFeature) {
+			// No need to manually remove socket when socket room feature is enabled
+			return false;
+		}
 		const tokenIds = this.socketIdToTokenIdMap.get(socketId);
 
 		if (tokenIds) {
@@ -80,6 +94,9 @@ export class WebSocketTracker implements IWebSocketTracker {
 	}
 
 	public getAllSockets(): IWebSocket[] {
+		if (this.useSocketIoRoomFeature) {
+			throw new Error("Method not supported when socket room feature is enabled.");
+		}
 		return Array.from(this.socketIdToSocketMap.values());
 	}
 }