Make it work in nomad

src/consul.rs

@@ -2,10 +2,14 @@
 
 use std::collections::HashMap;
 
+use log::debug;
+use log::info;
+use log::trace;
 use log::warn;
 use reqwest::header::HeaderMap;
 use serde::{Deserialize, Serialize};
 
+use crate::config::Config;
 use crate::config::ConfigProvider;
 use crate::config::Target;
 
@@ -17,9 +21,13 @@ pub struct ConsulConfigProvider {
 }
 
 impl ConsulConfigProvider {
-    pub fn new() -> Self {
+    pub fn new(config: Option<&Config>) -> Self {
         Self {
-            consul_config: ConsulConfig::from_env(),
+            consul_config: if let Some(config) = config {
+                ConsulConfig::from_config_or_env(config)
+            } else {
+                ConsulConfig::from_env()
+            },
             interval: tokio::time::interval(tokio::time::Duration::from_secs(10)),
         }
     }
@@ -28,16 +36,21 @@ impl ConsulConfigProvider {
 #[async_trait::async_trait]
 impl ConfigProvider for ConsulConfigProvider {
     async fn get_targets(&self) -> Result<Vec<Target>> {
+        info!("Getting targets from consul");
         let mut targets: Vec<Target> = Vec::new();
 
+        debug!("Calling consul_get_services");
         let services = consul_get_services(&self.consul_config).await?;
 
         // Find consul services and tags
         // Format of tags:  rustocat:udp:port
         //                  rustocat:tcp:port
+
+        debug!("Iterating over services: {:?}", services);
         for (name, tags) in services {
             for tag in tags {
                 if tag.starts_with("rustocat") {
+                    trace!("Found rustocat tag: {}", tag);
                     let parts = tag.split(":").collect::<Vec<&str>>();
                     if parts.len() != 3 {
                         warn!("Invalid tag: {} on service {}", tag, name);
@@ -45,17 +58,19 @@ impl ConfigProvider for ConsulConfigProvider {
                     }
 
                     let port = parts[2];
+                    trace!("Getting nodes for service: {}", name);
                     let nodes = consul_get_service_nodes(&self.consul_config, &name).await?;
 
                     let mut t = vec![];
                     for node in nodes {
-                        t.push(format!("{}:{}", node.Service.Address, node.Service.Port));
+                        t.push(format!("{}:{}", node.ServiceAddress, node.ServicePort));
                     }
                     let target = Target {
                         udp: Some(parts[1] == "udp"),
                         source: format!("0.0.0.0:{}", port),
                         targets: t,
                     };
+                    trace!("Adding target: {:?}", target);
                     targets.push(target);
                 }
             }
@@ -65,6 +80,7 @@ impl ConfigProvider for ConsulConfigProvider {
     }
 
     async fn wait_for_change(&mut self) -> Result<()> {
+        info!("Waiting for consul config change");
         self.interval.tick().await;
 
         Ok(())
@@ -74,7 +90,7 @@ impl ConfigProvider for ConsulConfigProvider {
 async fn consul_get_services(config: &ConsulConfig) -> Result<HashMap<String, Vec<String>>> {
     let mut headers = HeaderMap::new();
     if let Some(token) = config.token.clone() {
-        headers.insert("X-Consul-Token", token.parse().unwrap());
+        headers.insert("X-Consul-Token", token.parse()?);
     }
 
     return Ok(reqwest::Client::new()
@@ -86,21 +102,23 @@ async fn consul_get_services(config: &ConsulConfig) -> Result<HashMap<String, Ve
         .await?);
 }
 
-async fn consul_get_service_nodes(
-    config: &ConsulConfig,
-    service: &str,
-) -> Result<Vec<ServiceEntry>> {
+async fn consul_get_service_nodes(config: &ConsulConfig, service: &str) -> Result<Vec<Node>> {
     let mut headers = HeaderMap::new();
     if let Some(token) = config.token.clone() {
-        headers.insert("X-Consul-Token", token.parse().unwrap());
+        headers.insert("X-Consul-Token", token.parse()?);
     }
 
+    trace!(
+        "Calling consul_get_service_nodes: {}/v1/catalog/service/{service}",
+        config.baseurl
+    );
+
     return Ok(reqwest::Client::new()
-        .get(format!("{}/v1/catalog/services/{service}", config.baseurl))
+        .get(format!("{}/v1/catalog/service/{service}", config.baseurl))
         .headers(headers)
         .send()
         .await?
-        .json::<Vec<ServiceEntry>>()
+        .json::<Vec<Node>>()
         .await?);
 }
 
@@ -136,7 +154,8 @@ struct HealthCheck {
 struct Node {
     ID: String,
     Node: String,
-    Address: String,
+    ServiceAddress: String,
+    ServicePort: u16,
     Datacenter: Option<String>,
     TaggedAddresses: Option<HashMap<String, String>>,
     Meta: Option<HashMap<String, String>>,
@@ -166,4 +185,20 @@ impl ConsulConfig {
             token: option_env!("CONSUL_HTTP_TOKEN").map(|s| s.to_string()),
         }
     }
+
+    fn from_config_or_env(config: &Config) -> Self {
+        let baseurl = match config.consul_http_addr {
+            Some(ref s) => s.clone(),
+            None => option_env!("CONSUL_HTTP_ADDR")
+                .expect("CONSUL_HTTP_ADDR not set")
+                .to_string(),
+        };
+
+        let token = match config.consul_http_token {
+            Some(ref s) => Some(s.clone()),
+            None => option_env!("CONSUL_HTTP_TOKEN").map(|s| s.to_string()),
+        };
+
+        Self { baseurl, token }
+    }
 }