add client ids to system state (#1581)

* add client type * update tests * update sample system state * fix countRoomClients function * actually fix tests * add unit test * ran typeshare script * reapply client id change * update current state serving --------- Co-authored-by: Carson McManus <[email protected]>
dyc3 · Apr 3, 2024 · c15d8f4 · c15d8f4
1 parent 141c4e8
commit c15d8f4
Show file tree

Hide file tree

Showing 10 changed files with 375 additions and 54 deletions.
diff --git a/crates/ott-balancer-protocol/src/collector.rs b/crates/ott-balancer-protocol/src/collector.rs
@@ -23,7 +23,13 @@ pub struct MonolithState {
 #[typeshare]
 pub struct RoomState {
     pub name: RoomName,
-    pub clients: u32,
+    pub clients: Vec<ClientState>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[typeshare]
+pub struct ClientState {
+    pub id: ClientId,
 }
 
 pub enum Event {

diff --git a/crates/ott-balancer/src/balancer.rs b/crates/ott-balancer/src/balancer.rs
@@ -15,6 +15,7 @@ use tokio_tungstenite::tungstenite::protocol::CloseFrame;
 use tokio_tungstenite::tungstenite::Message;
 use tracing::{debug, error, info, instrument, trace, warn};
 
+use crate::balancer::collector::ClientState;
 use crate::client::ClientLink;
 use crate::config::BalancerConfig;
 use crate::connection::BALANCER_ID;
@@ -428,7 +429,11 @@ impl BalancerContext {
                     .iter()
                     .map(|(name, room)| RoomState {
                         name: name.clone(),
-                        clients: room.clients().len() as u32,
+                        clients: room
+                            .clients()
+                            .iter()
+                            .map(|c| ClientState { id: *c })
+                            .collect(),
                     })
                     .collect(),
             })

diff --git a/packages/ott-vis-panel/src/aggregate.spec.ts b/packages/ott-vis-panel/src/aggregate.spec.ts
@@ -16,19 +16,44 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 2 },
-					{ name: "bar", clients: 0 },
+					{
+						name: "foo",
+						clients: [
+							{
+								id: "e7229053-89df-428d-a37c-4b669fd57788",
+							},
+							{
+								id: "3fc0f726-2ad7-438b-8b2c-bae675dc1178",
+							},
+						],
+					},
+					{ name: "bar", clients: [] },
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 3 }],
+				rooms: [
+					{
+						name: "baz",
+						clients: [
+							{
+								id: "a90a98eb-5c82-44b3-90e0-1d117a9444c4",
+							},
+							{
+								id: "a7a40762-0308-408a-b954-d3f7dc2e5732",
+							},
+							{
+								id: "0ef93318-4b39-4b56-9180-637a9abeae9e",
+							},
+						],
+					},
+				],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 0 }],
+				rooms: [{ name: "qux", clients: [] }],
 			},
 			{
 				id: "f21df607-b572-4bdd-aa2f-3fead21bba86",
@@ -45,19 +70,32 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 1 },
-					{ name: "bar", clients: 2 },
+					{
+						name: "foo",
+						clients: [
+							{
+								id: "e7229053-89df-428d-a37c-4b669fd57788",
+							},
+						],
+					},
+					{
+						name: "bar",
+						clients: [
+							{ id: "4a6fe051-3247-4cad-860a-cb455ee65923" },
+							{ id: "33bbcd19-2af5-4244-9d71-cb647acc1b06" },
+						],
+					},
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 0 }],
+				rooms: [{ name: "baz", clients: [] }],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 0 }],
+				rooms: [{ name: "qux", clients: [] }],
 			},
 			{
 				id: "f21df607-b572-4bdd-aa2f-3fead21bba86",
@@ -74,19 +112,29 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 0 },
-					{ name: "bar", clients: 0 },
+					{ name: "foo", clients: [] },
+					{ name: "bar", clients: [] },
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 0 }],
+				rooms: [{ name: "baz", clients: [] }],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 4 }],
+				rooms: [
+					{
+						name: "qux",
+						clients: [
+							{ id: "d3be3464-efd5-41a1-b145-7d54378b02e3" },
+							{ id: "acc449cc-4748-435d-96b8-63530beac3d8" },
+							{ id: "9d2ff554-8388-4021-8467-5dfb208bd66e" },
+							{ id: "ff68188f-f739-46df-9bd7-dd25c1026651" },
+						],
+					},
+				],
 			},
 			{
 				id: "f21df607-b572-4bdd-aa2f-3fead21bba86",
@@ -125,34 +173,109 @@ describe("aggregation helpers", () => {
 
 	it("dedupes rooms", () => {
 		const rooms = [
-			{ name: "foo", clients: 1 },
-			{ name: "bar", clients: 2 },
-			{ name: "foo", clients: 1 },
+			{ name: "foo", clients: [{ id: "ff0ac5e0-caa8-4d5f-aba1-0c4aaa2d6f9e" }] },
+			{
+				name: "bar",
+				clients: [
+					{ id: "e36d4eb5-f526-4566-b94f-8cfc6dbf8548" },
+					{ id: "e842eeef-ef6c-4095-acc7-6342fb8c8b8c" },
+				],
+			},
+			{ name: "foo", clients: [{ id: "f7d5d57f-d15f-48b0-b30c-9bb378ce4943" }] },
 		];
 		expect(dedupeRooms(rooms)).toEqual([
-			{ name: "foo", clients: 2 },
-			{ name: "bar", clients: 2 },
+			{
+				name: "foo",
+				clients: [
+					{ id: "ff0ac5e0-caa8-4d5f-aba1-0c4aaa2d6f9e" },
+					{ id: "f7d5d57f-d15f-48b0-b30c-9bb378ce4943" },
+				],
+			},
+			{
+				name: "bar",
+				clients: [
+					{ id: "e36d4eb5-f526-4566-b94f-8cfc6dbf8548" },
+					{ id: "e842eeef-ef6c-4095-acc7-6342fb8c8b8c" },
+				],
+			},
 		]);
 	});
 
 	it("dedupes rooms using sample data", () => {
 		const rooms = sampleSystemState.flatMap(b => b.monoliths.flatMap(m => m.rooms));
 		expect(dedupeRooms(rooms)).toEqual([
-			{ name: "foo", clients: 3 },
-			{ name: "bar", clients: 2 },
-			{ name: "baz", clients: 3 },
-			{ name: "qux", clients: 4 },
+			{
+				name: "foo",
+				clients: [
+					{ id: "e7229053-89df-428d-a37c-4b669fd57788" },
+					{ id: "3fc0f726-2ad7-438b-8b2c-bae675dc1178" },
+					{ id: "e7229053-89df-428d-a37c-4b669fd57788" },
+				],
+			},
+			{
+				name: "bar",
+				clients: [
+					{ id: "4a6fe051-3247-4cad-860a-cb455ee65923" },
+					{ id: "33bbcd19-2af5-4244-9d71-cb647acc1b06" },
+				],
+			},
+			{
+				name: "baz",
+				clients: [
+					{ id: "a90a98eb-5c82-44b3-90e0-1d117a9444c4" },
+					{ id: "a7a40762-0308-408a-b954-d3f7dc2e5732" },
+					{ id: "0ef93318-4b39-4b56-9180-637a9abeae9e" },
+				],
+			},
+			{
+				name: "qux",
+				clients: [
+					{ id: "d3be3464-efd5-41a1-b145-7d54378b02e3" },
+					{ id: "acc449cc-4748-435d-96b8-63530beac3d8" },
+					{ id: "9d2ff554-8388-4021-8467-5dfb208bd66e" },
+					{ id: "ff68188f-f739-46df-9bd7-dd25c1026651" },
+				],
+			},
 		]);
 	});
 
 	it("dedupes monoliths", () => {
 		const monoliths = [
-			{ id: "a", region: "x", rooms: [{ name: "foo", clients: 2 }] },
+			{
+				id: "a",
+				region: "x",
+				rooms: [
+					{
+						name: "foo",
+						clients: [
+							{ id: "b379bce7-bd7a-4d79-a6bd-010e4fba1789" },
+							{ id: "a4505c5f-4856-49af-be53-77cabcb13aad" },
+						],
+					},
+				],
+			},
 			{ id: "b", region: "x", rooms: [] },
-			{ id: "a", region: "x", rooms: [{ name: "foo", clients: 1 }] },
+			{
+				id: "a",
+				region: "x",
+				rooms: [{ name: "foo", clients: [{ id: "379fdf91-e1e5-47b3-ac0c-0380a51c3479" }] }],
+			},
 		];
 		expect(dedupeMonoliths(monoliths)).toEqual([
-			{ id: "a", region: "x", rooms: [{ name: "foo", clients: 3 }] },
+			{
+				id: "a",
+				region: "x",
+				rooms: [
+					{
+						name: "foo",
+						clients: [
+							{ id: "b379bce7-bd7a-4d79-a6bd-010e4fba1789" },
+							{ id: "a4505c5f-4856-49af-be53-77cabcb13aad" },
+							{ id: "379fdf91-e1e5-47b3-ac0c-0380a51c3479" },
+						],
+					},
+				],
+			},
 			{ id: "b", region: "x", rooms: [] },
 		]);
 	});

diff --git a/packages/ott-vis-panel/src/aggregate.ts b/packages/ott-vis-panel/src/aggregate.ts
@@ -9,7 +9,7 @@ export function countRoomClients(state: SystemState): Record<string, number> {
 	for (const balancer of state) {
 		for (const monolith of balancer.monoliths) {
 			for (const room of monolith.rooms) {
-				roomClients[room.name] = (roomClients[room.name] ?? 0) + room.clients;
+				roomClients[room.name] = (roomClients[room.name] ?? 0) + room.clients.length;
 			}
 		}
 	}
@@ -60,7 +60,7 @@ function reduceRoom(rA: Room, rB: Room): Room {
 	// FIXME: (perf) This is a potentially hot path, and we should avoid creating a new object here.
 	return {
 		name: rA.name,
-		clients: rA.clients + rB.clients,
+		clients: [...rA.clients, ...rB.clients],
 	};
 }
 

diff --git a/packages/ott-vis-panel/src/components/CorePanel.tsx b/packages/ott-vis-panel/src/components/CorePanel.tsx
@@ -160,19 +160,34 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 2 },
-					{ name: "bar", clients: 0 },
+					{
+						name: "foo",
+						clients: [
+							{ id: "caa15370-8861-459e-997d-3e97f08f37d0" },
+							{ id: "12726a0c-02de-49f0-ab59-d87baf9c289f" },
+						],
+					},
+					{ name: "bar", clients: [] },
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 3 }],
+				rooms: [
+					{
+						name: "baz",
+						clients: [
+							{ id: "f3207419-b1d6-4c55-bc9d-799b4d1a70d7" },
+							{ id: "0ecd5456-ba1f-4585-b64e-e76d2c515c17" },
+							{ id: "6ed66113-cbd4-46ec-8b56-dae6e80d4f31" },
+						],
+					},
+				],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 0 }],
+				rooms: [{ name: "qux", clients: [] }],
 			},
 		],
 	},
@@ -184,19 +199,19 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 1 },
-					{ name: "bar", clients: 2 },
+					{ name: "foo", clients: [{ id: "4ac25d42-b0d3-49ff-9c43-cf98e1fde1d8" }] },
+					{ name: "bar", clients: [{ id: "f2e74aa2-8dbe-44bc-a2ef-75d201bb7387" }] },
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 0 }],
+				rooms: [{ name: "baz", clients: [] }],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 0 }],
+				rooms: [{ name: "qux", clients: [] }],
 			},
 		],
 	},
@@ -208,19 +223,29 @@ const sampleSystemState: SystemState = [
 				id: "2bd5e4a7-14f6-4da4-bedd-72946864a7bf",
 				region: "ewr",
 				rooms: [
-					{ name: "foo", clients: 0 },
-					{ name: "bar", clients: 0 },
+					{ name: "foo", clients: [] },
+					{ name: "bar", clients: [] },
 				],
 			},
 			{
 				id: "419580cb-f576-4314-8162-45340c94bae1",
 				region: "ewr",
-				rooms: [{ name: "baz", clients: 0 }],
+				rooms: [{ name: "baz", clients: [] }],
 			},
 			{
 				id: "0c85b46e-d343-46a3-ae4f-5f2aa1a8bdac",
 				region: "cdg",
-				rooms: [{ name: "qux", clients: 4 }],
+				rooms: [
+					{
+						name: "qux",
+						clients: [
+							{ id: "fe768adf-730a-4cc5-a7e3-2c3438a538c6" },
+							{ id: "c4a6362a-61c8-45dd-9913-49f1bbccaeb9" },
+							{ id: "3e3389ff-d3c2-4814-8089-44cc7ec01eb4" },
+							{ id: "ddf9309a-8ace-4c53-9dd8-f742e9f282c3" },
+						],
+					},
+				],
 			},
 		],
 	},