From b5bb613adae51fd41dad2a6fe93aac7a90514ac2 Mon Sep 17 00:00:00 2001
From: Carson McManus <carson.mcmanus1@gmail.com>
Date: Tue, 9 Apr 2024 09:35:14 -0400
Subject: [PATCH] balancer: add room arg to monolith selection

---
 crates/ott-balancer/src/balancer.rs  | 6 +++---
 crates/ott-balancer/src/selection.rs | 6 +++++-
 crates/ott-balancer/src/service.rs   | 2 +-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/crates/ott-balancer/src/balancer.rs b/crates/ott-balancer/src/balancer.rs
index bedd9729b..07522d913 100644
--- a/crates/ott-balancer/src/balancer.rs
+++ b/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(),
diff --git a/crates/ott-balancer/src/selection.rs b/crates/ott-balancer/src/selection.rs
index b3ef93c88..f75162d31 100644
--- a/crates/ott-balancer/src/selection.rs
+++ b/crates/ott-balancer/src/selection.rs
@@ -1,5 +1,6 @@
 use crate::monolith::BalancerMonolith;
 use enum_dispatch::enum_dispatch;
+use ott_balancer_protocol::RoomName;
 use rand::seq::IteratorRandom;
 use serde::Deserialize;
 
@@ -7,6 +8,7 @@ use serde::Deserialize;
 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())
diff --git a/crates/ott-balancer/src/service.rs b/crates/ott-balancer/src/service.rs
index 0e45f313c..addd16e93 100644
--- a/crates/ott-balancer/src/service.rs
+++ b/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());