balancer: add room arg to monolith selection

Author: dyc3

crates/ott-balancer/src/balancer.rs

@@ -386,9 +386,9 @@ impl BalancerContext {
         self.monoliths.values().collect()
     }
 
-    pub fn select_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
+    pub fn select_monolith(&self, room: &RoomName) -> anyhow::Result<&BalancerMonolith> {
         let filtered = self.filter_monoliths();
-        return self.monolith_selection.select_monolith(filtered);
+        return self.monolith_selection.select_monolith(room, filtered);
     }
 
     pub fn random_monolith(&self) -> anyhow::Result<&BalancerMonolith> {
@@ -453,7 +453,7 @@ pub async fn join_client(
         }
         None => {
             // the room is not loaded, randomly select a monolith
-            let selected = ctx_write.select_monolith()?;
+            let selected = ctx_write.select_monolith(&new_client.room)?;
             debug!(
                 "room is not loaded, selected monolith: {:?} (region: {:?})",
                 selected.id(),

crates/ott-balancer/src/selection.rs

@@ -1,12 +1,14 @@
 use crate::monolith::BalancerMonolith;
 use enum_dispatch::enum_dispatch;
+use ott_balancer_protocol::RoomName;
 use rand::seq::IteratorRandom;
 use serde::Deserialize;
 
 #[enum_dispatch(MonolithSelectionStrategy)]
 pub trait MonolithSelection: std::fmt::Debug {
     fn select_monolith<'a>(
         &'a self,
+        room: &RoomName,
         monoliths: Vec<&'a BalancerMonolith>,
     ) -> anyhow::Result<&BalancerMonolith>;
 
@@ -40,6 +42,7 @@ pub struct MinRoomsSelector;
 impl MonolithSelection for MinRoomsSelector {
     fn select_monolith<'a>(
         &'a self,
+        _room: &RoomName,
         monoliths: Vec<&'a BalancerMonolith>,
     ) -> anyhow::Result<&BalancerMonolith> {
         fn cmp(x: &BalancerMonolith, y: &BalancerMonolith) -> std::cmp::Ordering {
@@ -121,8 +124,9 @@ mod test {
 
         let monoliths: Vec<&BalancerMonolith> = vec![&monolith_one, &monolith_two];
 
+        let room: RoomName = "foo".into();
         let selected = MinRoomsSelector
-            .select_monolith(monoliths)
+            .select_monolith(&room, monoliths)
             .expect("failed to select monolith");
 
         assert_eq!(selected.id(), monolith_two.id())

crates/ott-balancer/src/service.rs

@@ -242,7 +242,7 @@ impl Service<Request<IncomingBody>> for BalancerService {
                                 );
                                 ctx_read.monoliths.get(&locator.monolith_id())
                             } else {
-                                ctx_read.select_monolith().ok()
+                                ctx_read.select_monolith(&room_name).ok()
                             };
                         if let Some(monolith) = monolith {
                             info!("proxying request to monolith {}", monolith.id());