path: root/src/consul_actor.rs

use std::cmp;
use std::collections::HashSet;
use std::time::Duration;

use anyhow::Result;
use log::*;
use serde::{Deserialize, Serialize};
use serde_lexpr::{error, from_str};
use tokio::sync::watch;
use tokio::time::delay_for;

use crate::config::RuntimeConfigConsul;
use crate::consul;
use crate::messages;

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "snake_case")]
pub enum DiplonatParameter {
  TcpPort(HashSet<u16>),
  UdpPort(HashSet<u16>),
  Acme(HashSet<String>),
}

#[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()
  // };
  let mut op = messages::PublicExposedPorts::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),
        DiplonatParameter::Acme(urls) => op.acme.extend(urls.clone()),
      };
    }
  }

  return op;
}

impl ConsulActor {
  pub fn new(config: RuntimeConfigConsul) -> Self {
    let (tx, rx) = watch::channel(messages::PublicExposedPorts::new());
    // 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)?;
    }
  }
}