feat(bootstrap): limit address during fetch from network contacts

maidsafe · Dec 11, 2024 · 6fe87e7 · 6fe87e7
1 parent 8014c92
commit 6fe87e7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/ant-bootstrap/src/contacts.rs b/ant-bootstrap/src/contacts.rs
@@ -33,6 +33,8 @@ const MAX_RETRIES_ON_FETCH_FAILURE: usize = 3;
 
 /// Discovers initial peers from a list of endpoints
 pub struct ContactsFetcher {
+    /// The number of addrs to fetch
+    max_addrs: usize,
     /// The list of endpoints
     endpoints: Vec<Url>,
     /// Reqwest Client
@@ -58,12 +60,18 @@ impl ContactsFetcher {
         let request_client = Client::builder().build()?;
 
         Ok(Self {
+            max_addrs: usize::MAX,
             endpoints,
             request_client,
             ignore_peer_id: false,
         })
     }
 
+    /// Set the number of addrs to fetch
+    pub fn set_max_addrs(&mut self, max_addrs: usize) {
+        self.max_addrs = max_addrs;
+    }
+
     /// Create a new struct with the mainnet endpoints
     pub fn with_mainnet_endpoints() -> Result<Self> {
         let mut fetcher = Self::new()?;
@@ -133,6 +141,9 @@ impl ContactsFetcher {
                             .collect::<Vec<_>>()
                     );
                     bootstrap_addresses.append(&mut endpoing_bootstrap_addresses);
+                    if bootstrap_addresses.len() >= self.max_addrs {
+                        break;
+                    }
                 }
                 Err(e) => {
                     warn!("Failed to fetch bootstrap addrs from {}: {}", endpoint, e);

diff --git a/ant-bootstrap/src/initial_peers.rs b/ant-bootstrap/src/initial_peers.rs
@@ -182,7 +182,10 @@ impl PeersArgs {
                 .iter()
                 .map(|url| url.parse::<Url>().map_err(|_| Error::FailedToParseUrl))
                 .collect::<Result<Vec<Url>>>()?;
-            let contacts_fetcher = ContactsFetcher::with_endpoints(addrs)?;
+            let mut contacts_fetcher = ContactsFetcher::with_endpoints(addrs)?;
+            if let Some(count) = count {
+                contacts_fetcher.set_max_addrs(count);
+            }
             let addrs = contacts_fetcher.fetch_bootstrap_addresses().await?;
             bootstrap_addresses.extend(addrs);
 
@@ -197,7 +200,10 @@ impl PeersArgs {
         }
 
         if !self.disable_mainnet_contacts {
-            let contacts_fetcher = ContactsFetcher::with_mainnet_endpoints()?;
+            let mut contacts_fetcher = ContactsFetcher::with_mainnet_endpoints()?;
+            if let Some(count) = count {
+                contacts_fetcher.set_max_addrs(count);
+            }
             let addrs = contacts_fetcher.fetch_bootstrap_addresses().await?;
             bootstrap_addresses.extend(addrs);
         }