Add compatibility with Docker dns servers (#1290)

* added dns.rs, adds compatibility with dns servers

* Adjusted format for Docker DNS

* added missing field

* made suggested changes, working on resolving
compiler issues

* fixed compiler issues

* ran cargo fmt

* made most suggested changes

* fixing compiler issues

* update balancer config for docker

* Update crates/ott-balancer/src/discovery/dns.rs

---------

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

Author: moreno-michael

crates/ott-balancer/src/config.rs

@@ -4,7 +4,9 @@ use clap::{Parser, ValueEnum};
 use figment::providers::Format;
 use serde::Deserialize;
 
-use crate::discovery::{FlyDiscoveryConfig, HarnessDiscoveryConfig, ManualDiscoveryConfig};
+use crate::discovery::{
+    DnsDiscoveryConfig, FlyDiscoveryConfig, HarnessDiscoveryConfig, ManualDiscoveryConfig,
+};
 
 static mut CONFIG: Option<BalancerConfig> = None;
 
@@ -32,6 +34,7 @@ impl Default for BalancerConfig {
 #[derive(Debug, Deserialize)]
 #[serde(tag = "method", rename_all = "lowercase")]
 pub enum DiscoveryConfig {
+    Dns(DnsDiscoveryConfig),
     Fly(FlyDiscoveryConfig),
     Manual(ManualDiscoveryConfig),
     Harness(HarnessDiscoveryConfig),

crates/ott-balancer/src/discovery.rs

@@ -4,10 +4,12 @@ use std::net::IpAddr;
 use std::time::Duration;
 use std::{collections::HashSet, net::SocketAddr};
 
+mod dns;
 mod fly;
 mod harness;
 mod manual;
 
+pub use dns::*;
 pub use fly::*;
 pub use harness::*;
 pub use manual::*;

crates/ott-balancer/src/discovery/dns.rs

@@ -0,0 +1,52 @@
+use async_trait::async_trait;
+use tracing::info;
+use trust_dns_resolver::TokioAsyncResolver;
+
+use super::*;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct DnsDiscoveryConfig {
+    /// The port that monoliths should be listening on for load balancer connections.
+    pub monolith_port: u16,
+    /// The DNS server to query. Optional. If not provided, the system configuration will be used instead.
+    pub dns_server: Option<String>,
+    /// The A record to query. If using docker-compose, this should be the service name for the monolith.
+    pub query: String,
+}
+
+pub struct DnsMonolithDiscoverer {
+    config: DnsDiscoveryConfig,
+}
+
+impl DnsMonolithDiscoverer {
+    pub fn new(config: DnsDiscoveryConfig) -> Self {
+        info!(
+            "Creating DnsMonolithDiscoverer, DNS server: {:?}",
+            config.dns_server
+        );
+        Self { config }
+    }
+}
+
+#[async_trait]
+impl MonolithDiscoverer for DnsMonolithDiscoverer {
+    async fn discover(&mut self) -> anyhow::Result<Vec<MonolithConnectionConfig>> {
+        let resolver =
+            TokioAsyncResolver::tokio_from_system_conf().expect("failed to create resolver");
+
+        let lookup = resolver.ipv4_lookup(&self.config.query).await?;
+        let monoliths = lookup
+            .iter()
+            .map(|ip| MonolithConnectionConfig {
+                host: HostOrIp::Ip(IpAddr::V4(*ip)),
+                port: self.config.monolith_port,
+            })
+            .collect::<Vec<_>>();
+
+        Ok(monoliths)
+    }
+
+    fn mode(&self) -> DiscoveryMode {
+        DiscoveryMode::Polling(Duration::from_secs(10))
+    }
+}

crates/ott-balancer/src/lib.rs

@@ -65,6 +65,10 @@ 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());
+            start_discovery_task(discovery, discovery_tx)
+        }
         DiscoveryConfig::Fly(config) => {
             let discovery = discovery::FlyMonolithDiscoverer::new(config.clone());
             start_discovery_task(discovery, discovery_tx)

deploy/ott-balancer-docker.toml

@@ -1,4 +1,4 @@
 [discovery]
-method = "docker"
+method = "dns"
 monolith_port = 3002
-monolith_container_name = "opentogethertube"
+query = "opentogethertube"