use std::collections::{HashMap, HashSet}; use std::fmt; use std::sync::Arc; use std::{cmp, time::Duration}; use anyhow::Result; use futures::future::BoxFuture; use futures::stream::{FuturesUnordered, StreamExt}; use tokio::{select, sync::watch, time::sleep}; use tracing::*; use df_consul::*; const IPV4_TARGET_METADATA_TAG: &str = "public_ipv4"; const IPV6_TARGET_METADATA_TAG: &str = "public_ipv6"; const CNAME_TARGET_METADATA_TAG: &str = "cname_target"; // ---- Extract DNS config from Consul catalog ---- #[derive(Debug, Eq, PartialEq, Default)] pub struct DnsConfig { pub entries: HashMap, } #[derive(Clone, Debug, Hash, PartialEq, Eq)] pub struct DnsEntryKey { pub dns_path: String, pub record_type: DnsRecordType, } #[derive(Debug, PartialEq, Eq)] pub struct DnsEntryValue { pub targets: HashSet, } #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] #[allow(clippy::upper_case_acronyms)] pub enum DnsRecordType { A, AAAA, CNAME, } impl DnsConfig { pub fn new() -> Self { Self { entries: HashMap::new(), } } fn add(&mut self, k: DnsEntryKey, v: DnsEntryValue) { if let Some(ent) = self.entries.get_mut(&k) { ent.targets.extend(v.targets); } else { self.entries.insert(k, v); } } } fn parse_d53_tag(tag: &str, node: &ConsulNode) -> Option<(DnsEntryKey, DnsEntryValue)> { let splits = tag.split(' ').collect::>(); if splits.len() != 2 { return None; } let (record_type, targets) = match splits[0] { "d53-a" => match node.meta.get(IPV4_TARGET_METADATA_TAG) { Some(tgt) => (DnsRecordType::A, [tgt.to_string()].into_iter().collect()), None => { warn!("Got d53-a tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, IPV4_TARGET_METADATA_TAG); return None; } }, "d53-aaaa" => match node.meta.get(IPV6_TARGET_METADATA_TAG) { Some(tgt) => (DnsRecordType::AAAA, [tgt.to_string()].into_iter().collect()), None => { warn!("Got d53-aaaa tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, IPV6_TARGET_METADATA_TAG); return None; } }, "d53-cname" => match node.meta.get(CNAME_TARGET_METADATA_TAG) { Some(tgt) => ( DnsRecordType::CNAME, [tgt.to_string()].into_iter().collect(), ), None => { warn!("Got d53-cname tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, CNAME_TARGET_METADATA_TAG); return None; } }, _ => return None, }; Some(( DnsEntryKey { dns_path: splits[1].to_string(), record_type, }, DnsEntryValue { targets }, )) } fn parse_consul_catalog(catalog: &ConsulNodeCatalog, dns_config: &mut DnsConfig) { trace!("Parsing node catalog: {:#?}", catalog); for (_, svc) in catalog.services.iter() { for tag in svc.tags.iter() { if let Some((k, v)) = parse_d53_tag(tag, &catalog.node) { dns_config.add(k, v); } } } } #[derive(Default)] struct NodeWatchState { last_idx: Option, last_catalog: Option, retries: u32, } pub fn spawn_dns_config_task( consul: Consul, mut must_exit: watch::Receiver, ) -> watch::Receiver> { let (tx, rx) = watch::channel(Arc::new(DnsConfig::new())); let consul = Arc::new(consul); tokio::spawn(async move { let mut nodes = HashMap::new(); let mut watches = FuturesUnordered::)>>::new(); while !*must_exit.borrow() { let list_nodes = select! { ln = consul.list_nodes() => ln, _ = must_exit.changed() => continue, }; match list_nodes { Ok(consul_nodes) => { debug!("Watched consul nodes: {:?}", consul_nodes); for consul_node in consul_nodes { let node = &consul_node.node; 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) => { error!("Could not get Consul node list: {}", e); } } let next_watch = select! { nw = watches.next() => nw, _ = must_exit.changed() => continue, }; let (node, res): (String, Result<_>) = match next_watch { Some(v) => v, None => { warn!("No nodes currently watched in dns_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 = 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 dns_config = DnsConfig::new(); for (_, watch_state) in nodes.iter() { if let Some(catalog) = &watch_state.last_catalog { parse_consul_catalog(catalog, &mut dns_config); } } tx.send(Arc::new(dns_config)).expect("Internal error"); } }); rx } 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 Duration::from_secs(cmp::min( max_time.as_secs(), 1.2f64.powf(retries as f64) as u64, )) } // ---- Display impls ---- impl std::fmt::Display for DnsRecordType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { DnsRecordType::A => write!(f, "A"), DnsRecordType::AAAA => write!(f, "AAAA"), DnsRecordType::CNAME => write!(f, "CNAME"), } } } impl std::fmt::Display for DnsEntryKey { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{} IN {}", self.dns_path, self.record_type) } } impl std::fmt::Display for DnsEntryValue { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "[")?; for (i, tgt) in self.targets.iter().enumerate() { if i > 0 { write!(f, " ")?; } write!(f, "{}", tgt)?; } write!(f, "]") } }