karyon_p2p/discovery/kademlia/
mod.rs

1mod bloom;
2mod lookup;
3mod messages;
4mod refresh;
5pub mod routing_table;
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use log::error;
11use rand::seq::IndexedRandom;
12
13use karyon_core::{
14    async_runtime::Executor,
15    async_util::{AsyncQueue, Backoff, TaskGroup, TaskResult},
16    crypto::KeyPair,
17};
18
19use karyon_net::Endpoint;
20
21use crate::{
22    config::Config,
23    discovery::{DiscoveredPeer, Discovery, PeerConnectionEvent},
24    message::{pick_endpoint, PeerAddr, Protocol},
25    monitor::Monitor,
26    protocol::ProtocolFlags,
27    PeerID, Result,
28};
29
30use bloom::BloomRef;
31use lookup::{LookupEndpoints, LookupService};
32use refresh::RefreshService;
33use routing_table::{
34    RoutingTable, CONNECTED_ENTRY, DISCONNECTED_ENTRY, PENDING_ENTRY, UNREACHABLE_ENTRY,
35    UNSTABLE_ENTRY,
36};
37
38/// Bound on the discovered-peer queue between Kademlia and the Node.
39const DISCOVERED_PEER_QUEUE_SIZE: usize = 128;
40
41/// Transports the lookup service can dial. Computed once here so the
42/// rest of the module never needs an inline cfg.
43#[cfg(feature = "quic")]
44pub(crate) const SUPPORTED_LOOKUP_PROTOCOLS: &[Protocol] = &[Protocol::Tcp, Protocol::Quic];
45#[cfg(not(feature = "quic"))]
46pub(crate) const SUPPORTED_LOOKUP_PROTOCOLS: &[Protocol] = &[Protocol::Tcp];
47
48pub struct KademliaDiscovery {
49    /// Routing table
50    table: Arc<RoutingTable>,
51
52    /// Lookup Service
53    lookup_service: Arc<LookupService>,
54
55    /// Refresh Service
56    refresh_service: Arc<RefreshService>,
57
58    /// Discovered peers queued for the Node to dial. Producers
59    /// (connect_loop, manual dial_peers) push; the Node's
60    /// connect_discovered_peers task drains via `recv()`.
61    peer_queue: Arc<AsyncQueue<DiscoveredPeer>>,
62
63    /// Managing spawned tasks.
64    task_group: TaskGroup,
65
66    /// Holds the configuration for the P2P network.
67    config: Arc<Config>,
68
69    /// Local bloom of advertised items. Written by `advertise`,
70    /// snapshotted on every connect-loop iteration to filter
71    /// routing-table entries, and stamped on outgoing PeerMsgs via
72    /// the lookup service.
73    bloom: BloomRef,
74}
75
76impl KademliaDiscovery {
77    /// Creates a new KademliaDiscovery
78    pub fn new(
79        key_pair: &KeyPair,
80        peer_id: &PeerID,
81        config: Arc<Config>,
82        monitor: Arc<Monitor>,
83        ex: Executor,
84    ) -> Arc<Self> {
85        let table = Arc::new(RoutingTable::new(peer_id.0));
86        let bloom = BloomRef::new();
87
88        // Pick the lookup endpoint (tcp/tls/quic) and the refresh
89        // endpoint (udp) from `discovery_endpoints` (any order).
90        let lookup_ep = config
91            .discovery_endpoints
92            .iter()
93            .find(|e| is_lookup_proto(e))
94            .cloned();
95        let refresh_ep = config
96            .discovery_endpoints
97            .iter()
98            .find(|e| e.is_udp())
99            .cloned();
100
101        let refresh_service = Arc::new(RefreshService::new(
102            config.clone(),
103            table.clone(),
104            monitor.clone(),
105            refresh_ep.clone(),
106            ex.clone(),
107        ));
108
109        let lookup_endpoints = LookupEndpoints {
110            listen: config.listen_endpoints.clone(),
111            lookup: lookup_ep,
112            refresh: refresh_ep,
113        };
114        let lookup_service = Arc::new(LookupService::new(
115            key_pair,
116            table.clone(),
117            config.clone(),
118            monitor.clone(),
119            bloom.clone(),
120            lookup_endpoints,
121            ex.clone(),
122        ));
123
124        let task_group = TaskGroup::with_executor(ex);
125
126        let peer_queue = AsyncQueue::new(DISCOVERED_PEER_QUEUE_SIZE);
127
128        Arc::new(Self {
129            refresh_service,
130            lookup_service,
131            table,
132            peer_queue,
133            task_group,
134            config,
135            bloom,
136        })
137    }
138
139    /// This method will attempt to find a peer in the routing table and
140    /// send it as a DiscoveredPeer for the Node to connect to.
141    /// If the routing table is empty, it will start the seeding process.
142    async fn connect_loop(self: Arc<Self>) -> Result<()> {
143        let backoff = Backoff::new(500, self.config.seeding_interval * 1000);
144        loop {
145            let local = self.bloom.snapshot();
146            match self.table.random_entry_filtered(PENDING_ENTRY, &local) {
147                Some(entry) => {
148                    backoff.reset();
149                    let key = entry.key;
150                    let peer = DiscoveredPeer {
151                        peer_id: Some(key.into()),
152                        addrs: entry.addrs.clone(),
153                        discovery_addrs: entry.discovery_addrs.clone(),
154                    };
155                    // Mark as connected to prevent re-picking while
156                    // the Node is still connecting.
157                    self.table.update_entry(&key, CONNECTED_ENTRY);
158                    self.peer_queue.push(peer).await;
159                }
160                None => {
161                    backoff.sleep().await;
162                    self.start_seeding().await;
163                }
164            }
165        }
166    }
167
168    /// Starts seeding process.
169    ///
170    /// This method randomly selects a peer from the routing table and
171    /// attempts to connect to that peer for the initial lookup. If the routing
172    /// table doesn't have an available entry, it will connect to one of the
173    /// provided bootstrap endpoints in the `Config` and initiate the lookup.
174    async fn start_seeding(&self) {
175        match self.table.random_entry(PENDING_ENTRY | CONNECTED_ENTRY) {
176            Some(entry) => {
177                let endpoint =
178                    match pick_endpoint(&entry.discovery_addrs, SUPPORTED_LOOKUP_PROTOCOLS) {
179                        Some(ep) => ep,
180                        None => return,
181                    };
182                let peer_id = Some(entry.key.into());
183                if let Err(err) = self.lookup_service.start_lookup(&endpoint, peer_id).await {
184                    self.table.update_entry(&entry.key, UNSTABLE_ENTRY);
185                    error!("Failed to do lookup: {endpoint}: {err}");
186                }
187            }
188            None => {
189                let peers = &self.config.bootstrap_peers;
190                let shuffled: Vec<_> = peers
191                    .choose_multiple(&mut rand::rng(), peers.len())
192                    .collect();
193                for endpoint in shuffled {
194                    if let Err(err) = self.lookup_service.start_lookup(endpoint, None).await {
195                        error!("Failed to do lookup: {endpoint}: {err}");
196                    }
197                }
198            }
199        }
200    }
201}
202
203#[async_trait]
204impl Discovery for KademliaDiscovery {
205    async fn start(self: Arc<Self>) -> Result<()> {
206        // Start the lookup service
207        self.lookup_service.start().await?;
208        // Start the refresh service
209        self.refresh_service.start().await?;
210
211        // Send manual dial peers as DiscoveredPeer
212        for endpoint in self.config.dial_peers.iter() {
213            if let Some(addr) = PeerAddr::from_endpoint(endpoint, 0) {
214                let peer = DiscoveredPeer {
215                    peer_id: None,
216                    addrs: vec![addr],
217                    discovery_addrs: vec![],
218                };
219                self.peer_queue.push(peer).await;
220            }
221        }
222
223        // Start connect loop
224        self.task_group.spawn_then(
225            {
226                let this = self.clone();
227                async move { this.connect_loop().await }
228            },
229            |res| async move {
230                if let TaskResult::Completed(Err(err)) = res {
231                    error!("Connect loop stopped: {err}");
232                }
233            },
234        );
235
236        Ok(())
237    }
238
239    async fn shutdown(&self) {
240        self.task_group.cancel().await;
241        self.refresh_service.shutdown().await;
242        self.lookup_service.shutdown().await;
243    }
244
245    async fn recv(&self) -> DiscoveredPeer {
246        self.peer_queue.recv().await
247    }
248
249    fn on_event(&self, event: PeerConnectionEvent) {
250        match event {
251            PeerConnectionEvent::Connected(pid) => {
252                self.table.update_entry(&pid.0, CONNECTED_ENTRY);
253            }
254            PeerConnectionEvent::Disconnected(pid) => {
255                self.table.update_entry(&pid.0, DISCONNECTED_ENTRY);
256            }
257            PeerConnectionEvent::ConnectFailed(Some(pid)) => {
258                self.table.update_entry(&pid.0, UNREACHABLE_ENTRY);
259            }
260            PeerConnectionEvent::ConnectFailed(None) => {}
261        }
262    }
263
264    fn advertise(&self, item: &[u8], flags: ProtocolFlags) {
265        self.bloom.add(item, flags);
266    }
267
268    fn find_peers_with(&self, item: &[u8]) -> Vec<DiscoveredPeer> {
269        self.table
270            .entries_with_item(item)
271            .into_iter()
272            .map(|e| DiscoveredPeer {
273                peer_id: Some(e.key.into()),
274                addrs: e.addrs,
275                discovery_addrs: e.discovery_addrs,
276            })
277            .collect()
278    }
279}
280
281/// Returns true if the endpoint can serve the Kademlia lookup service.
282fn is_lookup_proto(e: &Endpoint) -> bool {
283    if e.is_tcp() {
284        return true;
285    }
286    #[cfg(feature = "quic")]
287    if e.is_quic() {
288        return true;
289    }
290    false
291}
karyon_p2p/discovery/kademlia/mod.rs

karyon_p2p/discovery/kademlia/
mod.rs
