urma_runtime/p2p/
discovery.rs

1use crate::config::{P2P_ADDRESS_POOL, P2P_REDIAL_BACKOFF_SECS, P2P_SEED_INTERVAL_SECS};
2use std::collections::{HashSet, VecDeque};
3use std::net::{IpAddr, SocketAddr};
4use std::time::{Duration, Instant};
5
6#[derive(Debug, PartialEq, Eq)]
7pub enum Next {
8    Address(SocketAddr),
9    Exhausted,
10}
11
12enum Seeded {
13    Never,
14    At(Instant),
15}
16
17enum Drought {
18    Clear,
19    Since(Instant),
20}
21
22pub struct Discovery {
23    pool: VecDeque<SocketAddr>,
24    tried: HashSet<SocketAddr>,
25    banned: HashSet<IpAddr>,
26    redial: VecDeque<(Instant, SocketAddr)>,
27    skip_v6: bool,
28    seeded: Seeded,
29    drought: Drought,
30}
31
32impl Discovery {
33    pub fn new() -> Self {
34        Self {
35            pool: VecDeque::new(),
36            tried: HashSet::new(),
37            banned: HashSet::new(),
38            redial: VecDeque::new(),
39            skip_v6: false,
40            seeded: Seeded::Never,
41            drought: Drought::Clear,
42        }
43    }
44
45    pub fn is_empty(&self) -> bool {
46        self.pool.is_empty()
47    }
48
49    pub fn len(&self) -> usize {
50        self.pool.len()
51    }
52
53    pub fn seeding_due(&self, now: Instant) -> bool {
54        if !self.pool.is_empty() {
55            return false;
56        }
57        match self.seeded {
58            Seeded::Never => true,
59            Seeded::At(when) => {
60                now.duration_since(when) >= Duration::from_secs(P2P_SEED_INTERVAL_SECS)
61            }
62        }
63    }
64
65    pub fn mark_seeded(&mut self, now: Instant) {
66        self.seeded = Seeded::At(now);
67    }
68
69    pub fn drought_warns(&mut self, now: Instant) -> bool {
70        if !self.pool.is_empty() {
71            self.drought = Drought::Clear;
72            return false;
73        }
74        match self.drought {
75            Drought::Clear => {
76                self.drought = Drought::Since(now);
77                false
78            }
79            Drought::Since(start) => {
80                if now.duration_since(start) >= Duration::from_secs(P2P_SEED_INTERVAL_SECS) {
81                    self.drought = Drought::Since(now);
82                    return true;
83                }
84                false
85            }
86        }
87    }
88
89    fn admissible(&self, address: &SocketAddr) -> bool {
90        !(self.skip_v6 && address.is_ipv6())
91            && !self.tried.contains(address)
92            && !self.banned.contains(&address.ip())
93            && !self.pool.contains(address)
94            && self.pool.len() < P2P_ADDRESS_POOL
95    }
96
97    pub fn offer_seeded(&mut self, addresses: Vec<SocketAddr>) {
98        for address in addresses.into_iter().rev() {
99            if self.admissible(&address) {
100                self.pool.push_front(address);
101            }
102        }
103    }
104
105    pub fn readmit(&mut self, addresses: Vec<SocketAddr>) {
106        for address in &addresses {
107            self.tried.remove(address);
108        }
109        self.offer_seeded(addresses);
110    }
111
112    pub fn offer_gossip(&mut self, address: SocketAddr) {
113        if self.admissible(&address) {
114            self.pool.push_back(address);
115        }
116    }
117
118    pub fn dropped(&mut self, address: SocketAddr, now: Instant) {
119        if self.banned.contains(&address.ip()) || self.redial.iter().any(|entry| entry.1 == address)
120        {
121            return;
122        }
123        let due = now + Duration::from_secs(P2P_REDIAL_BACKOFF_SECS);
124        self.redial.push_back((due, address));
125    }
126
127    pub fn redial_pending(&self) -> usize {
128        self.redial.len()
129    }
130
131    pub fn next(&mut self, now: Instant) -> Next {
132        let ready = self.redial.front().copied().filter(|entry| entry.0 <= now);
133        for entry in ready.iter().copied() {
134            self.redial.pop_front();
135            let address = entry.1;
136            if !self.banned.contains(&address.ip()) && !(self.skip_v6 && address.is_ipv6()) {
137                self.tried.insert(address);
138                return Next::Address(address);
139            }
140        }
141        let Some(address) = self.pool.pop_front() else {
142            return Next::Exhausted;
143        };
144        self.tried.insert(address);
145        Next::Address(address)
146    }
147
148    pub fn ban(&mut self, ip: IpAddr) {
149        self.banned.insert(ip);
150        self.pool.retain(|address| address.ip() != ip);
151        self.redial.retain(|entry| entry.1.ip() != ip);
152    }
153
154    pub fn unreachable_v6(&mut self) {
155        self.skip_v6 = true;
156        self.pool.retain(|address| !address.is_ipv6());
157        self.redial.retain(|entry| !entry.1.is_ipv6());
158    }
159
160    pub fn skips_v6(&self) -> bool {
161        self.skip_v6
162    }
163}
urma_runtime/p2p/discovery.rs

urma_runtime/p2p/
discovery.rs
