abstract balancing logic (#1243)

* add new trait & struct

* minor changes to struct & trait

* more changes to trait

* attempt to add type that implements trait to BalancerContext

* add traits to MonolithSelection

* manually implement default for BalancerContext

* fixed errors and changed function logic

* moved trait to a new module

* moved impl to new mod and fixed imports

* created filter_monoliths

* Move filtering logic into dedicated function

* implemented new functions

* implemented new methods properly

* added snippet for unit test

* Made some changes to unit test

* changed unit test

Author: Victor-M-Giraldo

crates/ott-balancer/src/balancer.rs

@@ -4,7 +4,6 @@ use ott_balancer_protocol::monolith::{
     B2MClientMsg, B2MJoin, B2MLeave, B2MUnload, MsgB2M, MsgM2B, RoomMetadata,
 };
 use ott_balancer_protocol::*;
-use rand::seq::IteratorRandom;
 use serde_json::value::RawValue;
 use tokio::sync::RwLock;
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
@@ -16,6 +15,7 @@ use crate::client::ClientLink;
 use crate::config::BalancerConfig;
 use crate::monolith::Room;
 use crate::room::RoomLocator;
+use crate::selection::{MinRoomsSelector, MonolithSelection};
 use crate::{
     client::{BalancerClient, NewClient},
     messages::*,
@@ -171,12 +171,25 @@ impl BalancerLink {
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct BalancerContext {
     pub clients: HashMap<ClientId, BalancerClient>,
     pub monoliths: HashMap<MonolithId, BalancerMonolith>,
     pub rooms_to_monoliths: HashMap<RoomName, RoomLocator>,
     pub monoliths_by_region: HashMap<String, Vec<MonolithId>>,
+    pub monolith_selection: Box<dyn MonolithSelection + Send + Sync + 'static>,
+}
+
+impl Default for BalancerContext {
+    fn default() -> Self {
+        BalancerContext {
+            clients: HashMap::default(),
+            monoliths: HashMap::default(),
+            rooms_to_monoliths: HashMap::default(),
+            monoliths_by_region: HashMap::default(),
+            monolith_selection: Box::<MinRoomsSelector>::default(),
+        }
+    }
 }
 
 impl BalancerContext {
@@ -356,50 +369,29 @@ impl BalancerContext {
         Ok(monolith)
     }
 
-    /// When loading a room, call this to select the best monolith to load it on.
-    pub fn select_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
-        fn cmp(x: &BalancerMonolith, y: &BalancerMonolith) -> std::cmp::Ordering {
-            x.rooms().len().cmp(&y.rooms().len())
-        }
-
+    /// Prioritizes monoliths in the same region
+    pub fn filter_monoliths(&self) -> Vec<&BalancerMonolith> {
         let in_region = self
             .monoliths_by_region
             .get(BalancerConfig::get().region.as_str());
         if let Some(in_region) = in_region {
-            let selected = in_region
+            return in_region
                 .iter()
                 .flat_map(|id| self.monoliths.get(id))
-                .min_by(|x, y| cmp(x, y));
-            if let Some(s) = selected {
-                return Ok(s);
-            }
-        }
-        let selected = self.monoliths.values().min_by(|x, y| cmp(x, y));
-        match selected {
-            Some(s) => Ok(s),
-            None => anyhow::bail!("no monoliths available"),
+                .collect();
         }
+
+        self.monoliths.values().collect()
     }
 
-    pub fn random_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
-        let in_region = self
-            .monoliths_by_region
-            .get(BalancerConfig::get().region.as_str());
-        if let Some(in_region) = in_region {
-            let selected = in_region.iter().choose(&mut rand::thread_rng());
-            if let Some(s) = selected {
-                if let Some(m) = self.monoliths.get(s) {
-                    return Ok(m);
-                }
-            }
-        }
+    pub fn select_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
+        let filtered = self.filter_monoliths();
+        return self.monolith_selection.select_monolith(filtered);
+    }
 
-        let selected = self
-            .monoliths
-            .values()
-            .choose(&mut rand::thread_rng())
-            .ok_or(anyhow::anyhow!("no monoliths available"))?;
-        Ok(selected)
+    pub fn random_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
+        let filtered = self.filter_monoliths();
+        return self.monolith_selection.random_monolith(filtered);
     }
 
     #[instrument(skip(self, monolith), err, fields(monolith_id = %monolith))]

crates/ott-balancer/src/lib.rs

@@ -24,6 +24,7 @@ pub mod discovery;
 pub mod messages;
 pub mod monolith;
 pub mod room;
+pub mod selection;
 pub mod service;
 
 pub async fn run() -> anyhow::Result<()> {

crates/ott-balancer/src/selection.rs

@@ -0,0 +1,120 @@
+use rand::seq::IteratorRandom;
+
+use crate::monolith::BalancerMonolith;
+
+#[derive(Debug, Default)]
+pub struct MinRoomsSelector;
+pub trait MonolithSelection: std::fmt::Debug {
+    fn select_monolith<'a>(
+        &'a self,
+        monolith: Vec<&'a BalancerMonolith>,
+    ) -> anyhow::Result<&BalancerMonolith>;
+
+    fn random_monolith<'a>(
+        &'a self,
+        monolith: Vec<&'a BalancerMonolith>,
+    ) -> anyhow::Result<&BalancerMonolith>;
+}
+
+impl MonolithSelection for MinRoomsSelector {
+    fn select_monolith<'a>(
+        &'a self,
+        monolith: Vec<&'a BalancerMonolith>,
+    ) -> anyhow::Result<&BalancerMonolith> {
+        fn cmp(x: &BalancerMonolith, y: &BalancerMonolith) -> std::cmp::Ordering {
+            x.rooms().len().cmp(&y.rooms().len())
+        }
+
+        let selected = monolith.iter().min_by(|x, y| cmp(x, y));
+        match selected {
+            Some(s) => Ok(s),
+            None => anyhow::bail!("no monoliths available"),
+        }
+    }
+
+    fn random_monolith<'a>(
+        &'a self,
+        monolith: Vec<&'a BalancerMonolith>,
+    ) -> anyhow::Result<&BalancerMonolith> {
+        let selected = monolith
+            .iter()
+            .choose(&mut rand::thread_rng())
+            .ok_or_else(|| anyhow::anyhow!("no monoliths available"))?;
+        Ok(selected)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::net::Ipv4Addr;
+    use std::sync::Arc;
+
+    use crate::discovery::{HostOrIp, MonolithConnectionConfig};
+    use crate::monolith::{BalancerMonolith, NewMonolith};
+    use ott_balancer_protocol::*;
+
+    use super::{MinRoomsSelector, MonolithSelection};
+
+    #[tokio::test]
+    async fn test_min_by() {
+        let room_one = RoomName::from("room one");
+        let room_two = RoomName::from("room two");
+        let room_three = RoomName::from("room three");
+        let (monolith_outbound_tx, _monolith_outbound_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_outbound_tx_one = Arc::new(monolith_outbound_tx);
+        let (client_inbound_tx_one, _client_inbound_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_id_one = uuid::Uuid::new_v4().into();
+
+        let mut monolith_one = BalancerMonolith::new(
+            NewMonolith {
+                id: monolith_id_one,
+                region: "unknown".into(),
+                config: MonolithConnectionConfig {
+                    host: HostOrIp::Ip(Ipv4Addr::LOCALHOST.into()),
+                    port: 3002,
+                },
+                proxy_port: 3000,
+            },
+            monolith_outbound_tx_one,
+            client_inbound_tx_one,
+        );
+
+        monolith_one
+            .add_room(&room_one)
+            .expect("failed to add room");
+        monolith_one
+            .add_room(&room_two)
+            .expect("failed to add room");
+
+        let (monolith_outbound_tx, _monolith_outbound_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_outbound_tx_two = Arc::new(monolith_outbound_tx);
+        let (client_inbound_tx_two, _client_inbound_rx) = tokio::sync::mpsc::channel(100);
+        let monolith_id_two = uuid::Uuid::new_v4().into();
+
+        let mut monolith_two = BalancerMonolith::new(
+            NewMonolith {
+                id: monolith_id_two,
+                region: "unknown".into(),
+                config: MonolithConnectionConfig {
+                    host: HostOrIp::Ip(Ipv4Addr::LOCALHOST.into()),
+                    port: 3002,
+                },
+                proxy_port: 3000,
+            },
+            monolith_outbound_tx_two,
+            client_inbound_tx_two,
+        );
+
+        monolith_two
+            .add_room(&room_three)
+            .expect("failed to add room");
+
+        let monoliths: Vec<&BalancerMonolith> = vec![&monolith_one, &monolith_two];
+
+        let selected = MinRoomsSelector
+            .select_monolith(monoliths)
+            .expect("failed to select monolith");
+
+        assert_eq!(selected.id(), monolith_two.id())
+    }
+}