Skip to main content

tatara_engine/cluster/
discovery.rs

1use anyhow::Result;
2use mdns_sd::{ServiceDaemon, ServiceEvent, ServiceInfo};
3use std::collections::HashSet;
4use std::net::IpAddr;
5use std::time::Duration;
6use tracing::{debug, info};
7
8const SERVICE_TYPE: &str = "_tatara._tcp.local.";
9
10/// Announces this node's presence via mDNS on the local network.
11pub struct MdnsAnnouncer {
12    daemon: ServiceDaemon,
13}
14
15impl MdnsAnnouncer {
16    pub fn new(
17        instance_name: &str,
18        hostname: &str,
19        ip: IpAddr,
20        gossip_port: u16,
21        http_port: u16,
22        raft_port: u16,
23        cluster_id: &str,
24    ) -> Result<Self> {
25        let daemon = ServiceDaemon::new()
26            .map_err(|e| anyhow::anyhow!("Failed to create mDNS daemon: {}", e))?;
27
28        let properties = [
29            ("cluster", cluster_id),
30            ("gossip_port", &gossip_port.to_string()),
31            ("http_port", &http_port.to_string()),
32            ("raft_port", &raft_port.to_string()),
33        ];
34
35        let host_label = format!("{}.", hostname);
36        let service = ServiceInfo::new(
37            SERVICE_TYPE,
38            instance_name,
39            &host_label,
40            ip,
41            gossip_port,
42            &properties[..],
43        )
44        .map_err(|e| anyhow::anyhow!("Failed to create mDNS service info: {}", e))?;
45
46        daemon
47            .register(service)
48            .map_err(|e| anyhow::anyhow!("Failed to register mDNS service: {}", e))?;
49
50        info!(
51            instance = instance_name,
52            ip = %ip,
53            gossip_port = gossip_port,
54            "mDNS service announced"
55        );
56
57        Ok(Self { daemon })
58    }
59
60    pub fn shutdown(self) -> Result<()> {
61        self.daemon
62            .shutdown()
63            .map(|_| ())
64            .map_err(|e| anyhow::anyhow!("mDNS shutdown error: {}", e))
65    }
66}
67
68/// Discovers tatara peers on the local network via mDNS.
69/// Returns a list of gossip addresses (ip:port).
70pub async fn discover_peers(cluster_id: &str, timeout: Duration) -> Result<Vec<String>> {
71    let daemon = ServiceDaemon::new()
72        .map_err(|e| anyhow::anyhow!("Failed to create mDNS browser: {}", e))?;
73
74    let receiver = daemon
75        .browse(SERVICE_TYPE)
76        .map_err(|e| anyhow::anyhow!("Failed to browse mDNS: {}", e))?;
77
78    let mut peers = HashSet::new();
79    let deadline = tokio::time::Instant::now() + timeout;
80
81    info!(timeout = ?timeout, "Discovering peers via mDNS...");
82
83    loop {
84        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
85        if remaining.is_zero() {
86            break;
87        }
88
89        match tokio::time::timeout(
90            remaining,
91            tokio::task::spawn_blocking({
92                let receiver = receiver.clone();
93                move || receiver.recv_timeout(Duration::from_millis(500))
94            }),
95        )
96        .await
97        {
98            Ok(Ok(Ok(ServiceEvent::ServiceResolved(info)))) => {
99                // Check cluster ID matches
100                let props = info.get_properties();
101                let svc_cluster = props.get_property_val_str("cluster").unwrap_or("");
102
103                if svc_cluster != cluster_id {
104                    debug!(
105                        found_cluster = svc_cluster,
106                        our_cluster = cluster_id,
107                        "Ignoring peer from different cluster"
108                    );
109                    continue;
110                }
111
112                for addr in info.get_addresses() {
113                    let gossip_port = info.get_port();
114                    let peer = format!("{}:{}", addr, gossip_port);
115                    if peers.insert(peer.clone()) {
116                        info!(peer = %peer, "Discovered peer via mDNS");
117                    }
118                }
119            }
120            Ok(Ok(Ok(_))) => {}  // Other events
121            Ok(Ok(Err(_))) => {} // Timeout on recv
122            Ok(Err(_)) | Err(_) => break,
123        }
124    }
125
126    let _ = daemon.shutdown();
127    Ok(peers.into_iter().collect())
128}