thenodes/network/
peer_store.rs

1// src/network/peer_store.rs
2// In-memory peer candidate and reconnect-attempt store.
3
4use crate::config::Config;
5use rand::seq::IteratorRandom;
6use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet};
8use std::net::SocketAddr;
9use std::sync::Arc;
10use tokio::fs;
11use tokio::io::AsyncWriteExt as _;
12use tokio::sync::RwLock;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15pub enum PeerSource {
16    Bootstrap,
17    Handshake,
18    Gossip,
19    Manual,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct PeerRecord {
24    pub addr: SocketAddr,
25    pub source: PeerSource,
26    pub failures: u32,
27    pub last_success_epoch: Option<u64>,
28    /// Optional node identifier (populated on successful handshake)
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub node_id: Option<String>,
31    /// Optional capability flags advertised by peer
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub capabilities: Option<Vec<String>>,
34}
35
36#[derive(Clone)]
37pub struct PeerStore {
38    inner: Arc<RwLock<HashMap<SocketAddr, PeerRecord>>>,
39    reconnect_attempts: Arc<RwLock<HashMap<SocketAddr, u32>>>,
40}
41
42impl Default for PeerStore {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl PeerStore {
49    pub fn new() -> Self {
50        Self {
51            inner: Arc::new(RwLock::new(HashMap::new())),
52            reconnect_attempts: Arc::new(RwLock::new(HashMap::new())),
53        }
54    }
55
56    pub async fn insert(&self, addr: SocketAddr, source: PeerSource) {
57        let mut map = self.inner.write().await;
58        map.entry(addr).or_insert(PeerRecord {
59            addr,
60            source,
61            failures: 0,
62            last_success_epoch: None,
63            node_id: None,
64            capabilities: None,
65        });
66    }
67
68    pub async fn bulk_insert<I: IntoIterator<Item = SocketAddr>>(&self, it: I, source: PeerSource) {
69        let mut map = self.inner.write().await;
70        for addr in it {
71            map.entry(addr).or_insert(PeerRecord {
72                addr,
73                source,
74                failures: 0,
75                last_success_epoch: None,
76                node_id: None,
77                capabilities: None,
78            });
79        }
80    }
81
82    pub async fn mark_success(&self, addr: &SocketAddr) {
83        let mut map = self.inner.write().await;
84        if let Some(rec) = map.get_mut(addr) {
85            rec.last_success_epoch = Some(Self::epoch());
86            rec.failures = 0;
87        }
88        drop(map);
89        self.reset_reconnect(addr).await;
90    }
91
92    /// Mark peer as successfully connected and update metadata from HELLO
93    pub async fn mark_success_with_meta(
94        &self,
95        addr: &SocketAddr,
96        node_id: Option<String>,
97        capabilities: Option<Vec<String>>,
98    ) {
99        let mut map = self.inner.write().await;
100        if let Some(rec) = map.get_mut(addr) {
101            rec.last_success_epoch = Some(Self::epoch());
102            rec.failures = 0;
103            if node_id.is_some() {
104                rec.node_id = node_id;
105            }
106            if capabilities.is_some() {
107                rec.capabilities = capabilities;
108            }
109        } else {
110            // Insert new record if not present
111            map.insert(
112                *addr,
113                PeerRecord {
114                    addr: *addr,
115                    source: PeerSource::Handshake,
116                    failures: 0,
117                    last_success_epoch: Some(Self::epoch()),
118                    node_id,
119                    capabilities,
120                },
121            );
122        }
123        drop(map);
124        self.reset_reconnect(addr).await;
125    }
126
127    pub async fn mark_failure(&self, addr: &SocketAddr) {
128        let mut map = self.inner.write().await;
129        if let Some(rec) = map.get_mut(addr) {
130            rec.failures = rec.failures.saturating_add(1);
131        }
132    }
133
134    /// Reserve the next reconnect attempt number, starting at one.
135    ///
136    /// Returns `None` when the configured consecutive-attempt budget is exhausted.
137    pub async fn next_reconnect_attempt(
138        &self,
139        addr: &SocketAddr,
140        max_attempts: u32,
141    ) -> Option<u32> {
142        let mut attempts = self.reconnect_attempts.write().await;
143        let current = attempts.get(addr).copied().unwrap_or(0);
144        if current >= max_attempts {
145            return None;
146        }
147        let next = current.saturating_add(1);
148        attempts.insert(*addr, next);
149        Some(next)
150    }
151
152    pub async fn reset_reconnect(&self, addr: &SocketAddr) {
153        self.reconnect_attempts.write().await.remove(addr);
154    }
155
156    pub async fn reconnect_attempts(&self, addr: &SocketAddr) -> u32 {
157        self.reconnect_attempts
158            .read()
159            .await
160            .get(addr)
161            .copied()
162            .unwrap_or(0)
163    }
164
165    pub async fn sample(&self, k: usize, exclude: &HashSet<SocketAddr>) -> Vec<SocketAddr> {
166        let map = self.inner.read().await;
167        let mut rng = rand::thread_rng();
168        map.values()
169            .filter(|r| !exclude.contains(&r.addr))
170            .map(|r| r.addr)
171            .choose_multiple(&mut rng, k)
172    }
173
174    pub async fn all(&self) -> Vec<SocketAddr> {
175        self.inner.read().await.keys().cloned().collect()
176    }
177
178    // Persistence API
179    const DEFAULT_FILENAME: &str = "peers.json";
180
181    pub async fn load_from_file(path: &str, ttl_secs: u64, max_entries: usize) -> Self {
182        let store = Self::new();
183        if let Ok(bytes) = fs::read(path).await {
184            let entries: Vec<PeerRecord> = serde_json::from_slice(&bytes).unwrap_or_default();
185            let now = Self::epoch();
186            let mut count = 0usize;
187            for rec in entries.into_iter() {
188                let age_ok = rec
189                    .last_success_epoch
190                    .map(|t| now.saturating_sub(t) <= ttl_secs)
191                    .unwrap_or(true);
192                if age_ok {
193                    store.insert(rec.addr, rec.source).await;
194                    count += 1;
195                    if count >= max_entries {
196                        break;
197                    }
198                }
199            }
200        }
201        store
202    }
203
204    pub async fn save_to_file(&self, path: &str) -> Result<(), std::io::Error> {
205        let map = self.inner.read().await;
206        let mut entries: Vec<PeerRecord> = map.values().cloned().collect();
207        // Prefer recent successes first
208        entries.sort_by_key(|r| std::cmp::Reverse(r.last_success_epoch.unwrap_or(0)));
209        let json = serde_json::to_vec_pretty(&entries).map_err(std::io::Error::other)?;
210        if let Some(parent) = std::path::Path::new(path).parent() {
211            fs::create_dir_all(parent).await.ok();
212        }
213        let mut f = fs::File::create(path).await?;
214        f.write_all(&json).await?;
215        Ok(())
216    }
217
218    pub async fn spawn_periodic_save(self, path: String, interval_secs: u64) {
219        tokio::spawn(async move {
220            let mut ticker = tokio::time::interval(std::time::Duration::from_secs(interval_secs));
221            loop {
222                ticker.tick().await;
223                let _ = self.save_to_file(&path).await;
224            }
225        });
226    }
227
228    /// Construct and optionally start periodic saves based on runtime config.
229    /// Uses `node.state_dir` to derive path when not explicitly provided.
230    pub async fn from_config(cfg: &Config) -> Self {
231        if let Some(p) = cfg.network.as_ref().and_then(|n| n.persistence.as_ref()) {
232            let enabled = p.enabled.unwrap_or(false);
233            let max_entries = p.max_entries.unwrap_or(1024);
234            let ttl_secs = p.ttl_secs.unwrap_or(7 * 24 * 3600);
235            let base_dir = cfg
236                .node
237                .as_ref()
238                .and_then(|n| n.state_dir.clone())
239                .unwrap_or_else(|| "data".to_string());
240            let path = p
241                .path
242                .clone()
243                .unwrap_or_else(|| format!("{}/{}", base_dir, Self::DEFAULT_FILENAME));
244            let store = Self::load_from_file(&path, ttl_secs, max_entries).await;
245            if enabled {
246                let interval = p.save_interval_secs.unwrap_or(60);
247                store.clone().spawn_periodic_save(path, interval).await;
248            }
249            store
250        } else {
251            Self::new()
252        }
253    }
254
255    /// Save to derived path if persistence enabled in config.
256    pub async fn save_if_enabled(&self, cfg: &Config) {
257        if let Some(p) = cfg.network.as_ref().and_then(|n| n.persistence.as_ref()) {
258            if p.enabled.unwrap_or(false) {
259                let base_dir = cfg
260                    .node
261                    .as_ref()
262                    .and_then(|n| n.state_dir.clone())
263                    .unwrap_or_else(|| "data".to_string());
264                let path = p
265                    .path
266                    .clone()
267                    .unwrap_or_else(|| format!("{}/{}", base_dir, Self::DEFAULT_FILENAME));
268                let _ = self.save_to_file(&path).await;
269            }
270        }
271    }
272
273    fn epoch() -> u64 {
274        use std::time::{SystemTime, UNIX_EPOCH};
275        SystemTime::now()
276            .duration_since(UNIX_EPOCH)
277            .unwrap()
278            .as_secs()
279    }
280}
thenodes/network/peer_store.rs

thenodes/network/
peer_store.rs
