forked from Deuxfleurs/diplonat
114 lines
3 KiB
Rust
114 lines
3 KiB
Rust
use std::cmp;
|
|
use std::collections::HashSet;
|
|
use std::time::Duration;
|
|
|
|
use anyhow::Result;
|
|
use log::*;
|
|
use serde::{Serialize, Deserialize};
|
|
use serde_lexpr::{from_str,error};
|
|
use tokio::sync::watch;
|
|
use tokio::time::delay_for;
|
|
|
|
use crate::config::RuntimeConfigConsul;
|
|
use crate::consul;
|
|
use crate::messages;
|
|
|
|
#[derive(Serialize, Deserialize, Debug)]
|
|
pub enum DiplonatParameter {
|
|
tcp_port(HashSet<u16>),
|
|
udp_port(HashSet<u16>)
|
|
}
|
|
|
|
#[derive(Serialize, Deserialize, Debug)]
|
|
pub enum DiplonatConsul {
|
|
diplonat(Vec<DiplonatParameter>)
|
|
}
|
|
|
|
pub struct ConsulActor {
|
|
pub rx_open_ports: watch::Receiver<messages::PublicExposedPorts>,
|
|
|
|
consul: consul::Consul,
|
|
node: String,
|
|
retries: u32,
|
|
|
|
tx_open_ports: watch::Sender<messages::PublicExposedPorts>
|
|
}
|
|
|
|
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<DiplonatConsul> {
|
|
let mut r = Vec::new();
|
|
|
|
for (_, service_info) in &catalog.Services {
|
|
for tag in &service_info.Tags {
|
|
let diplo_conf: error::Result<DiplonatConsul> = 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<DiplonatConsul>) -> 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::tcp_port(p) => op.tcp_ports.extend(p),
|
|
DiplonatParameter::udp_port(p) => op.udp_ports.extend(p),
|
|
};
|
|
}
|
|
}
|
|
|
|
return op;
|
|
}
|
|
|
|
impl ConsulActor {
|
|
pub fn new(config: RuntimeConfigConsul) -> Self {
|
|
let (tx, rx) = watch::channel(messages::PublicExposedPorts{
|
|
tcp_ports: HashSet::new(),
|
|
udp_ports: HashSet::new()
|
|
});
|
|
|
|
return Self {
|
|
consul: consul::Consul::new(&config.url),
|
|
node: config.node_name,
|
|
retries: 0,
|
|
rx_open_ports: rx,
|
|
tx_open_ports: tx,
|
|
};
|
|
}
|
|
|
|
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);
|
|
delay_for(will_retry_in).await;
|
|
continue;
|
|
}
|
|
};
|
|
self.retries = 0;
|
|
let msg = to_open_ports(&to_parameters(&catalog));
|
|
debug!("Extracted configuration: {:#?}", msg);
|
|
|
|
self.tx_open_ports.broadcast(msg)?;
|
|
}
|
|
}
|
|
}
|