Add livelihood checking

src/udp_listen.rs

@@ -2,6 +2,7 @@ use crate::orchestration::PeerHandle;
 use crate::{discovery::DiscoveryMessage, node_service::HealthCheckRequest};
 use std::collections::HashMap;
 use std::net::SocketAddr;
+use std::time::Duration;
 use system_configuration::sys::libc::disconnectx;
 use tokio::net::UdpSocket;
 use tokio::select;
@@ -15,6 +16,7 @@ struct NodeInfo {
     listen_port: u16,
     allowed_peer_ids: Option<Vec<String>>,
     allowed_interfaces: Option<Vec<String>>,
+    discovery_timeout: Duration,
 }
 
 pub async fn listen_for_discovery(
@@ -79,13 +81,24 @@ struct PeerInfo {
     priority: u8,
 }
 
+enum Action {
+    HealthChecks,
+    NewPeer(SocketAddr, DiscoveryMessage),
+}
+
 pub async fn manage_discovery(node_info: NodeInfo) {
     let mut peers: HashMap<String, PeerHandle> = HashMap::new();
     let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<(SocketAddr, DiscoveryMessage)>();
     tokio::spawn(listen_for_discovery(node_info.clone(), tx));
 
-    while let Some((addr, message)) = rx.recv().await {
+    loop {
-        info!("Received discovery message from {}", message.node_id);
+        let action = select! {
+            _ = tokio::time::sleep(node_info.discovery_timeout) => Action::HealthChecks,
+            Some((addr, message)) = rx.recv() => Action::NewPeer(addr, message),
+        };
+
+        match action {
+            Action::NewPeer(addr, message) => {
                 let existing = peers.get(&message.node_id);
                 let insert_new = match existing {
                     None => true,
@@ -137,4 +150,26 @@ pub async fn manage_discovery(node_info: NodeInfo) {
 
                 peers.insert(message.node_id, a);
             }
+
+            Action::HealthChecks => {
+                let mut to_remove = Vec::new();
+                
+                for peer in peers.values() {
+                    let is_healthy = peer.client.lock().await.health_check(HealthCheckRequest::default())
+                        .await
+                        .ok()
+                        .map(|x| x.into_inner().is_healthy)
+                        .unwrap_or(false);
+
+                    if !is_healthy {
+                        to_remove.push(peer.node_id.clone());
+                    }
+                }
+
+                for peer_id in to_remove {
+                    peers.remove(&peer_id);
+                }
+            }
+        }
+    }
 }