addrman: add functionality to select by network

Add an optional parameter to the addrman Select function that allows callers to
specify which network the returned address should be on. Ensure that the proper
table is selected with different cases of whether the new or tried table has
network addresses that match.

Co-authored-by: Martin Zumsande <[email protected]>

Author: amitiuttarwar

src/addrman.cpp

@@ -714,28 +714,41 @@ void AddrManImpl::Attempt_(const CService& addr, bool fCountFailure, NodeSeconds
     }
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, std::optional<Network> network) const
 {
     AssertLockHeld(cs);
 
     if (vRandom.empty()) return {};
-    if (new_only && nNew == 0) return {};
+
+    size_t new_count = nNew;
+    size_t tried_count = nTried;
+
+    if (network.has_value()) {
+        auto it = m_network_counts.find(*network);
+        if (it == m_network_counts.end()) return {};
+
+        auto counts = it->second;
+        new_count = counts.n_new;
+        tried_count = counts.n_tried;
+    }
+
+    if (new_only && new_count == 0) return {};
+    if (new_count + tried_count == 0) return {};
 
     // Decide if we are going to search the new or tried table
+    // If either option is viable, use a 50% chance to choose
     bool search_tried;
-    int bucket_count;
-
-    // Use a 50% chance for choosing between tried and new table entries.
-    if (!new_only &&
-       (nTried > 0 &&
-        (nNew == 0 || insecure_rand.randbool() == 0))) {
+    if (new_only || tried_count == 0) {
+        search_tried = false;
+    } else if (new_count == 0) {
         search_tried = true;
-        bucket_count = ADDRMAN_TRIED_BUCKET_COUNT;
     } else {
-        search_tried = false;
-        bucket_count = ADDRMAN_NEW_BUCKET_COUNT;
+        search_tried = insecure_rand.randbool();
     }
 
+    const int bucket_count{search_tried ? ADDRMAN_TRIED_BUCKET_COUNT : ADDRMAN_NEW_BUCKET_COUNT};
+
+    // Loop through the addrman table until we find an appropriate entry
     double chance_factor = 1.0;
     while (1) {
         // Pick a bucket, and an initial position in that bucket.
@@ -748,7 +761,16 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only) const
         for (i = 0; i < ADDRMAN_BUCKET_SIZE; ++i) {
             int position = (initial_position + i) % ADDRMAN_BUCKET_SIZE;
             int node_id = GetEntry(search_tried, bucket, position);
-            if (node_id != -1) break;
+            if (node_id != -1) {
+                if (network.has_value()) {
+                    const auto it{mapInfo.find(node_id)};
+                    assert(it != mapInfo.end());
+                    const auto info{it->second};
+                    if (info.GetNetwork() == *network) break;
+                } else {
+                    break;
+                }
+            }
         }
 
         // If the bucket is entirely empty, start over with a (likely) different one.
@@ -1168,11 +1190,11 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::SelectTriedCollision()
     return ret;
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only, std::optional<Network> network) const
 {
     LOCK(cs);
     Check();
-    auto addrRet = Select_(new_only);
+    auto addrRet = Select_(new_only, network);
     Check();
     return addrRet;
 }
@@ -1266,9 +1288,9 @@ std::pair<CAddress, NodeSeconds> AddrMan::SelectTriedCollision()
     return m_impl->SelectTriedCollision();
 }
 
-std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only, std::optional<Network> network) const
 {
-    return m_impl->Select(new_only);
+    return m_impl->Select(new_only, network);
 }
 
 std::vector<CAddress> AddrMan::GetAddr(size_t max_addresses, size_t max_pct, std::optional<Network> network) const

src/addrman.h

@@ -146,11 +146,12 @@ class AddrMan
     /**
      * Choose an address to connect to.
      *
-     * @param[in] new_only  Whether to only select addresses from the new table.
+     * @param[in] new_only Whether to only select addresses from the new table.
+     * @param[in] network  Select only addresses of this network (nullopt = all)
      * @return    CAddress The record for the selected peer.
      *            seconds  The last time we attempted to connect to that peer.
      */
-    std::pair<CAddress, NodeSeconds> Select(bool new_only = false) const;
+    std::pair<CAddress, NodeSeconds> Select(bool new_only = false, std::optional<Network> network = std::nullopt) const;
 
     /**
      * Return all or many randomly selected addresses, optionally by network.

src/addrman_impl.h

@@ -127,7 +127,7 @@ class AddrManImpl
 
     std::pair<CAddress, NodeSeconds> SelectTriedCollision() EXCLUSIVE_LOCKS_REQUIRED(!cs);
 
-    std::pair<CAddress, NodeSeconds> Select(bool new_only) const
+    std::pair<CAddress, NodeSeconds> Select(bool new_only, std::optional<Network> network) const
         EXCLUSIVE_LOCKS_REQUIRED(!cs);
 
     std::vector<CAddress> GetAddr(size_t max_addresses, size_t max_pct, std::optional<Network> network) const
@@ -251,7 +251,7 @@ class AddrManImpl
 
     void Attempt_(const CService& addr, bool fCountFailure, NodeSeconds time) EXCLUSIVE_LOCKS_REQUIRED(cs);
 
-    std::pair<CAddress, NodeSeconds> Select_(bool new_only) const EXCLUSIVE_LOCKS_REQUIRED(cs);
+    std::pair<CAddress, NodeSeconds> Select_(bool new_only, std::optional<Network> network) const EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     /** Helper to generalize looking up an addrman entry from either table.
      *