wip

src/main.rs

@@ -1,3 +1,132 @@
+use futures::{Future, Stream};
+use maplit::hashmap;
+use mdns::{Record, RecordKind};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::mpsc;
+use std::sync::mpsc::Sender;
+use std::thread;
+use std::thread::sleep;
+use std::time::Instant;
+use std::{net::IpAddr, time::Duration};
+
+/// The hostname of the devices we are searching for.
+const SERVICE_NAME: &'static str = "_prometheus-http._tcp.local";
+
+struct Service {
+    name: String,
+    addr: IpAddr,
+    port: u16,
+    last_seen: Instant,
+}
+
+#[derive(Serialize, Deserialize)]
+struct PrometheusService {
+    targets: Vec<String>,
+    labels: HashMap<String, String>,
+}
+
+impl From<&Service> for PrometheusService {
+    fn from(service: &Service) -> Self {
+        PrometheusService {
+            targets: vec![format!("{}:{}", service.addr, service.port)],
+            labels: hashmap! {
+                "name".to_string() => service.name.clone()
+            },
+        }
+    }
+}
+
+const TIMEOUT: Duration = Duration::from_secs(60);
+
 fn main() {
-    println!("Hello, world!");
+    let (tx, rx) = mpsc::channel();
+
+    thread::spawn(move || {
+        discover(tx);
+    });
+
+    let mut services: HashMap<IpAddr, Service> = HashMap::new();
+
+    loop {
+        let start_count = services.len();
+
+        while let Ok(service) = rx.try_recv() {
+            services.insert(service.addr, service);
+        }
+
+        let added_count = services.len();
+
+        services.retain(|_, service| Instant::now().duration_since(service.last_seen) < TIMEOUT);
+
+        let removed_count = services.len();
+
+        if start_count != added_count || added_count != removed_count {
+            let output: Vec<PrometheusService> =
+                services.iter().map(|(_, service)| service.into()).collect();
+            println!("{}", serde_json::to_string(&output).unwrap());
+        }
+
+        sleep(Duration::from_secs(15));
+    }
+}
+
+fn discover(tx: Sender<Service>) {
+    tokio::run(
+        mdns::discover::all(SERVICE_NAME, Duration::from_secs(15))
+            .unwrap()
+            .for_each(move |response| {
+                if response
+                    .records()
+                    .any(|record| record.name.as_str() == SERVICE_NAME)
+                {
+                    let addr = response.records().filter_map(self::to_ip_addr).next();
+                    let port = response.records().filter_map(self::to_port).next();
+                    let name = response.records().filter_map(self::to_name).next();
+
+                    dbg!(response);
+
+                    if let (Some(addr), Some(name), Some(port)) = (addr, name, port) {
+                        let _ = tx.send(Service {
+                            name,
+                            addr,
+                            port,
+                            last_seen: Instant::now(),
+                        });
+                    }
+                }
+
+                Ok(())
+            })
+            .map_err(|e| eprintln!("{:?}", e)),
+    );
+}
+
+fn to_ip_addr(record: &Record) -> Option<IpAddr> {
+    match record.kind {
+        RecordKind::A(addr) => Some(addr.into()),
+        RecordKind::AAAA(addr) => Some(addr.into()),
+        _ => None,
+    }
+}
+
+fn to_port(record: &Record) -> Option<u16> {
+    match record.kind {
+        RecordKind::SRV { port, .. } if record.name.contains(SERVICE_NAME) => Some(port),
+        _ => None,
+    }
+}
+
+fn to_name(record: &Record) -> Option<String> {
+    if let RecordKind::TXT(txt) = &record.kind {
+        for pair in txt {
+            let mut parts = pair.split('=');
+            if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
+                if key == "name" {
+                    return Some(value.to_string());
+                }
+            }
+        }
+    }
+    None
 }