balancer: reorganize monolith selection (#1648)

dyc3 · Apr 9, 2024 · f65a985 · f65a985
1 parent 808b884
commit f65a985
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 34 deletions.
diff --git a/crates/ott-balancer/src/balancer.rs b/crates/ott-balancer/src/balancer.rs
@@ -17,11 +17,11 @@ use tracing::{debug, error, info, instrument, trace, warn};
 
 use crate::balancer::collector::ClientState;
 use crate::client::ClientLink;
-use crate::config::{BalancerConfig, MonolithSelectionStrategy};
+use crate::config::BalancerConfig;
 use crate::connection::BALANCER_ID;
 use crate::monolith::Room;
 use crate::room::RoomLocator;
-use crate::selection::MonolithSelection;
+use crate::selection::{MonolithSelection, MonolithSelectionStrategy};
 use crate::{
     client::{BalancerClient, NewClient},
     messages::*,

diff --git a/crates/ott-balancer/src/config.rs b/crates/ott-balancer/src/config.rs
@@ -1,29 +1,17 @@
 use std::{borrow::BorrowMut, path::PathBuf, sync::Once};
 
 use clap::{Parser, ValueEnum};
-use enum_dispatch::enum_dispatch;
 use figment::providers::Format;
 use serde::Deserialize;
 
 use ott_common::discovery::DiscoveryConfig;
 
-use crate::selection::MinRoomsSelector;
+use crate::selection::MonolithSelectionStrategy;
 
 static mut CONFIG: Option<BalancerConfig> = None;
 
 static CONFIG_INIT: Once = Once::new();
 
-#[derive(Debug, Deserialize, Copy, Clone)]
-#[enum_dispatch]
-pub enum MonolithSelectionStrategy {
-    MinRooms(MinRoomsSelector),
-}
-
-impl Default for MonolithSelectionStrategy {
-    fn default() -> Self {
-        MonolithSelectionStrategy::MinRooms(MinRoomsSelector)
-    }
-}
 #[derive(Debug, Deserialize)]
 #[serde(default)]
 pub struct BalancerConfig {

diff --git a/crates/ott-balancer/src/selection.rs b/crates/ott-balancer/src/selection.rs
@@ -1,50 +1,57 @@
-use crate::config::MonolithSelectionStrategy;
 use crate::monolith::BalancerMonolith;
 use enum_dispatch::enum_dispatch;
 use rand::seq::IteratorRandom;
 use serde::Deserialize;
 
-#[derive(Debug, Default, Deserialize, Copy, Clone)]
-pub struct MinRoomsSelector;
 #[enum_dispatch(MonolithSelectionStrategy)]
 pub trait MonolithSelection: std::fmt::Debug {
     fn select_monolith<'a>(
         &'a self,
-        monolith: Vec<&'a BalancerMonolith>,
+        monoliths: Vec<&'a BalancerMonolith>,
     ) -> anyhow::Result<&BalancerMonolith>;
 
     fn random_monolith<'a>(
         &'a self,
-        monolith: Vec<&'a BalancerMonolith>,
-    ) -> anyhow::Result<&BalancerMonolith>;
+        monoliths: Vec<&'a BalancerMonolith>,
+    ) -> anyhow::Result<&BalancerMonolith> {
+        let selected = monoliths
+            .iter()
+            .choose(&mut rand::thread_rng())
+            .ok_or_else(|| anyhow::anyhow!("no monoliths available"))?;
+        Ok(selected)
+    }
+}
+
+#[derive(Debug, Deserialize, Copy, Clone)]
+#[enum_dispatch]
+pub enum MonolithSelectionStrategy {
+    MinRooms(MinRoomsSelector),
+}
+
+impl Default for MonolithSelectionStrategy {
+    fn default() -> Self {
+        MonolithSelectionStrategy::MinRooms(MinRoomsSelector)
+    }
 }
 
+#[derive(Debug, Default, Deserialize, Copy, Clone)]
+pub struct MinRoomsSelector;
+
 impl MonolithSelection for MinRoomsSelector {
     fn select_monolith<'a>(
         &'a self,
-        monolith: Vec<&'a BalancerMonolith>,
+        monoliths: 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));
+        let selected = monoliths.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)]