diff --git a/src/rpc/consul.rs b/src/rpc/consul.rs index 760e9fcb..a9e0ec93 100644 --- a/src/rpc/consul.rs +++ b/src/rpc/consul.rs @@ -128,36 +128,44 @@ impl ConsulDiscovery { // ---- READING FROM CONSUL CATALOG ---- pub async fn get_consul_nodes(&self) -> Result, ConsulError> { - let url = format!( - "{}/v1/catalog/service/{}", - self.config.consul_http_addr, self.config.service_name - ); - - let http = self.client.get(&url).send().await?; - let entries: Vec = http.json().await?; - let mut ret = vec![]; - for ent in entries { - let ip = ent.address.parse::().ok(); - let pubkey = ent - .meta - .get(&format!("{}-pubkey", META_PREFIX)) - .and_then(|k| hex::decode(k).ok()) - .and_then(|k| NodeID::from_slice(&k[..])); - if let (Some(ip), Some(pubkey)) = (ip, pubkey) { - ret.push((pubkey, SocketAddr::new(ip, ent.service_port))); - } else { - warn!( - "Could not process node spec from Consul: {:?} (invalid IP address or node ID/pubkey)", - ent - ); + + let dcs_to_query: Vec> = match &self.config.datacenters { + Some(dcs) if !dcs.is_empty() => dcs.iter().map(|dc| Some(dc.as_str())).collect(), + _ => vec![None], + }; + + for dc in dcs_to_query { + let url = match dc { + Some(datacenter) => format!( + "{}/v1/catalog/service/{}?dc={}", + self.config.consul_http_addr, self.config.service_name, datacenter + ), + None => format!( + "{}/v1/catalog/service/{}", + self.config.consul_http_addr, self.config.service_name + ), + }; + + let http = self.client.get(&url).send().await?; + let entries: Vec = http.json().await?; + + for ent in entries { + let ip = ent.address.parse::().ok(); + let pubkey = ent + .meta + .get(&format!("{}-pubkey", META_PREFIX)) + .and_then(|k| hex::decode(k).ok()) + .and_then(|k| NodeID::from_slice(&k[..])); + if let (Some(ip), Some(pubkey)) = (ip, pubkey) { + ret.push((pubkey, SocketAddr::new(ip, ent.service_port))); + } } } - debug!("Got nodes from Consul: {:?}", ret); + debug!("Got {} nodes from Consul", ret.len()); Ok(ret) } - // ---- PUBLISHING TO CONSUL CATALOG ---- pub async fn publish_consul_service( diff --git a/src/util/config.rs b/src/util/config.rs index a4521eb5..29047f8c 100644 --- a/src/util/config.rs +++ b/src/util/config.rs @@ -257,6 +257,8 @@ pub struct ConsulDiscoveryConfig { /// Additional service metadata to add #[serde(default)] pub meta: Option>, + #[serde(default)] + pub datacenters: Option>, } #[derive(Deserialize, Debug, Clone)]