partition_sim/
consul.rs

1use std::{
2    net::IpAddr,
3    process::{Command, Stdio},
4};
5
6/// The main Consul api that the supervisor uses to query a list of peer
7/// ip addresses and ports they registered their services with.
8///
9pub fn query_consul_for_peers(
10    dns_addr: &str,
11    dns_port: u16,
12    service_name: &str,
13) -> crate::Result<Vec<(IpAddr, u16)>> {
14    let child = Command::new("dig")
15        .arg("+short")
16        .arg(format!("@{}", dns_addr))
17        .arg("-p")
18        .arg(format!("{}", dns_port))
19        .arg("-t")
20        .arg("SRV")
21        .arg(format!("{}.service.consul", service_name))
22        .stdout(Stdio::piped())
23        .spawn()?;
24
25    let child = Command::new("awk")
26        .arg("{print $4,$3;}")
27        .stdin(child.stdout.expect("No stdout of dig command?"))
28        .stdout(Stdio::piped())
29        .output()?;
30
31    let output = String::from_utf8_lossy(&child.stdout);
32
33    let dns_names_and_ports: Vec<_> = output
34        .lines()
35        .map(|line| {
36            let mut parts = line.split_whitespace();
37            let dns_name = parts.next().unwrap();
38            let port = parts.next().unwrap();
39            (dns_name, port)
40        })
41        .collect();
42
43    let dns_names = dns_names_and_ports
44        .iter()
45        .map(|&(dns_name, _)| dns_name)
46        .collect::<Vec<_>>();
47
48    let output = Command::new("dig")
49        .arg("+short")
50        .arg(format!("@{}", dns_addr))
51        .arg("-p")
52        .arg(format!("{}", dns_port))
53        .args(&dns_names)
54        .output()?;
55
56    let output = String::from_utf8_lossy(&output.stdout);
57
58    Ok(output
59        .lines()
60        .map(|line| line.to_string())
61        .zip(dns_names_and_ports.into_iter().map(|(_, port)| port))
62        .map(|(ip_addr_raw, port)| {
63            let ip_addr = ip_addr_raw.parse::<IpAddr>().unwrap();
64            (ip_addr, port.parse::<u16>().unwrap())
65        })
66        .collect())
67}
68
69#[cfg(test)]
70mod tests {
71    use std::collections::HashSet;
72    use std::net::IpAddr;
73
74    use super::query_consul_for_peers;
75
76    /// This test must be run after `docker-compose run -d` has been run in the root of the project.
77    #[test]
78    pub fn stuff() {
79        let res = query_consul_for_peers("127.0.0.1", 8600, "test-node-base").unwrap();
80        println!("{:#?}", res);
81        let observed: HashSet<(IpAddr, u16)> =
82            std::collections::HashSet::from_iter(res.into_iter());
83        let expected = std::collections::HashSet::from_iter(vec![
84            ("192.168.192.4".parse().unwrap(), 9001),
85            ("192.168.192.6".parse().unwrap(), 9001),
86            ("192.168.192.7".parse().unwrap(), 9001),
87            ("192.168.192.8".parse().unwrap(), 9001),
88            ("192.168.192.5".parse().unwrap(), 9001),
89        ]);
90        assert_eq!(observed, expected);
91    }
92}
partition_sim/consul.rs

partition_sim/
consul.rs
