use std::{cmp, collections::HashSet, time::Duration}; use anyhow::Result; use log::*; use serde::{Deserialize, Serialize}; use serde_lexpr::{error, from_str}; use tokio::{sync::watch, time::sleep}; use crate::config::RuntimeConfigConsul; use crate::{consul, messages}; #[derive(Serialize, Deserialize, Debug)] pub enum DiplonatParameter { #[serde(rename = "tcp_port")] TcpPort(HashSet), #[serde(rename = "udp_port")] UdpPort(HashSet), } #[derive(Serialize, Deserialize, Debug)] pub enum DiplonatConsul { #[serde(rename = "diplonat")] Diplonat(Vec), } pub struct ConsulActor { pub rx_open_ports: watch::Receiver, consul: consul::Consul, node: String, retries: u32, tx_open_ports: watch::Sender, } 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 to_parameters(catalog: &consul::CatalogNode) -> Vec { let mut r = Vec::new(); for (_, service_info) in &catalog.services { for tag in &service_info.tags { let diplo_conf: error::Result = from_str(tag); match diplo_conf { Ok(conf) => r.push(conf), Err(e) => debug!("Failed to parse entry {}. {}", tag, e), }; } } return r; } fn to_open_ports(params: &Vec) -> messages::PublicExposedPorts { let mut op = messages::PublicExposedPorts { tcp_ports: HashSet::new(), udp_ports: HashSet::new(), }; for conf in params { let DiplonatConsul::Diplonat(c) = conf; for parameter in c { match parameter { DiplonatParameter::TcpPort(p) => op.tcp_ports.extend(p), DiplonatParameter::UdpPort(p) => op.udp_ports.extend(p), }; } } return op; } impl ConsulActor { pub fn new(config: &RuntimeConfigConsul, node: &str) -> Self { let (tx, rx) = watch::channel(messages::PublicExposedPorts { tcp_ports: HashSet::new(), udp_ports: HashSet::new(), }); return Self { consul: consul::Consul::new(config), rx_open_ports: rx, tx_open_ports: tx, node: node.to_string(), retries: 0, }; } pub async fn listen(&mut self) -> Result<()> { loop { let catalog = match self.consul.watch_node(&self.node).await { Ok(c) => c, Err(e) => { self.consul.watch_node_reset(); self.retries = cmp::min(std::u32::MAX - 1, self.retries) + 1; let will_retry_in = retry_to_time(self.retries, Duration::from_secs(600)); error!( "Failed to query consul. Will retry in {}s. {}", will_retry_in.as_secs(), e ); sleep(will_retry_in).await; continue; } }; self.retries = 0; let msg = to_open_ports(&to_parameters(&catalog)); debug!("Extracted configuration: {:#?}", msg); self.tx_open_ports.send(msg)?; } } }