addrman: change Select to support multiple networks

Author: brunoerg

src/addrman.cpp

@@ -710,7 +710,7 @@ void AddrManImpl::Attempt_(const CService& addr, bool fCountFailure, NodeSeconds
     }
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, std::optional<Network> network) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, const std::unordered_set<Network>& networks) const
 {
     AssertLockHeld(cs);
 
@@ -719,13 +719,18 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, std::option
     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 (!networks.empty()) {
+        new_count = 0;
+        tried_count = 0;
+        for (auto& network : networks) {
+            auto it = m_network_counts.find(network);
+            if (it == m_network_counts.end()) {
+                continue;
+            }
+            auto counts = it->second;
+            new_count += counts.n_new;
+            tried_count += counts.n_tried;
+        }
     }
 
     if (new_only && new_count == 0) return {};
@@ -758,9 +763,9 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, std::option
             position = (initial_position + i) % ADDRMAN_BUCKET_SIZE;
             node_id = GetEntry(search_tried, bucket, position);
             if (node_id != -1) {
-                if (network.has_value()) {
+                if (!networks.empty()) {
                     const auto it{mapInfo.find(node_id)};
-                    if (Assume(it != mapInfo.end()) && it->second.GetNetwork() == *network) break;
+                    if (Assume(it != mapInfo.end()) && networks.contains(it->second.GetNetwork())) break;
                 } else {
                     break;
                 }
@@ -1208,11 +1213,11 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::SelectTriedCollision()
     return ret;
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only, std::optional<Network> network) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only, const std::unordered_set<Network>& networks) const
 {
     LOCK(cs);
     Check();
-    auto addrRet = Select_(new_only, network);
+    auto addrRet = Select_(new_only, networks);
     Check();
     return addrRet;
 }
@@ -1315,9 +1320,9 @@ std::pair<CAddress, NodeSeconds> AddrMan::SelectTriedCollision()
     return m_impl->SelectTriedCollision();
 }
 
-std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only, std::optional<Network> network) const
+std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only, const std::unordered_set<Network>& networks) const
 {
-    return m_impl->Select(new_only, network);
+    return m_impl->Select(new_only, networks);
 }
 
 std::vector<CAddress> AddrMan::GetAddr(size_t max_addresses, size_t max_pct, std::optional<Network> network, const bool filtered) const

src/addrman.h

@@ -15,6 +15,7 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -154,12 +155,12 @@ class AddrMan
      *                     an address from the new table or an empty pair. Passing `false` will return an
      *                     empty pair or an address from either the new or tried table (it does not
      *                     guarantee a tried entry).
-     * @param[in] network  Select only addresses of this network (nullopt = all). Passing a network may
+     * @param[in] networks Select only addresses of these networks (empty = all). Passing networks may
      *                     slow down the search.
      * @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, std::optional<Network> network = std::nullopt) const;
+    std::pair<CAddress, NodeSeconds> Select(bool new_only = false, const std::unordered_set<Network>& networks = {}) const;
 
     /**
      * Return all or many randomly selected addresses, optionally by network.

src/addrman_impl.h

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

src/bench/addrman.cpp

@@ -133,7 +133,7 @@ static void AddrManSelectByNetwork(benchmark::Bench& bench)
     FillAddrMan(addrman);
 
     bench.run([&] {
-        (void)addrman.Select(/*new_only=*/false, NET_I2P);
+        (void)addrman.Select(/*new_only=*/false, {NET_I2P});
     });
 }
 

src/net.cpp

@@ -2739,7 +2739,11 @@ void CConnman::ThreadOpenConnections(const std::vector<std::string> connect, Spa
                 // If preferred_net has a value set, pick an extra outbound
                 // peer from that network. The eviction logic in net_processing
                 // ensures that a peer from another network will be evicted.
-                std::tie(addr, addr_last_try) = addrman.Select(false, preferred_net);
+                std::unordered_set<Network> preferred_nets;
+                if (preferred_net.has_value()) {
+                    preferred_nets = {*preferred_net};
+                }
+                std::tie(addr, addr_last_try) = addrman.Select(false, preferred_nets);
             }
 
             // Require outbound IPv4/IPv6 connections, other than feelers, to be to distinct network groups

src/test/addrman_tests.cpp

@@ -196,55 +196,59 @@ BOOST_AUTO_TEST_CASE(addrman_select)
 BOOST_AUTO_TEST_CASE(addrman_select_by_network)
 {
     auto addrman = std::make_unique<AddrMan>(EMPTY_NETGROUPMAN, DETERMINISTIC, GetCheckRatio(m_node));
-    BOOST_CHECK(!addrman->Select(/*new_only=*/true, NET_IPV4).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_IPV4).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/true, {NET_IPV4}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_IPV4}).first.IsValid());
 
     // add ipv4 address to the new table
     CNetAddr source = ResolveIP("252.2.2.2");
     CService addr1 = ResolveService("250.1.1.1", 8333);
     BOOST_CHECK(addrman->Add({CAddress(addr1, NODE_NONE)}, source));
 
-    BOOST_CHECK(addrman->Select(/*new_only=*/true, NET_IPV4).first == addr1);
-    BOOST_CHECK(addrman->Select(/*new_only=*/false, NET_IPV4).first == addr1);
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_IPV6).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_ONION).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_I2P).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_CJDNS).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/true, NET_CJDNS).first.IsValid());
+    BOOST_CHECK(addrman->Select(/*new_only=*/true, {NET_IPV4}).first == addr1);
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, {NET_IPV4}).first == addr1);
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_IPV6}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_ONION}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_I2P}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_CJDNS}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/true, {NET_CJDNS}).first.IsValid());
     BOOST_CHECK(addrman->Select(/*new_only=*/false).first == addr1);
 
     // add I2P address to the new table
     CAddress i2p_addr;
     i2p_addr.SetSpecial("udhdrtrcetjm5sxzskjyr5ztpeszydbh4dpl3pl4utgqqw2v4jna.b32.i2p");
     BOOST_CHECK(addrman->Add({i2p_addr}, source));
 
-    BOOST_CHECK(addrman->Select(/*new_only=*/true, NET_I2P).first == i2p_addr);
-    BOOST_CHECK(addrman->Select(/*new_only=*/false, NET_I2P).first == i2p_addr);
-    BOOST_CHECK(addrman->Select(/*new_only=*/false, NET_IPV4).first == addr1);
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_IPV6).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_ONION).first.IsValid());
-    BOOST_CHECK(!addrman->Select(/*new_only=*/false, NET_CJDNS).first.IsValid());
+    BOOST_CHECK(addrman->Select(/*new_only=*/true, {NET_I2P}).first == i2p_addr);
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, {NET_I2P}).first == i2p_addr);
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, {NET_IPV4}).first == addr1);
+    std::unordered_set<Network> nets_with_entries = {NET_IPV4, NET_I2P};
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, nets_with_entries).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_IPV6}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_ONION}).first.IsValid());
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, {NET_CJDNS}).first.IsValid());
+    std::unordered_set<Network> nets_without_entries = {NET_IPV6, NET_ONION, NET_CJDNS};
+    BOOST_CHECK(!addrman->Select(/*new_only=*/false, nets_without_entries).first.IsValid());
 
     // bump I2P address to tried table
     BOOST_CHECK(addrman->Good(i2p_addr));
 
-    BOOST_CHECK(!addrman->Select(/*new_only=*/true, NET_I2P).first.IsValid());
-    BOOST_CHECK(addrman->Select(/*new_only=*/false, NET_I2P).first == i2p_addr);
+    BOOST_CHECK(!addrman->Select(/*new_only=*/true, {NET_I2P}).first.IsValid());
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, {NET_I2P}).first == i2p_addr);
 
     // add another I2P address to the new table
     CAddress i2p_addr2;
     i2p_addr2.SetSpecial("c4gfnttsuwqomiygupdqqqyy5y5emnk5c73hrfvatri67prd7vyq.b32.i2p");
     BOOST_CHECK(addrman->Add({i2p_addr2}, source));
 
-    BOOST_CHECK(addrman->Select(/*new_only=*/true, NET_I2P).first == i2p_addr2);
+    BOOST_CHECK(addrman->Select(/*new_only=*/true, {NET_I2P}).first == i2p_addr2);
 
     // ensure that both new and tried table are selected from
     bool new_selected{false};
     bool tried_selected{false};
     int counter = 256;
 
     while (--counter > 0 && (!new_selected || !tried_selected)) {
-        const CAddress selected{addrman->Select(/*new_only=*/false, NET_I2P).first};
+        const CAddress selected{addrman->Select(/*new_only=*/false, {NET_I2P}).first};
         BOOST_REQUIRE(selected == i2p_addr || selected == i2p_addr2);
         if (selected == i2p_addr) {
             tried_selected = true;
@@ -277,7 +281,7 @@ BOOST_AUTO_TEST_CASE(addrman_select_special)
     // since the only ipv4 address is on the new table, ensure that the new
     // table gets selected even if new_only is false. if the table was being
     // selected at random, this test will sporadically fail
-    BOOST_CHECK(addrman->Select(/*new_only=*/false, NET_IPV4).first == addr1);
+    BOOST_CHECK(addrman->Select(/*new_only=*/false, {NET_IPV4}).first == addr1);
 }
 
 BOOST_AUTO_TEST_CASE(addrman_new_collisions)

src/test/fuzz/addrman.cpp

@@ -285,7 +285,15 @@ FUZZ_TARGET(addrman, .init = initialize_addrman)
     auto max_pct = fuzzed_data_provider.ConsumeIntegralInRange<size_t>(0, 4096);
     auto filtered = fuzzed_data_provider.ConsumeBool();
     (void)const_addr_man.GetAddr(max_addresses, max_pct, network, filtered);
-    (void)const_addr_man.Select(fuzzed_data_provider.ConsumeBool(), network);
+
+    std::unordered_set<Network> nets;
+    for (const auto& net : ALL_NETWORKS) {
+        if (fuzzed_data_provider.ConsumeBool()) {
+            nets.insert(net);
+        }
+    }
+    (void)const_addr_man.Select(fuzzed_data_provider.ConsumeBool(), nets);
+
     std::optional<bool> in_new;
     if (fuzzed_data_provider.ConsumeBool()) {
         in_new = fuzzed_data_provider.ConsumeBool();