From f904146583020bab71d0ef3f06285aa703a1c23f Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Mon, 17 May 2021 15:27:21 +1000
Subject: [PATCH] Optimise getting the next connection attempt address

---
 zebra-network/src/address_book.rs           | 22 ++++++---------------
 zebra-network/src/peer_set/candidate_set.rs |  2 +-
 2 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/zebra-network/src/address_book.rs b/zebra-network/src/address_book.rs
index 6a22efa228b..6db5966e75e 100644
--- a/zebra-network/src/address_book.rs
+++ b/zebra-network/src/address_book.rs
@@ -1,12 +1,7 @@
 //! The `AddressBook` manages information about what peers exist, when they were
 //! seen, and what services they provide.
 
-use std::{
-    collections::{BTreeSet, HashMap},
-    iter::Extend,
-    net::SocketAddr,
-    time::Instant,
-};
+use std::{collections::HashMap, iter::Extend, net::SocketAddr, time::Instant};
 
 use chrono::{DateTime, Utc};
 use tracing::Span;
@@ -323,19 +318,14 @@ impl AddressBook {
             .cloned()
     }
 
-    /// Return an iterator over peers that are due for a connection attempt,
-    /// in connection attempt order.
-    pub fn connection_peers_ordered(&'_ self) -> impl Iterator<Item = MetaAddr> + '_ {
+    /// Return the next peer that is due for a connection attempt.
+    pub fn next_attempt_peer(&self) -> Option<MetaAddr> {
         let _guard = self.span.enter();
 
-        // TODO: optimise, if needed, or get rid of older peers
-
-        // Skip recently used peers (including live peers), then sort using BTreeSet
-        self.by_addr
-            .values()
+        // Skip recently used peers (including live peers)
+        self.peers_unordered()
             .filter(move |peer| !self.recently_used_addr(&peer.addr))
-            .collect::<BTreeSet<_>>()
-            .into_iter()
+            .min()
             .cloned()
     }
 
diff --git a/zebra-network/src/peer_set/candidate_set.rs b/zebra-network/src/peer_set/candidate_set.rs
index fc6b764c8d3..1884ae92ccd 100644
--- a/zebra-network/src/peer_set/candidate_set.rs
+++ b/zebra-network/src/peer_set/candidate_set.rs
@@ -295,7 +295,7 @@ where
             let mut guard = self.address_book.lock().unwrap();
             // It's okay to return without sleeping here, because we're returning
             // `None`. We only need to sleep before yielding an address.
-            let connect = guard.connection_peers_ordered().next()?;
+            let connect = guard.next_attempt_peer()?;
 
             let connect = MetaAddr::update_attempt(&connect.addr);
             guard.update(connect)?