1  use std::collections::HashMap;
 2  use std::sync::Arc;
 3  
 4  use iroh::PublicKey;
 5  use rand::Rng;
 6  use serde::{Deserialize, Serialize};
 7  use tokio::sync::RwLock;
 8  
 9  #[derive(Debug, Clone, Serialize, Deserialize)]
10  pub struct NodeInfo {
11      pub name: String,
12      pub id: PublicKey,
13  }
14  
15  #[derive(Clone)]
16  pub struct NodeMap {
17      inner: Arc<RwLock<HashMap<String, NodeInfo>>>,
18  }
19  
20  impl NodeMap {
21      /// Create a new node map containing only our own entry.
22      pub fn new(name: String, id: PublicKey) -> Self {
23          let mut map = HashMap::new();
24          map.insert(name.clone(), NodeInfo { name, id });
25          Self {
26              inner: Arc::new(RwLock::new(map)),
27          }
28      }
29  
30      /// Insert or update a node.
31      pub async fn insert(&self, node: NodeInfo) {
32          self.inner.write().await.insert(node.name.clone(), node);
33      }
34  
35      /// Pick a random node, optionally excluding one by public key.
36      pub async fn pick_random(&self, exclude: Option<&PublicKey>) -> Option<NodeInfo> {
37          let map = self.inner.read().await;
38          let candidates: Vec<_> = map
39              .values()
40              .filter(|n| exclude.is_none_or(|pk| n.id != *pk))
41              .collect();
42          if candidates.is_empty() {
43              None
44          } else {
45              let idx = rand::rng().random_range(0..candidates.len());
46              Some(candidates[idx].clone())
47          }
48      }
49  
50      /// Get a node by name.
51      pub async fn get(&self, name: &str) -> Option<NodeInfo> {
52          self.inner.read().await.get(name).cloned()
53      }
54  
55      /// Return all known nodes.
56      pub async fn list(&self) -> Vec<NodeInfo> {
57          self.inner.read().await.values().cloned().collect()
58      }
59  
60      /// Collect gossip candidate public keys, excluding `own_name`.
61      pub async fn gossip_candidates(
62          &self,
63          own_name: &str,
64          initial_peers: &[PublicKey],
65      ) -> Vec<PublicKey> {
66          let map = self.inner.read().await;
67          let mut keys: Vec<PublicKey> = Vec::new();
68  
69          // Add initial peers.
70          for pk in initial_peers {
71              keys.push(*pk);
72          }
73  
74          // Add all other known nodes not already in the list.
75          for node in map.values() {
76              if node.name == own_name {
77                  continue;
78              }
79              if keys.contains(&node.id) {
80                  continue;
81              }
82              keys.push(node.id);
83          }
84  
85          keys
86      }
87  }