Merge branch 'main' into healthcheck

src/main.rs

@@ -3,9 +3,12 @@ mod balancer;
 mod config;
 mod proxy;
 
-use crate::balancer::{Balancer, ConnectionInfo};
+use std::collections::HashMap;
+use crate::balancer::{ConnectionInfo};
 use crate::proxy::tcp::proxy_tcp_connection;
 use std::fs::File;
+use std::path::PathBuf;
+use std::net::IpAddr;
 use std::sync::atomic::{AtomicU64, Ordering};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::io::{AsyncBufReadExt, AsyncReadExt};
@@ -93,28 +96,54 @@ async fn start_healthcheck_listener(addr: &str, healths: HashMap<IpAddr, Arc<RwL
     }
 
     Ok(())
+use std::sync::{Arc, Mutex, RwLock};
+use std::time::Duration;
+use anywho::Error;
+use tokio::net::TcpListener;
+use tokio::sync::mpsc;
+use crate::backend::ServerMetrics;
+use crate::config::loader::{build_lb, RoutingTable};
+
+use notify::{Watcher, RecursiveMode, Event};
+use clap::Parser;
+
+static NEXT_CONN_ID: AtomicU64 = AtomicU64::new(1);
+
+struct ProgramState {
+    tx_rt_map: HashMap<u16, mpsc::UnboundedSender<RoutingTable>>,
+    healths: HashMap<IpAddr, Arc<RwLock<ServerMetrics>>>,
+}
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    #[arg(short, long, default_value = "config.yaml")]
+    config: PathBuf,
 }
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let f = File::open("config.yaml").expect("couldn't open config.yaml");
-    let app_config: config::AppConfig = serde_saphyr::from_reader(f)?;
+    let args = Args::parse();
 
-    println!(
-        "Loaded {} backends, {} rules.",
-        app_config.backends.len(),
-        app_config.rules.len()
-    );
-
-    let (listeners, healths) = config::loader::build_lb(app_config);
-
-    if listeners.is_empty() {
-        eprintln!("its a lawless land");
-        return Ok(());
+    if !args.config.is_file() {
+        eprintln!("config file not found or not accessible");
+        std::process::exit(1);
     }
 
-    let mut handles = Vec::new();
+    println!("reading config from {:?}", args.config);
 
+    let state = Arc::new(Mutex::new(ProgramState {
+        tx_rt_map: HashMap::new(),
+        healths: HashMap::new(),
+    }));
+
+    if let Err(e) = load_config(&args.config, state.clone()).await {
+        eprintln!("config file loading failed: {}", e);
+    }
+
+    let config_path = args.config.clone();
+    let state_clone = state.clone();
+  
     handles.push(
         tokio::spawn(async {
             start_healthcheck_listener("0.0.0.0:8080", healths).await.unwrap();
@@ -127,53 +156,136 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         })
     );
 
-    for (port, mut routing_table) in listeners {
-        handles.push(tokio::spawn(async move {
-            let addr = format!("0.0.0.0:{}", port);
-            println!("Starting tcp listener on {}", addr);
+    tokio::spawn(async move {
+        let (tx, mut rx) = mpsc::channel(1);
+      
+        let mut watcher = notify::recommended_watcher(move |res: Result<Event, notify::Error>| {
+            if let Ok(event) = res {
+                if event.kind.is_modify() {
+                    let _ = tx.blocking_send(());
+                }
+            }
+        }).unwrap();
 
-            let listener = TcpListener::bind(&addr).await.expect("Failed to bind port");
+        watcher.watch(&config_path, RecursiveMode::NonRecursive).unwrap();
+        println!("watching for changes to {:?}", config_path);
 
-            loop {
-                let (socket, remote_addr) = match listener.accept().await {
-                    Ok(v) => v,
-                    Err(e) => {
-                        eprintln!("error: listener port {}: {}", port, e);
-                        continue;
+        while rx.recv().await.is_some() {
+            if let Err(e) = load_config(&config_path, state_clone.clone()).await {
+                eprintln!("loading config failed: {}", e);
+            }
+        }
+    });
+
+    loop { tokio::time::sleep(Duration::from_hours(1)).await; }
+}
+
+async fn load_config(path: &PathBuf, state: Arc<Mutex<ProgramState>>) -> Result<(), Error> {
+    let f = File::open(path)?;
+    let app_config: config::AppConfig = match serde_saphyr::from_reader(f) {
+        Ok(app_config) => app_config,
+        Err(e) => { eprintln!("error parsing config {}", e); return Ok(()); }
+    };
+
+    println!(
+        "Loaded config, with {} backends, {} rules.",
+        app_config.backends.len(),
+        app_config.rules.len()
+    );
+
+    let (mut listeners, health_monitors) = match build_lb(app_config) {
+        Ok(v) => v,
+        Err(e) => {
+            eprintln!("config has logical errors: {}", e);
+            return Ok(());
+        }
+    };
+    let mut prog_state = state.lock().unwrap();
+
+    let ports_to_remove: Vec<u16> = prog_state.tx_rt_map
+        .keys()
+        .cloned()
+        .filter(|port| !listeners.contains_key(port))
+        .collect();
+
+    for port in ports_to_remove {
+        prog_state.tx_rt_map.remove(&port);
+    }
+
+    prog_state.healths = health_monitors;
+    for (port, routing_table) in listeners.drain() {
+        if let Some(x) = prog_state.tx_rt_map.get_mut(&port) {
+            x.send(routing_table)?;
+            println!("updated rules on port {}", port);
+        } else {
+            let (tx_rt, rx_rt) = mpsc::unbounded_channel();
+            prog_state.tx_rt_map.insert(port, tx_rt);
+
+            tokio::spawn(run_listener(port, rx_rt, routing_table));
+        }
+    }
+
+    println!("reload complete");
+    Ok(())
+}
+
+async fn run_listener(
+    port: u16,
+    mut rx_rt: mpsc::UnboundedReceiver<RoutingTable>,
+    mut current_table: RoutingTable
+) {
+    let addr = format!("0.0.0.0:{}", port);
+    println!("Starting tcp listener on {}", addr);
+
+    let listener = TcpListener::bind(&addr).await.expect("Failed to bind port");
+
+    loop {
+        tokio::select! {
+            msg = rx_rt.recv() => {
+                match msg {
+                    Some(new_table) => {
+                        current_table = new_table;
                     }
-                };
-
-                let remote_ip = remote_addr.ip();
-                let conn_id = NEXT_CONN_ID.fetch_add(1, Ordering::Relaxed);
-
-                let mut chosen_backend = None;
-
-                for (cidr, balancer_idx) in &mut routing_table.entries {
-                    if cidr.contains(&remote_ip) {
-                        let balancer = &mut routing_table.balancers[*balancer_idx];
-                        chosen_backend = balancer.choose_backend(ConnectionInfo {
-                            client_ip: remote_ip,
-                        });
+                    None => {
+                        println!("Unbinding listener on port {}", port);
                         break;
                     }
                 }
+            }
+            accept_result = listener.accept() => {
+                match accept_result {
+                    Ok((socket, remote_addr)) => {
+                        let remote_ip = remote_addr.ip();
+                        let conn_id = NEXT_CONN_ID.fetch_add(1, Ordering::Relaxed);
 
-                if let Some(backend) = chosen_backend {
-                    tokio::spawn(async move {
-                        if let Err(e) = proxy_tcp_connection(conn_id, socket, backend).await {
-                            eprintln!("error: conn_id={} proxy failed: {}", conn_id, e);
+                        let mut chosen_backend = None;
+
+                        for (cidr, balancer_idx) in &mut current_table.entries {
+                            if cidr.contains(&remote_ip) {
+                                let balancer = &mut current_table.balancers[*balancer_idx];
+                                chosen_backend = balancer.choose_backend(ConnectionInfo {
+                                    client_ip: remote_ip,
+                                });
+                                break;
+                            }
                         }
-                    });
-                } else {
-                    println!("error: no matching rule for {} on port {}", remote_ip, port);
+
+                        if let Some(backend) = chosen_backend {
+                            tokio::spawn(async move {
+                                if let Err(e) = proxy_tcp_connection(conn_id, socket, backend).await {
+                                    eprintln!("error: conn_id={} proxy failed: {}", conn_id, e);
+                                }
+                            });
+                        } else {
+                            println!("error: no matching rule for {} on port {}", remote_ip, port);
+                        }
+                    }
+                    Err(e) => {
+                        eprintln!("error: listener port {}: {}", port, e);
+                    continue;
+                    }
                 }
             }
-        }));
+        }
     }
-
-    for h in handles {
-        let _ = h.await;
-    }
-
-    Ok(())
 }