remove stupid cache and use proper watch on consul kv prefix

6 files changed, 212 insertions, 77 deletions

diff --git a/Cargo.lock b/Cargo.lock
index addbd5c..699e98c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3,6 +3,15 @@
 version = 3
 
 [[package]]
+name = "aho-corasick"
+version = "0.7.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc936419f96fa211c1b9166887b38e5e40b19958e5b895be7c1f93adec7071ac"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
 name = "ansi_term"
 version = "0.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -104,6 +113,7 @@ dependencies = [
  "async-trait",
  "df-consul",
  "futures",
+ "regex",
  "reqwest",
  "serde",
  "serde_json",
@@ -115,11 +125,12 @@ dependencies = [
 
 [[package]]
 name = "df-consul"
-version = "0.3.3"
+version = "0.3.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0e38cfbab431b53dfd2d09f2a9902510c636d3d7397645bac5cf1959cfde2999"
+checksum = "565fcd7efcbdc3e3420e70bc38187a8dd6d5f22759858c32e6af14329bf27ff3"
 dependencies = [
  "anyhow",
+ "base64",
  "bytes",
  "futures",
  "log",
@@ -548,6 +559,8 @@ version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e076559ef8e241f2ae3479e36f97bd5741c0330689e217ad51ce2c76808b868a"
 dependencies = [
+ "aho-corasick",
+ "memchr",
  "regex-syntax",
 ]
 
diff --git a/Cargo.nix b/Cargo.nix
index 4b369b1..464d61d 100644
--- a/Cargo.nix
+++ b/Cargo.nix
@@ -23,7 +23,7 @@ args@{
   ignoreLockHash,
 }:
 let
-  nixifiedLockHash = "f8fd87706eb1709f2cf3a695f4400e1a5e130e3c599cdfebc00bba51c80f494f";
+  nixifiedLockHash = "d1a9c3cd406d87d45242e62c6697855d0e5e3c787d92d2fcd084b41469659da1";
   workspaceSrc = if args.workspaceSrc == null then ./. else args.workspaceSrc;
   currentLockHash = builtins.hashFile "sha256" (workspaceSrc + /Cargo.lock);
   lockHashIgnored = if ignoreLockHash
@@ -47,6 +47,20 @@ in
   workspace = {
     d53 = rustPackages.unknown.d53."0.1.0";
   };
+  "registry+https://github.com/rust-lang/crates.io-index".aho-corasick."0.7.20" = overridableMkRustCrate (profileName: rec {
+    name = "aho-corasick";
+    version = "0.7.20";
+    registry = "registry+https://github.com/rust-lang/crates.io-index";
+    src = fetchCratesIo { inherit name version; sha256 = "cc936419f96fa211c1b9166887b38e5e40b19958e5b895be7c1f93adec7071ac"; };
+    features = builtins.concatLists [
+      [ "default" ]
+      [ "std" ]
+    ];
+    dependencies = {
+      memchr = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".memchr."2.5.0" { inherit profileName; }).out;
+    };
+  });
+  
   "registry+https://github.com/rust-lang/crates.io-index".ansi_term."0.12.1" = overridableMkRustCrate (profileName: rec {
     name = "ansi_term";
     version = "0.12.1";
@@ -188,8 +202,9 @@ in
     dependencies = {
       anyhow = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".anyhow."1.0.66" { inherit profileName; }).out;
       async_trait = (buildRustPackages."registry+https://github.com/rust-lang/crates.io-index".async-trait."0.1.59" { profileName = "__noProfile"; }).out;
-      df_consul = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.3" { inherit profileName; }).out;
+      df_consul = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.4" { inherit profileName; }).out;
       futures = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".futures."0.3.25" { inherit profileName; }).out;
+      regex = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".regex."1.7.0" { inherit profileName; }).out;
       reqwest = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".reqwest."0.11.13" { inherit profileName; }).out;
       serde = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".serde."1.0.149" { inherit profileName; }).out;
       serde_json = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".serde_json."1.0.89" { inherit profileName; }).out;
@@ -200,13 +215,14 @@ in
     };
   });
   
-  "registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.3" = overridableMkRustCrate (profileName: rec {
+  "registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.4" = overridableMkRustCrate (profileName: rec {
     name = "df-consul";
-    version = "0.3.3";
+    version = "0.3.4";
     registry = "registry+https://github.com/rust-lang/crates.io-index";
-    src = fetchCratesIo { inherit name version; sha256 = "0e38cfbab431b53dfd2d09f2a9902510c636d3d7397645bac5cf1959cfde2999"; };
+    src = fetchCratesIo { inherit name version; sha256 = "565fcd7efcbdc3e3420e70bc38187a8dd6d5f22759858c32e6af14329bf27ff3"; };
     dependencies = {
       anyhow = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".anyhow."1.0.66" { inherit profileName; }).out;
+      base64 = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".base64."0.13.1" { inherit profileName; }).out;
       bytes = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".bytes."1.3.0" { inherit profileName; }).out;
       futures = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".futures."0.3.25" { inherit profileName; }).out;
       log = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".log."0.4.17" { inherit profileName; }).out;
@@ -810,9 +826,27 @@ in
     registry = "registry+https://github.com/rust-lang/crates.io-index";
     src = fetchCratesIo { inherit name version; sha256 = "e076559ef8e241f2ae3479e36f97bd5741c0330689e217ad51ce2c76808b868a"; };
     features = builtins.concatLists [
+      [ "aho-corasick" ]
+      [ "default" ]
+      [ "memchr" ]
+      [ "perf" ]
+      [ "perf-cache" ]
+      [ "perf-dfa" ]
+      [ "perf-inline" ]
+      [ "perf-literal" ]
       [ "std" ]
+      [ "unicode" ]
+      [ "unicode-age" ]
+      [ "unicode-bool" ]
+      [ "unicode-case" ]
+      [ "unicode-gencat" ]
+      [ "unicode-perl" ]
+      [ "unicode-script" ]
+      [ "unicode-segment" ]
     ];
     dependencies = {
+      aho_corasick = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".aho-corasick."0.7.20" { inherit profileName; }).out;
+      memchr = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".memchr."2.5.0" { inherit profileName; }).out;
       regex_syntax = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".regex-syntax."0.6.28" { inherit profileName; }).out;
     };
   });
diff --git a/Cargo.toml b/Cargo.toml
index 0b26b63..5d19dc0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,7 +15,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 structopt = "0.3"
 tokio = { version = "1.22", default-features = false, features = ["rt", "rt-multi-thread", "io-util", "net", "time", "macros", "sync", "signal", "fs"] }
 reqwest = { version = "0.11", default-features = false, features = ["json", "rustls-tls-webpki-roots" ] }
+regex = "1"
 serde = { version = "1.0.107", features = ["derive"] }
 serde_json = "1.0"
 
-df-consul = "0.3.3"
+df-consul = "0.3.4"
diff --git a/src/autodiscovery.rs b/src/autodiscovery.rs
new file mode 100644
index 0000000..9fcc094
--- /dev/null
+++ b/src/autodiscovery.rs
@@ -0,0 +1,110 @@
+//! Fetch autodiscoverd IP addresses stored by Diplonat into Consul
+
+use std::collections::HashMap;
+use std::net::{Ipv4Addr, Ipv6Addr};
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
+
+use anyhow::{anyhow, bail, Result};
+use regex::Regex;
+use serde::{Deserialize, Serialize};
+use tokio::{select, sync::watch};
+use tracing::*;
+
+use df_consul::*;
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct DiplonatAutodiscoveryResult<A> {
+    pub timestamp: u64,
+    pub address: Option<A>,
+}
+
+#[derive(Default, Debug)]
+pub struct AutodiscoveredAddresses {
+    pub ipv4: HashMap<String, DiplonatAutodiscoveryResult<Ipv4Addr>>,
+    pub ipv6: HashMap<String, DiplonatAutodiscoveryResult<Ipv6Addr>>,
+}
+
+pub fn watch_autodiscovered_ips(
+    consul: Consul,
+    mut must_exit: watch::Receiver<bool>,
+) -> watch::Receiver<Arc<AutodiscoveredAddresses>> {
+    let (tx, rx) = watch::channel(Arc::new(AutodiscoveredAddresses::default()));
+
+    tokio::spawn(async move {
+        let mut last_index = None;
+        let re = Regex::new(r".*autodiscovery/(\w+)/(\w+)$").unwrap();
+
+        while !*must_exit.borrow() {
+            let r = select! {
+                _ = must_exit.changed() => continue,
+                r = consul.kv_get_prefix("diplonat/autodiscovery/", last_index) => r,
+            };
+
+            let entries = match r {
+                Err(e) => {
+                    warn!("Error fetching diplonat autodiscovery consul prefix: {}", e);
+                    tokio::time::sleep(Duration::from_secs(30)).await;
+                    continue;
+                }
+                Ok(r) => {
+                    last_index = Some(r.index());
+                    r.into_inner()
+                }
+            };
+
+            let mut addresses = AutodiscoveredAddresses::default();
+
+            for (k, v) in entries {
+                if let Err(e) = parse_autodiscovered_address(&re, &mut addresses, &k, &v) {
+                    warn!(
+                        "Invalid k/v pair in diplonat autodiscovery results: {} = {} ({})",
+                        k,
+                        std::str::from_utf8(&v).unwrap_or("<???>"),
+                        e
+                    );
+                }
+            }
+
+            if tx.send(Arc::new(addresses)).is_err() {
+                info!("Autodiscovered addresses watcher terminating");
+            }
+        }
+    });
+
+    rx
+}
+
+fn parse_autodiscovered_address(
+    re: &Regex,
+    addresses: &mut AutodiscoveredAddresses,
+    k: &str,
+    v: &[u8],
+) -> Result<()> {
+    let caps = re.captures(k).ok_or(anyhow!("key does not match regex"))?;
+
+    if let (Some(family), Some(node)) = (caps.get(1), caps.get(2)) {
+        match family.as_str() {
+            "ipv4" => {
+                let r: DiplonatAutodiscoveryResult<Ipv4Addr> = serde_json::from_slice(v)?;
+                addresses.ipv4.insert(node.as_str().to_string(), r);
+            }
+            "ipv6" => {
+                let r: DiplonatAutodiscoveryResult<Ipv6Addr> = serde_json::from_slice(v)?;
+                addresses.ipv6.insert(node.as_str().to_string(), r);
+            }
+            _ => bail!("invalid address family {}", family.as_str()),
+        }
+    } else {
+        bail!("invalid regex captures {:?}", caps);
+    }
+
+    Ok(())
+}
+
+pub fn timestamp() -> u64 {
+    SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .expect("clock error")
+        .as_secs()
+}
diff --git a/src/dns_config.rs b/src/dns_config.rs
index 8e82c32..acee8d7 100644
--- a/src/dns_config.rs
+++ b/src/dns_config.rs
@@ -2,7 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::fmt;
 use std::net::{Ipv4Addr, Ipv6Addr};
 use std::sync::Arc;
-use std::time::{Duration, SystemTime};
+use std::time::Duration;
 
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
@@ -11,6 +11,8 @@ use tracing::*;
 
 use df_consul::*;
 
+use crate::autodiscovery::*;
+
 const IP_TARGET_METADATA_TAG_PREFIX: &str = "public_";
 const CNAME_TARGET_METADATA_TAG: &str = "cname_target";
 
@@ -66,29 +68,25 @@ pub fn spawn_dns_config_task(
 ) -> watch::Receiver<Arc<DnsConfig>> {
     let (tx, rx) = watch::channel(Arc::new(DnsConfig::new()));
 
-    let fetcher = DnsConfigFetcher {
-        consul,
-        node_ipv4_cache: HashMap::new(),
-        node_ipv6_cache: HashMap::new(),
-    };
+    let fetcher = DnsConfigTask { consul };
 
     tokio::spawn(fetcher.task(tx, must_exit));
 
     rx
 }
 
-struct DnsConfigFetcher {
+struct DnsConfigTask {
     consul: Consul,
-    node_ipv4_cache: HashMap<String, (u64, Option<Ipv4Addr>)>,
-    node_ipv6_cache: HashMap<String, (u64, Option<Ipv6Addr>)>,
 }
 
-impl DnsConfigFetcher {
+impl DnsConfigTask {
     async fn task(
         mut self,
         tx: watch::Sender<Arc<DnsConfig>>,
         mut must_exit: watch::Receiver<bool>,
     ) {
+        let mut autodiscovery_rx = watch_autodiscovered_ips(self.consul.clone(), must_exit.clone());
+
         let mut catalog_rx = self
             .consul
             .watch_all_service_health(Duration::from_secs(60));
@@ -96,11 +94,13 @@ impl DnsConfigFetcher {
         while !*must_exit.borrow() {
             select! {
                 _ = catalog_rx.changed() => (),
+                _ = autodiscovery_rx.changed() => (),
                 _ = must_exit.changed() => continue,
             };
 
             let services = catalog_rx.borrow_and_update().clone();
-            match self.parse_catalog(&services).await {
+            let autodiscovery = autodiscovery_rx.borrow_and_update().clone();
+            match self.parse_catalog(&services, &autodiscovery) {
                 Ok(dns_config) => tx.send(Arc::new(dns_config)).expect("Internal error"),
                 Err(e) => {
                     error!("Error when parsing tags: {}", e);
@@ -109,7 +109,11 @@ impl DnsConfigFetcher {
         }
     }
 
-    async fn parse_catalog(&mut self, services: &catalog::AllServiceHealth) -> Result<DnsConfig> {
+    fn parse_catalog(
+        &mut self,
+        services: &catalog::AllServiceHealth,
+        autodiscovery: &AutodiscoveredAddresses,
+    ) -> Result<DnsConfig> {
         let mut dns_config = DnsConfig::new();
         for (_svc, nodes) in services.iter() {
             for node in nodes.iter() {
@@ -118,7 +122,7 @@ impl DnsConfigFetcher {
                     continue;
                 }
                 for tag in node.service.tags.iter() {
-                    if let Some((k, v)) = self.parse_d53_tag(tag, &node.node).await? {
+                    if let Some((k, v)) = self.parse_d53_tag(tag, &node.node, autodiscovery)? {
                         dns_config.add(k, v);
                     }
                 }
@@ -128,10 +132,11 @@ impl DnsConfigFetcher {
         Ok(dns_config)
     }
 
-    async fn parse_d53_tag(
+    fn parse_d53_tag(
         &mut self,
         tag: &str,
         node: &catalog::Node,
+        autodiscovery: &AutodiscoveredAddresses,
     ) -> Result<Option<(DnsEntryKey, DnsEntryValue)>> {
         let splits = tag.split(' ').collect::<Vec<_>>();
         if splits.len() != 2 {
@@ -139,14 +144,14 @@ impl DnsConfigFetcher {
         }
 
         let (record_type, target) = match splits[0] {
-            "d53-a" => match self.get_node_ipv4(&node).await? {
+            "d53-a" => match self.get_node_ipv4(&autodiscovery, &node)? {
                 Some(tgt) => (DnsRecordType::A, tgt.to_string()),
                 None => {
                     warn!("Got d53-a tag `{}` but node {} does not appear to have a known public IPv4 address. Tag is ignored.", tag, node.node);
                     return Ok(None);
                 }
             },
-            "d53-aaaa" => match self.get_node_ipv6(&node).await? {
+            "d53-aaaa" => match self.get_node_ipv6(&autodiscovery, &node)? {
                 Some(tgt) => (DnsRecordType::AAAA, tgt.to_string()),
                 None => {
                     warn!("Got d53-aaaa tag `{}` but node {} does not appear to have a known public IPv6 address. Tag is ignored.", tag, node.node);
@@ -174,77 +179,48 @@ impl DnsConfigFetcher {
         )))
     }
 
-    async fn get_node_ipv4(&mut self, node: &catalog::Node) -> Result<Option<Ipv4Addr>> {
-        Self::get_node_ip(&self.consul, "ipv4", &mut self.node_ipv4_cache, node).await
+    fn get_node_ipv4(
+        &mut self,
+        autodiscovery: &AutodiscoveredAddresses,
+        node: &catalog::Node,
+    ) -> Result<Option<Ipv4Addr>> {
+        Self::get_node_ip("ipv4", &autodiscovery.ipv4, node)
     }
 
-    async fn get_node_ipv6(&mut self, node: &catalog::Node) -> Result<Option<Ipv6Addr>> {
-        Self::get_node_ip(&self.consul, "ipv6", &mut self.node_ipv6_cache, node).await
+    fn get_node_ipv6(
+        &mut self,
+        autodiscovery: &AutodiscoveredAddresses,
+        node: &catalog::Node,
+    ) -> Result<Option<Ipv6Addr>> {
+        Self::get_node_ip("ipv6", &autodiscovery.ipv6, node)
     }
 
-    async fn get_node_ip<A>(
-        consul: &Consul,
-        family: &str,
-        cache: &mut HashMap<String, (u64, Option<A>)>,
+    fn get_node_ip<A>(
+        family: &'static str,
+        autodiscovery: &HashMap<String, DiplonatAutodiscoveryResult<A>>,
         node: &catalog::Node,
     ) -> Result<Option<A>>
     where
         A: Serialize + for<'de> Deserialize<'de> + std::fmt::Debug + std::str::FromStr + Copy + Eq,
         <A as std::str::FromStr>::Err: Send + Sync + std::error::Error + 'static,
     {
-        match cache.get(&node.node) {
-            Some((t, a)) if timestamp() <= t + AUTODISCOVERY_CACHE_DURATION => Ok(*a),
-            _ => {
-                let kv_key = format!("diplonat/autodiscovery/{}/{}", family, node.node);
-                let autodiscovery = consul.kv_get(&kv_key).await?;
-
-                if let Some(json) = autodiscovery {
-                    let a = serde_json::from_slice::<DiplonatAutodiscoveryResult<A>>(&json)?;
-                    if timestamp() <= a.timestamp + AUTODISCOVERY_CACHE_DURATION {
-                        if cache.get(&node.node).map(|x| x.1) != Some(a.address) {
-                            info!(
-                                "Got {} address for {} from diplonat autodiscovery: {:?}",
-                                family, node.node, a.address
-                            );
-                        }
-                        cache.insert(node.node.clone(), (a.timestamp, a.address));
-                        return Ok(a.address);
-                    } else {
-                        warn!("{} address for {} from diplonat autodiscovery is outdated (value: {:?}), falling back on value from Consul node meta", family, node.node, a.address);
-                    }
+        match autodiscovery.get(&node.node) {
+            Some(ar) if timestamp() <= ar.timestamp + AUTODISCOVERY_CACHE_DURATION => {
+                Ok(ar.address)
+            }
+            x => {
+                if let Some(ar) = x {
+                    warn!("{} address for {} from diplonat autodiscovery is outdated (value: {:?}), falling back on value from Consul node meta", family, node.node, ar.address);
                 }
 
                 let meta_tag = format!("{}{}", IP_TARGET_METADATA_TAG_PREFIX, family);
-                let a = node.meta.get(&meta_tag).map(|x| x.parse()).transpose()?;
-
-                if cache.get(&node.node).map(|x| x.1) != Some(a) {
-                    info!(
-                        "Got {} address for {} from Consul node meta: {:?}",
-                        family, node.node, a
-                    );
-                }
-                cache.insert(node.node.clone(), (timestamp(), a));
-                Ok(a)
+                let addr = node.meta.get(&meta_tag).map(|x| x.parse()).transpose()?;
+                Ok(addr)
             }
         }
     }
 }
 
-// ---- util for interaction with diplonat ----
-
-#[derive(Serialize, Deserialize, Debug)]
-pub struct DiplonatAutodiscoveryResult<A> {
-    pub timestamp: u64,
-    pub address: Option<A>,
-}
-
-fn timestamp() -> u64 {
-    SystemTime::now()
-        .duration_since(SystemTime::UNIX_EPOCH)
-        .expect("clock error")
-        .as_secs()
-}
-
 // ---- Display impls ----
 
 impl std::fmt::Display for DnsRecordType {
diff --git a/src/main.rs b/src/main.rs
index a8a6b36..9461596 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,6 +5,7 @@ use tokio::select;
 use tokio::sync::watch;
 use tracing::*;
 
+mod autodiscovery;
 mod dns_config;
 mod dns_updater;
 mod provider;