allow ott-collector to discover balancers (#1419)

* allow ott-collector to discover balancers

* updated imports to reflect changes

---------

Co-authored-by: Carson McManus <[email protected]>

Author: moreno-michael

crates/ott-balancer/benches/joins.rs

@@ -8,9 +8,9 @@ use ott_balancer::{
     balancer::{start_dispatcher, Balancer, BalancerContext, BalancerLink},
     client::NewClient,
     config::BalancerConfig,
-    discovery::{HostOrIp, MonolithConnectionConfig},
     monolith::NewMonolith,
 };
+use ott_common::discovery::{HostOrIp, MonolithConnectionConfig};
 
 fn set_up_balancer() -> (BalancerLink, JoinHandle<()>) {
     let ctx = Arc::new(RwLock::new(BalancerContext::new()));

crates/ott-balancer/benches/latency.rs

@@ -15,10 +15,10 @@ use ott_balancer::{
     balancer::{start_dispatcher, Balancer, BalancerContext, BalancerLink},
     client::{ClientLink, NewClient},
     config::BalancerConfig,
-    discovery::{HostOrIp, MonolithConnectionConfig},
     messages::SocketMessage,
     monolith::NewMonolith,
 };
+use ott_common::discovery::{HostOrIp, MonolithConnectionConfig};
 
 async fn send_msg_monolith(
     link: &BalancerLink,

crates/ott-balancer/src/balancer.rs

@@ -701,10 +701,9 @@ pub async fn dispatch_monolith_message(
 
 #[cfg(test)]
 mod test {
+    use ott_common::discovery::{HostOrIp, MonolithConnectionConfig};
     use std::net::Ipv4Addr;
 
-    use crate::discovery::{HostOrIp, MonolithConnectionConfig};
-
     use super::*;
 
     #[tokio::test]

crates/ott-balancer/src/config.rs

@@ -4,7 +4,7 @@ use clap::{Parser, ValueEnum};
 use figment::providers::Format;
 use serde::Deserialize;
 
-use crate::discovery::{
+use ott_common::discovery::{
     DnsDiscoveryConfig, FlyDiscoveryConfig, HarnessDiscoveryConfig, ManualDiscoveryConfig,
 };
 

crates/ott-balancer/src/connection.rs

@@ -5,6 +5,7 @@ use std::time::Duration;
 
 use futures_util::{SinkExt, StreamExt};
 use ott_balancer_protocol::monolith::MsgM2B;
+use ott_common::discovery::{MonolithConnectionConfig, MonolithDiscoveryMsg};
 use tokio_tungstenite::{connect_async, WebSocketStream};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, warn};
@@ -13,7 +14,6 @@ use tungstenite::protocol::CloseFrame;
 use tungstenite::Message;
 
 use crate::balancer::BalancerLink;
-use crate::discovery::{MonolithConnectionConfig, MonolithDiscoveryMsg};
 use crate::messages::SocketMessage;
 use crate::monolith::NewMonolith;
 

crates/ott-balancer/src/discovery.rs

@@ -13,14 +13,15 @@ use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::EnvFilter;
 
 use crate::config::{BalancerConfig, DiscoveryConfig};
-use crate::discovery::start_discovery_task;
 use crate::service::BalancerService;
-
+use ott_common::discovery::{
+    start_discovery_task, DnsMonolithDiscoverer, FlyMonolithDiscoverer, HarnessMonolithDiscoverer,
+    ManualMonolithDiscoverer,
+};
 pub mod balancer;
 pub mod client;
 pub mod config;
 pub mod connection;
-pub mod discovery;
 pub mod messages;
 pub mod monolith;
 pub mod room;
@@ -81,19 +82,19 @@ pub async fn run() -> anyhow::Result<()> {
     info!("Starting monolith discovery");
     let _discovery_handle = match &config.discovery {
         DiscoveryConfig::Dns(config) => {
-            let discovery = discovery::DnsMonolithDiscoverer::new(config.clone());
+            let discovery = DnsMonolithDiscoverer::new(config.clone());
             start_discovery_task(discovery, discovery_tx)
         }
         DiscoveryConfig::Fly(config) => {
-            let discovery = discovery::FlyMonolithDiscoverer::new(config.clone());
+            let discovery = FlyMonolithDiscoverer::new(config.clone());
             start_discovery_task(discovery, discovery_tx)
         }
         DiscoveryConfig::Manual(config) => {
-            let discovery = discovery::ManualMonolithDiscoverer::new(config.clone());
+            let discovery = ManualMonolithDiscoverer::new(config.clone());
             start_discovery_task(discovery, discovery_tx)
         }
         DiscoveryConfig::Harness(config) => {
-            let discovery = discovery::HarnessMonolithDiscoverer::new(config.clone());
+            let discovery = HarnessMonolithDiscoverer::new(config.clone());
             start_discovery_task(discovery, discovery_tx)
         }
     };

crates/ott-balancer/src/lib.rs

@@ -4,10 +4,10 @@ use std::sync::Arc;
 use anyhow::bail;
 use ott_balancer_protocol::monolith::*;
 use ott_balancer_protocol::*;
+use ott_common::discovery::MonolithConnectionConfig;
 use tokio_tungstenite::tungstenite::Message;
 use tracing::error;
 
-use crate::discovery::MonolithConnectionConfig;
 use crate::messages::*;
 
 /// A Monolith refers to the NodeJS server that manages rooms and performs all business logic.

crates/ott-balancer/src/monolith.rs

@@ -49,9 +49,9 @@ 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 ott_common::discovery::{HostOrIp, MonolithConnectionConfig};
 
     use super::{MinRoomsSelector, MonolithSelection};
 

crates/ott-balancer/src/selection.rs

@@ -1,3 +1,4 @@
+use ott_common::discovery::{start_discovery_task, DnsDiscoveryConfig, DnsMonolithDiscoverer};
 use rocket::serde::json::Json;
 use serde::Serialize;
 
@@ -150,13 +151,22 @@ fn serve_state() -> Json<SystemState> {
     Json(return_sample_state())
 }
 
-#[launch]
-fn rocket() -> _ {
-    // TODO: spawn discovery tokio task here
-
-    rocket::build()
+#[rocket::main]
+async fn main() -> Result<(), rocket::Error> {
+    let (discovery_tx, _discovery_rx) = tokio::sync::mpsc::channel(2);
+    let discovery = DnsMonolithDiscoverer::new(DnsDiscoveryConfig {
+        monolith_port: 8081,
+        dns_server: None,
+        query: "balancer".to_string(),
+    });
+    start_discovery_task(discovery, discovery_tx);
+    let _rocket = rocket::build()
         .attach(cors::Cors)
         .mount("/", routes![status, cors::handle_preflight, serve_state])
+        .launch()
+        .await?;
+
+    Ok(())
 }
 
 #[get("/status")]

crates/ott-collector/src/main.rs

@@ -1,8 +1,10 @@
 //! Handles discovery of Monoliths.
 
+use std::collections::HashSet;
 use std::net::IpAddr;
 use std::net::SocketAddr;
 use std::time::Duration;
+use tracing::{debug, error, warn};
 
 mod dns;
 mod fly;
@@ -82,3 +84,86 @@ pub enum DiscoveryMode {
     Polling(Duration),
     Continuous,
 }
+pub struct DiscoveryTask {
+    discovery: Box<dyn MonolithDiscoverer + Send + Sync>,
+
+    monoliths: HashSet<MonolithConnectionConfig>,
+    discovery_tx: tokio::sync::mpsc::Sender<MonolithDiscoveryMsg>,
+}
+
+impl DiscoveryTask {
+    pub fn new(
+        discovery: impl MonolithDiscoverer + Send + Sync + 'static,
+        discovery_tx: tokio::sync::mpsc::Sender<MonolithDiscoveryMsg>,
+    ) -> Self {
+        Self {
+            discovery: Box::new(discovery),
+            monoliths: Default::default(),
+            discovery_tx,
+        }
+    }
+
+    pub async fn do_continuous_discovery(&mut self) {
+        loop {
+            if let Err(e) = self.do_discovery().await {
+                error!("Monolith Discovery failed: {:?}", e);
+                tokio::time::sleep(Duration::from_secs(5)).await;
+            }
+        }
+    }
+
+    async fn do_discovery(&mut self) -> anyhow::Result<()> {
+        let monoliths = self.discovery.discover().await?;
+        debug!("Discovered monoliths: {:?}", monoliths);
+        let monoliths_new: HashSet<_> = monoliths.into_iter().collect();
+        let msg = build_discovery_msg(&self.monoliths, &monoliths_new);
+        // apply the changes to our state
+        for m in &msg.removed {
+            self.monoliths.remove(m);
+        }
+        self.monoliths.extend(monoliths_new);
+        // send the message
+        self.discovery_tx.send(msg).await?;
+
+        if self.monoliths.is_empty() {
+            warn!("No monoliths discovered");
+        }
+
+        if let DiscoveryMode::Polling(d) = self.discovery.mode() {
+            tokio::time::sleep(d).await;
+        }
+
+        Ok(())
+    }
+}
+
+fn build_discovery_msg(
+    current: &HashSet<MonolithConnectionConfig>,
+    new: &HashSet<MonolithConnectionConfig>,
+) -> MonolithDiscoveryMsg {
+    let monoliths_added = new.difference(current).cloned().collect::<Vec<_>>();
+    let monoliths_removed = current.difference(new).cloned().collect::<Vec<_>>();
+    MonolithDiscoveryMsg {
+        added: monoliths_added,
+        removed: monoliths_removed,
+    }
+}
+
+pub fn start_discovery_task(
+    discovery: impl MonolithDiscoverer + Send + Sync + 'static,
+    discovery_tx: tokio::sync::mpsc::Sender<MonolithDiscoveryMsg>,
+) -> JoinHandle<()> {
+    tokio::task::Builder::new()
+        .name("discovery")
+        .spawn(async move {
+            let mut task = DiscoveryTask::new(discovery, discovery_tx);
+            task.do_continuous_discovery().await;
+        })
+        .expect("failed to spawn discovery task")
+}
+
+#[derive(Debug, Clone)]
+pub struct MonolithDiscoveryMsg {
+    pub added: Vec<MonolithConnectionConfig>,
+    pub removed: Vec<MonolithConnectionConfig>,
+}