Web socket tracker change (#18930)

This PR includes changes:
- Move socket tracker to a separate file
- Change socket tracker to track all connected sockets
- Add unit test for web socket tracker

This change is required for cluster draining feature.

---------

Co-authored-by: Yunpeng Hou <[email protected]>
Co-authored-by: Yunpeng Hou <[email protected]>

Author: 3 people

server/routerlicious/packages/lambdas/src/alfred/index.ts

@@ -612,11 +612,15 @@ export function configureWebSocketServices(
 			(connectedMessage as any).timestamp = connectedTimestamp;
 
 			// Track socket and tokens for this connection
-			if (socketTracker && claims.jti) {
-				socketTracker.addSocketForToken(
-					core.createCompositeTokenId(message.tenantId, message.id, claims.jti),
-					socket,
-				);
+			if (socketTracker) {
+				if (claims.jti) {
+					socketTracker.addSocketForToken(
+						core.createCompositeTokenId(message.tenantId, message.id, claims.jti),
+						socket,
+					);
+				} else {
+					socketTracker.addSocket(socket);
+				}
 			}
 
 			// Set up listener to forward signal to clients in the collaboration session when the broadcast-signal endpoint is called

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

@@ -132,6 +132,12 @@
 			"ClassDeclaration_RiddlerResources": {
 				"forwardCompat": false,
 				"backCompat": false
+			},
+			"ClassDeclaration_AlfredResources": {
+				"forwardCompat": false
+			},
+			"InterfaceDeclaration_IAlfredResourcesCustomizations": {
+				"forwardCompat": false
 			}
 		}
 	}

server/routerlicious/packages/routerlicious-base/src/test/types/validateServerRouterliciousBasePrevious.generated.ts

@@ -31,6 +31,7 @@ declare function get_old_ClassDeclaration_AlfredResources():
 declare function use_current_ClassDeclaration_AlfredResources(
     use: TypeOnly<current.AlfredResources>): void;
 use_current_ClassDeclaration_AlfredResources(
+    // @ts-expect-error compatibility expected to be broken
     get_old_ClassDeclaration_AlfredResources());
 
 /*
@@ -199,6 +200,7 @@ declare function get_old_InterfaceDeclaration_IAlfredResourcesCustomizations():
 declare function use_current_InterfaceDeclaration_IAlfredResourcesCustomizations(
     use: TypeOnly<current.IAlfredResourcesCustomizations>): void;
 use_current_InterfaceDeclaration_IAlfredResourcesCustomizations(
+    // @ts-expect-error compatibility expected to be broken
     get_old_InterfaceDeclaration_IAlfredResourcesCustomizations());
 
 /*

server/routerlicious/packages/services-core/package.json

@@ -94,6 +94,9 @@
 			"InterfaceDeclaration_ICheckpointService": {
 				"forwardCompat": false,
 				"backCompat": false
+			},
+			"InterfaceDeclaration_IWebSocketTracker": {
+				"forwardCompat": false
 			}
 		}
 	}

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

@@ -154,7 +154,6 @@ export { TokenGenerator } from "./token";
 export { clientConnectivityStorageId, IUsageData, signalUsageStorageId } from "./usageData";
 export { IZookeeperClient, ZookeeperClientConstructor } from "./zookeeper";
 export {
-	IWebSocketTracker,
 	ITokenRevocationManager,
 	IRevokedTokenChecker,
 	ITokenRevocationResponse,
@@ -164,3 +163,4 @@ export {
 	createCompositeTokenId,
 } from "./tokenRevocationManager";
 export { IServiceMessageResourceManager } from "./serviceMessage";
+export { IWebSocketTracker } from "./webSocketTracker";

server/routerlicious/packages/services-core/src/test/types/validateServerServicesCorePrevious.generated.ts

@@ -2872,6 +2872,7 @@ declare function get_old_InterfaceDeclaration_IWebSocketTracker():
 declare function use_current_InterfaceDeclaration_IWebSocketTracker(
     use: TypeOnly<current.IWebSocketTracker>): void;
 use_current_InterfaceDeclaration_IWebSocketTracker(
+    // @ts-expect-error compatibility expected to be broken
     get_old_InterfaceDeclaration_IWebSocketTracker());
 
 /*

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

@@ -3,24 +3,6 @@
  * Licensed under the MIT License.
  */
 import { NetworkError, INetworkErrorDetails } from "@fluidframework/server-services-client";
-import { IWebSocket } from "./http";
-
-/**
- * Interface of web socket tracker
- * it tracks the mapping of web socket and token used to establish the socket connection
- * @internal
- */
-export interface IWebSocketTracker {
-	// Add a socket to internal map
-	addSocketForToken(compositeTokenId: string, webSocket: IWebSocket);
-
-	// Get socket objects from internal map
-	getSocketsForToken(compositeTokenId: string): IWebSocket[];
-
-	// Remove socket from tracking
-	// Return true if socket is removed, false if socket is not found
-	removeSocket(socketId: string): boolean;
-}
 
 /**
  * @internal

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

@@ -0,0 +1,28 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+import { IWebSocket } from "./http";
+
+/**
+ * Interface of web socket tracker
+ * it tracks the mapping of web socket and token used to establish the socket connection
+ * @internal
+ */
+export interface IWebSocketTracker {
+	// Add a token to socket mapping
+	addSocketForToken(compositeTokenId: string, webSocket: IWebSocket);
+
+	// Get the socket objects with the given token
+	getSocketsForToken(compositeTokenId: string): IWebSocket[];
+
+	// Add a socket to tracking
+	addSocket(webSocket: IWebSocket);
+
+	// Remove a socket from tracking
+	// Return true if socket is removed, false if socket is not found
+	removeSocket(socketId: string): boolean;
+
+	// Get all tracked socket objects
+	getAllSockets(): IWebSocket[];
+}

server/routerlicious/packages/services-utils/package.json

@@ -27,6 +27,7 @@
 		"prettier:fix": "prettier --write . --cache --ignore-path ../../.prettierignore",
 		"test": "mocha --recursive \"dist/test/*.spec.*js\"",
 		"test:coverage": "c8 npm test -- --reporter xunit --reporter-option output=nyc/junit-report.xml",
+		"test:debug": "mocha inspect --recursive \"dist/test/*.spec.*js\"",
 		"tsc": "tsc",
 		"typetests:gen": "fluid-type-test-generator",
 		"typetests:prepare": "flub typetests --dir . --reset --previous --normalize"
@@ -106,6 +107,10 @@
 		}
 	},
 	"typeValidation": {
-		"broken": {}
+		"broken": {
+			"ClassDeclaration_WebSocketTracker": {
+				"forwardCompat": false
+			}
+		}
 	}
 }

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

@@ -52,9 +52,6 @@ export {
 export { bindTimeoutContext } from "./timeoutContext";
 export { IThrottleConfig, ISimpleThrottleConfig, getThrottleConfig } from "./throttlerConfigs";
 export { IThrottleMiddlewareOptions, throttle } from "./throttlerMiddleware";
-export {
-	WebSocketTracker,
-	DummyTokenRevocationManager,
-	DummyRevokedTokenChecker,
-} from "./tokenRevocationManager";
+export { DummyTokenRevocationManager, DummyRevokedTokenChecker } from "./tokenRevocationManager";
 export { WinstonLumberjackEngine } from "./winstonLumberjackEngine";
+export { WebSocketTracker } from "./webSocketTracker";

server/routerlicious/packages/services-utils/src/test/types/validateServerServicesUtilsPrevious.generated.ts

@@ -319,6 +319,7 @@ declare function get_old_ClassDeclaration_WebSocketTracker():
 declare function use_current_ClassDeclaration_WebSocketTracker(
     use: TypeOnly<current.WebSocketTracker>): void;
 use_current_ClassDeclaration_WebSocketTracker(
+    // @ts-expect-error compatibility expected to be broken
     get_old_ClassDeclaration_WebSocketTracker());
 
 /*

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

@@ -0,0 +1,86 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import assert from "assert";
+import { WebSocketTracker } from "../webSocketTracker";
+import { IWebSocket } from "@fluidframework/server-services-core";
+
+describe("WebSocketTracker", () => {
+	// let webSocketTracker: WebSocketTracker;
+	let webSocket1: IWebSocket;
+	let webSocket2: IWebSocket;
+
+	beforeEach(() => {
+		webSocket1 = {
+			id: "socketId1",
+		} as IWebSocket;
+		webSocket2 = {
+			id: "socketId2",
+		} as IWebSocket;
+	});
+
+	it("should add and retrieve a socket for a token", () => {
+		let webSocketTracker: WebSocketTracker = new WebSocketTracker();
+		const compositeTokenId = "token1";
+		webSocketTracker.addSocketForToken(compositeTokenId, webSocket1);
+
+		const sockets = webSocketTracker.getSocketsForToken(compositeTokenId);
+		assert.strictEqual(sockets.length, 1), "Should have 1 socket";
+		assert.strictEqual(sockets[0].id, webSocket1.id, "Socket id should match");
+	});
+
+	it("should work for multiple calls", () => {
+		let webSocketTracker: WebSocketTracker = new WebSocketTracker();
+		let webSocket1: IWebSocket = {
+			id: "socketId1",
+		} as IWebSocket;
+		const compositeTokenId1 = "token1";
+
+		// Add 1st socket
+		webSocketTracker.addSocketForToken(compositeTokenId1, webSocket1);
+
+		let sockets = webSocketTracker.getSocketsForToken(compositeTokenId1);
+		assert.strictEqual(sockets.length, 1, "Should have 1 socket");
+		assert.strictEqual(sockets[0].id, webSocket1.id, "Should be socket id 1");
+
+		// Add 2nd socket
+		const compositeTokenId2 = "token2";
+		webSocketTracker.addSocketForToken(compositeTokenId2, webSocket2);
+
+		// Check results: socket 2 should exist
+		sockets = webSocketTracker.getSocketsForToken(compositeTokenId2);
+		assert.strictEqual(sockets.length, 1);
+		assert.strictEqual(sockets[0].id, webSocket2.id, "Should be socket id 2");
+
+		sockets = webSocketTracker.getAllSockets();
+		assert.strictEqual(sockets.length, 2, "Should have 2 sockets");
+
+		// Remove the socket
+		webSocketTracker.removeSocket(webSocket1.id);
+
+		// Check results
+		// socket 1 should not exist
+		sockets = webSocketTracker.getSocketsForToken(compositeTokenId1);
+		assert.strictEqual(sockets.length, 0, "Should be empty");
+
+		// socket 2 should exist
+		sockets = webSocketTracker.getSocketsForToken(compositeTokenId2);
+		assert.strictEqual(sockets.length, 1);
+		assert.strictEqual(sockets[0].id, webSocket2.id, "Should be socket id 2");
+
+		// Should only have 1 socket
+		sockets = webSocketTracker.getAllSockets();
+		assert.strictEqual(sockets.length, 1), "Should only have 1 socket after removal";
+	});
+
+	it("should add socket", () => {
+		let webSocketTracker: WebSocketTracker = new WebSocketTracker();
+		webSocketTracker.addSocket(webSocket1);
+
+		const sockets = webSocketTracker.getAllSockets();
+		assert.strictEqual(sockets.length, 1), "Should have 1 socket";
+		assert.strictEqual(sockets[0].id, webSocket1.id, "Socket id should match");
+	});
+});

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

@@ -3,85 +3,12 @@
  * Licensed under the MIT License.
  */
 import {
-	IWebSocket,
-	IWebSocketTracker,
 	ITokenRevocationManager,
 	IRevokedTokenChecker,
 	ITokenRevocationResponse,
 } from "@fluidframework/server-services-core";
 import { NetworkError } from "@fluidframework/server-services-client";
 
-/**
- * @internal
- */
-export class WebSocketTracker implements IWebSocketTracker {
-	// Map of socket id to socket object
-	private readonly socketIdToSocketMap: Map<string, IWebSocket>;
-	// Map of composite token id to socket ids. It assumes one token could be used for multiple sockets
-	private readonly tokenIdToSocketIdMap: Map<string, Set<string>>;
-	// 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() {
-		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);
-		} 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]));
-		}
-
-		this.socketIdToSocketMap.set(webSocket.id, webSocket);
-	}
-
-	public getSocketsForToken(compositeTokenId: string): IWebSocket[] {
-		const socketIds = this.tokenIdToSocketIdMap.get(compositeTokenId);
-
-		if (!socketIds) {
-			return [];
-		}
-
-		const socketResult: IWebSocket[] = [];
-		socketIds.forEach((socketId: string) => {
-			const socketObj = this.socketIdToSocketMap.get(socketId);
-			if (socketObj) {
-				socketResult.push(socketObj);
-			}
-		});
-		return socketResult;
-	}
-
-	public removeSocket(socketId: string) {
-		const tokenIds = this.socketIdToTokenIdMap.get(socketId);
-
-		if (tokenIds) {
-			tokenIds.forEach((tokenId: string) => {
-				const socketIds = this.tokenIdToSocketIdMap.get(tokenId);
-				if (socketIds) {
-					socketIds.delete(socketId);
-					if (socketIds.size <= 0) {
-						this.tokenIdToSocketIdMap.delete(tokenId);
-					}
-				}
-			});
-		}
-		this.socketIdToTokenIdMap.delete(socketId);
-		return this.socketIdToSocketMap.delete(socketId);
-	}
-}
-
 /**
  * @internal
  */