From 72e2fc8b1daa73c8e2e5a9923d6ab3461e81a9c8 Mon Sep 17 00:00:00 2001
From: Carson McManus <dyc3@users.noreply.github.com>
Date: Thu, 11 Jan 2024 21:31:30 -0500
Subject: [PATCH] balancer: fix clients not being removed from balancer context
 on disconnect (#1212)

* balancer: fix clients not being removed from balancer context on disconnect
possibly related to #1207?

* add a unit test
---
 crates/ott-balancer/src/balancer.rs | 85 +++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/crates/ott-balancer/src/balancer.rs b/crates/ott-balancer/src/balancer.rs
index ce4ce9dd0..24082c512 100644
--- a/crates/ott-balancer/src/balancer.rs
+++ b/crates/ott-balancer/src/balancer.rs
@@ -212,6 +212,7 @@ impl BalancerContext {
         let monolith = self.find_monolith_mut(client_id)?;
         monolith.remove_client(client_id);
         monolith.send(B2MLeave { client: client_id }).await?;
+        self.clients.remove(&client_id);
 
         Ok(())
     }
@@ -695,3 +696,87 @@ pub async fn dispatch_monolith_message(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod test {
+    use std::net::Ipv4Addr;
+
+    use crate::discovery::{HostOrIp, MonolithConnectionConfig};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_clients_add_remove() {
+        // a bunch of setup
+        let room_name = RoomName::from("test");
+        let mut ctx = BalancerContext::new();
+        let (monolith_outbound_tx, _monolith_outbound_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_outbound_tx = Arc::new(monolith_outbound_tx);
+        let (client_inbound_tx, _client_inbound_rx) = tokio::sync::mpsc::channel(100);
+        let (client_unicast_tx, _client_unicast_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_id = uuid::Uuid::new_v4().into();
+        let monolith = BalancerMonolith::new(
+            NewMonolith {
+                id: monolith_id,
+                region: "unknown".into(),
+                config: MonolithConnectionConfig {
+                    host: HostOrIp::Ip(Ipv4Addr::LOCALHOST.into()),
+                    port: 3002,
+                },
+                proxy_port: 3000,
+            },
+            monolith_outbound_tx,
+            client_inbound_tx,
+        );
+        let client_id = uuid::Uuid::new_v4().into();
+        let client = BalancerClient::new(
+            NewClient {
+                id: client_id,
+                room: room_name.clone(),
+                token: "test".into(),
+            },
+            client_unicast_tx,
+        );
+        ctx.add_monolith(monolith);
+        ctx.add_room(room_name.clone(), RoomLocator::new(monolith_id, 0))
+            .expect("failed to add room");
+
+        // add a client
+        ctx.add_client(client, monolith_id)
+            .await
+            .expect("failed to add client");
+
+        // make sure the client is in the context
+        assert!(ctx.clients.contains_key(&client_id));
+
+        // make sure the client is in the monolith
+        assert!(ctx
+            .monoliths
+            .get(&monolith_id)
+            .unwrap()
+            .rooms()
+            .get(&room_name)
+            .unwrap()
+            .clients()
+            .contains(&client_id));
+
+        // remove the client
+        ctx.remove_client(client_id)
+            .await
+            .expect("failed to remove client");
+
+        // make sure the client is not in the context
+        assert!(!ctx.clients.contains_key(&client_id));
+
+        // make sure the client is not in the monolith
+        assert!(!ctx
+            .monoliths
+            .get(&monolith_id)
+            .unwrap()
+            .rooms()
+            .get(&room_name)
+            .unwrap()
+            .clients()
+            .contains(&client_id));
+    }
+}