use std::collections::HashMap; use std::net::SocketAddr; use std::sync::{atomic, Arc}; use std::{cmp, time::Duration}; use anyhow::Result; use futures::future::BoxFuture; use futures::stream::{FuturesUnordered, StreamExt}; use log::*; use tokio::{sync::watch, time::sleep}; use crate::consul::*; // ---- Extract proxy config from Consul catalog ---- #[derive(Debug)] pub enum HostDescription { Hostname(String), Pattern(glob::Pattern), } impl HostDescription { fn new(desc: &str) -> Result { if desc.chars().any(|x| matches!(x, '*' | '?' | '[' | ']')) { Ok(Self::Pattern(glob::Pattern::new(desc)?)) } else { Ok(Self::Hostname(desc.to_string())) } } pub fn matches(&self, v: &str) -> bool { match self { Self::Pattern(p) => p.matches(v), Self::Hostname(s) => s == v, } } } #[derive(Debug)] pub struct ProxyEntry { pub target_addr: SocketAddr, pub host: HostDescription, pub path_prefix: Option, pub priority: u32, pub add_headers: Vec<(String, String)>, // Counts the number of times this proxy server has been called to // This implements a round-robin load balancer if there are multiple // entries for the same host and same path prefix. pub calls: atomic::AtomicU64, } impl std::fmt::Display for ProxyEntry { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{} ", self.target_addr)?; match &self.host { HostDescription::Hostname(h) => write!(f, "{}", h)?, HostDescription::Pattern(p) => write!(f, "Pattern('{}')", p.as_str())?, } write!( f, "{} {}", self.path_prefix.as_ref().unwrap_or(&String::new()), self.priority )?; if !self.add_headers.is_empty() { write!(f, " +Headers: {:?}", self.add_headers)?; } Ok(()) } } #[derive(Debug)] pub struct ProxyConfig { pub entries: Vec, } fn retry_to_time(retries: u32, max_time: Duration) -> Duration { // 1.2^x seems to be a good value to exponentially increase time at a good pace // eg. 1.2^32 = 341 seconds ~= 5 minutes - ie. after 32 retries we wait 5 // minutes return Duration::from_secs(cmp::min( max_time.as_secs(), 1.2f64.powf(retries as f64) as u64, )); } fn parse_tricot_tag( tag: &str, target_addr: SocketAddr, add_headers: &[(String, String)], ) -> Option { let splits = tag.split(' ').collect::>(); if (splits.len() != 2 && splits.len() != 3) || splits[0] != "tricot" { return None; } let (host, path_prefix) = match splits[1].split_once('/') { Some((h, p)) => (h, Some(p.to_string())), None => (splits[1], None), }; let priority = match splits.len() { 3 => splits[2].parse().ok()?, _ => 100, }; let host = match HostDescription::new(host) { Ok(h) => h, Err(e) => { warn!("Invalid hostname pattern {}: {}", host, e); return None; } }; Some(ProxyEntry { target_addr, host, path_prefix, priority, add_headers: add_headers.to_vec(), calls: atomic::AtomicU64::from(0), }) } fn parse_tricot_add_header_tag(tag: &str) -> Option<(String, String)> { let splits = tag.split(' ').collect::>(); if splits.len() == 3 && splits[0] == "tricot-add-header" { Some((splits[1].to_string(), splits[2].to_string())) } else { None } } fn parse_consul_catalog(catalog: &ConsulNodeCatalog) -> Vec { trace!("Parsing node catalog: {:#?}", catalog); let mut entries = vec![]; for (_, svc) in catalog.services.iter() { let ip_addr = match svc.address.parse() { Ok(ip) => ip, _ => match catalog.node.address.parse() { Ok(ip) => ip, _ => { warn!( "Could not get address for service {} at node {}", svc.service, catalog.node.node ); continue; } }, }; let addr = SocketAddr::new(ip_addr, svc.port); let mut add_headers = vec![]; for tag in svc.tags.iter() { if let Some(pair) = parse_tricot_add_header_tag(tag) { add_headers.push(pair); } } for tag in svc.tags.iter() { if let Some(ent) = parse_tricot_tag(tag, addr, &add_headers[..]) { entries.push(ent); } } } trace!("Result of parsing catalog:"); for ent in entries.iter() { trace!(" {}", ent); } entries } #[derive(Default)] struct NodeWatchState { last_idx: Option, last_catalog: Option, retries: u32, } pub fn spawn_proxy_config_task(consul: Consul) -> watch::Receiver> { let (tx, rx) = watch::channel(Arc::new(ProxyConfig { entries: Vec::new(), })); let consul = Arc::new(consul); tokio::spawn(async move { let mut nodes = HashMap::new(); let mut watches = FuturesUnordered::)>>::new(); loop { match consul.list_nodes().await { Ok(consul_nodes) => { info!("Watched consul nodes: {:?}", consul_nodes); for node in consul_nodes { if !nodes.contains_key(&node) { nodes.insert(node.clone(), NodeWatchState::default()); let node = node.to_string(); let consul = consul.clone(); watches.push(Box::pin(async move { let res = consul.watch_node(&node, None).await; (node, res) })); } } } Err(e) => { warn!("Could not get Consul node list: {}", e); } } let (node, res): (String, Result<_>) = match watches.next().await { Some(v) => v, None => { warn!("No nodes currently watched in proxy_config.rs"); sleep(Duration::from_secs(10)).await; continue; } }; match res { Ok((catalog, new_idx)) => { let mut watch_state = nodes.get_mut(&node).unwrap(); watch_state.last_idx = Some(new_idx); watch_state.last_catalog = Some(catalog); watch_state.retries = 0; let idx = watch_state.last_idx; let consul = consul.clone(); watches.push(Box::pin(async move { let res = consul.watch_node(&node, idx).await; (node, res) })); } Err(e) => { let mut watch_state = nodes.get_mut(&node).unwrap(); watch_state.retries += 1; watch_state.last_idx = None; let will_retry_in = retry_to_time(watch_state.retries, Duration::from_secs(600)); error!( "Failed to query consul for node {}. Will retry in {}s. {}", node, will_retry_in.as_secs(), e ); let consul = consul.clone(); watches.push(Box::pin(async move { sleep(will_retry_in).await; let res = consul.watch_node(&node, None).await; (node, res) })); continue; } } let mut entries = vec![]; for (_, watch_state) in nodes.iter() { if let Some(catalog) = &watch_state.last_catalog { entries.extend(parse_consul_catalog(catalog)); } } let config = ProxyConfig { entries }; tx.send(Arc::new(config)).expect("Internal error"); } }); rx }