feat: replace bootstrap node if bucket full

sn_networking/src/cmd.rs

@@ -9,8 +9,8 @@
 use crate::{
     driver::{PendingGetClosestType, SwarmDriver},
     error::{Error, Result},
-    sort_peers_by_address, GetRecordError, MsgResponder, NetworkEvent, CLOSE_GROUP_SIZE,
-    REPLICATE_RANGE,
+    multiaddr_pop_p2p, sort_peers_by_address, GetRecordError, MsgResponder, NetworkEvent,
+    CLOSE_GROUP_SIZE, REPLICATE_RANGE,
 };
 use bytes::Bytes;
 use libp2p::{
@@ -486,6 +486,15 @@ impl SwarmDriver {
                 };
             }
             SwarmCmd::Dial { addr, sender } => {
+                let mut addr_copy = addr.clone();
+                if let Some(peer_id) = multiaddr_pop_p2p(&mut addr_copy) {
+                    // Only consider the dial peer is bootstrap node when proper PeerId is provided.
+                    if let Some(kbucket) = self.swarm.behaviour_mut().kademlia.kbucket(peer_id) {
+                        let ilog2 = kbucket.range().0.ilog2();
+                        let peers = self.bootstrap_peers.entry(ilog2).or_default();
+                        peers.insert(peer_id);
+                    }
+                }
                 let _ = match self.dial(addr) {
                     Ok(_) => sender.send(Ok(())),
                     Err(e) => sender.send(Err(e.into())),

sn_networking/src/driver.rs

@@ -52,7 +52,7 @@ use sn_protocol::{
     NetworkAddress, PrettyPrintKBucketKey, PrettyPrintRecordKey,
 };
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     fmt::Debug,
     net::SocketAddr,
     num::NonZeroUsize,
@@ -549,6 +549,7 @@ impl NetworkBuilder {
             dialed_peers: CircularVec::new(255),
             is_gossip_handler: false,
             network_discovery: NetworkDiscovery::new(&peer_id),
+            bootstrap_peers: Default::default(),
         };
 
         Ok((
@@ -594,6 +595,7 @@ pub struct SwarmDriver {
     // A list of random `PeerId` candidates that falls into kbuckets,
     // This is to ensure a more accurate network discovery.
     pub(crate) network_discovery: NetworkDiscovery,
+    pub(crate) bootstrap_peers: BTreeMap<Option<u32>, HashSet<PeerId>>,
 }
 
 impl SwarmDriver {

sn_networking/src/event.rs

@@ -277,7 +277,23 @@ impl SwarmDriver {
                             {
                                 let ilog2 = kbucket.range().0.ilog2();
                                 let num_peers = kbucket.num_entries();
-                                (num_peers >= K_VALUE.into(), ilog2)
+                                let mut is_bucket_full = num_peers >= K_VALUE.into();
+
+                                // If the bucket contains any of a bootstrap node,
+                                // consider the bucket is not full and dial back
+                                // so that the bootstrap nodes can be replaced.
+                                if is_bucket_full {
+                                    if let Some(peers) = self.bootstrap_peers.get(&ilog2) {
+                                        if kbucket
+                                            .iter()
+                                            .any(|entry| peers.contains(entry.node.key.preimage()))
+                                        {
+                                            is_bucket_full = false;
+                                        }
+                                    }
+                                }
+
+                                (is_bucket_full, ilog2)
                             } else {
                                 // Function will return `None` if the given key refers to self
                                 // hence return true to skip further action.
@@ -308,6 +324,8 @@ impl SwarmDriver {
 
                         // If we are not local, we care only for peers that we dialed and thus are reachable.
                         if self.local || has_dialed && peer_is_agent {
+                            self.remove_bootstrap_from_full(peer_id);
+
                             trace!(%peer_id, ?addrs, "identify: attempting to add addresses to routing table");
 
                             // Attempt to add the addresses to the routing table.
@@ -1245,6 +1263,32 @@ impl SwarmDriver {
 
         false
     }
+
+    // if target bucket is full, remove a bootstrap node if presents.
+    fn remove_bootstrap_from_full(&mut self, peer_id: PeerId) {
+        let mut shall_removed = None;
+
+        if let Some(kbucket) = self.swarm.behaviour_mut().kademlia.kbucket(peer_id) {
+            if kbucket.num_entries() >= K_VALUE.into() {
+                if let Some(peers) = self.bootstrap_peers.get(&kbucket.range().0.ilog2()) {
+                    for peer_entry in kbucket.iter() {
+                        if peers.contains(peer_entry.node.key.preimage()) {
+                            shall_removed = Some(*peer_entry.node.key.preimage());
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+        if let Some(to_be_removed_bootstrap) = shall_removed {
+            trace!("Bootstrap node {to_be_removed_bootstrap:?} to be replaced by peer {peer_id:?}");
+            let _entry = self
+                .swarm
+                .behaviour_mut()
+                .kademlia
+                .remove_peer(&to_be_removed_bootstrap);
+        }
+    }
 }
 
 /// Helper function to print formatted connection role info.

sn_networking/src/lib.rs

@@ -159,6 +159,7 @@ impl Network {
     }
 
     /// Dial the given peer at the given address.
+    /// This function will only be called for the bootstrap nodes.
     pub async fn dial(&self, addr: Multiaddr) -> Result<()> {
         let (sender, receiver) = oneshot::channel();
         self.send_swarm_cmd(SwarmCmd::Dial { addr, sender })?;