net/hosts/fetch_with_schemes(): limit added

src/net/hosts.rs

@@ -270,20 +270,36 @@ impl Hosts {
     /// Get up to n random hosts from the hosts set.
     pub async fn fetch_n_random(&self, n: u32) -> Vec<Url> {
         let n = n as usize;
+        if n == 0 {
+            return vec![]
+        }
         let addrs = self.addrs.read().await;
         let urls = addrs.iter().choose_multiple(&mut OsRng, n.min(addrs.len()));
         let urls = urls.iter().map(|&url| url.clone()).collect();
         urls
     }
 
-    /// Get all peers that match the given transport schemes from the hosts set.
-    /// TODO: add a limit: usize argument
-    pub async fn fetch_with_schemes(&self, schemes: &[String]) -> Vec<Url> {
+    /// Get up to limit peers that match the given transport schemes from the hosts set.
+    /// If limit was not provided, return all peers.
+    pub async fn fetch_with_schemes(&self, schemes: &[String], limit: Option<usize>) -> Vec<Url> {
+        let addrs = self.addrs.read().await;
+        let mut limit = match limit {
+            Some(l) => l.min(addrs.len()),
+            None => addrs.len(),
+        };
         let mut ret = vec![];
 
-        for addr in self.addrs.read().await.iter() {
+        if limit == 0 {
+            return ret
+        }
+
+        for addr in addrs.iter() {
             if schemes.contains(&addr.scheme().to_string()) {
                 ret.push(addr.clone());
+                limit -= 1;
+                if limit == 0 {
+                    return ret
+                }
             }
         }
 
@@ -292,6 +308,10 @@ impl Hosts {
             for addr in self.quarantine.read().await.keys() {
                 if schemes.contains(&addr.scheme().to_string()) {
                     ret.push(addr.clone());
+                    limit -= 1;
+                    if limit == 0 {
+                        break
+                    }
                 }
             }
         }

src/net/session/outbound_session.rs

@@ -360,7 +360,7 @@ impl Slot {
         macro_rules! mix_transport {
             ($a:expr, $b:expr) => {
                 if transports.contains(&$a.to_string()) && transport_mixing {
-                    let mut a_to_b = p2p.hosts().fetch_with_schemes(&[$b.to_string()]).await;
+                    let mut a_to_b = p2p.hosts().fetch_with_schemes(&[$b.to_string()], None).await;
                     for addr in a_to_b.iter_mut() {
                         addr.set_scheme($a).unwrap();
                         hosts.push(addr.clone());
@@ -374,7 +374,7 @@ impl Slot {
         mix_transport!("nym+tls", "tcp+tls");
 
         // And now the actual requested transports
-        for addr in p2p.hosts().fetch_with_schemes(transports).await {
+        for addr in p2p.hosts().fetch_with_schemes(transports, None).await {
             hosts.push(addr);
         }