diff --git a/ant-bootstrap/src/contacts.rs b/ant-bootstrap/src/contacts.rs index 17a08d514d..ca1faf9ec9 100644 --- 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, /// 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 { let mut fetcher = Self::new()?; @@ -133,6 +141,9 @@ impl ContactsFetcher { .collect::>() ); 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 index a6f52b5012..55b3f78e16 100644 --- 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::().map_err(|_| Error::FailedToParseUrl)) .collect::>>()?; - 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); }