From 6b229284fd2209938ee8fdffed4d080395b3aa05 Mon Sep 17 00:00:00 2001 From: Amiti Uttarwar Date: Sat, 18 Feb 2023 18:29:45 -0700 Subject: [PATCH] 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 --- src/addrman.cpp | 60 +++++++++++++++++++++++++++++++--------------- src/addrman.h | 5 ++-- src/addrman_impl.h | 4 ++-- 3 files changed, 46 insertions(+), 23 deletions(-) diff --git a/src/addrman.cpp b/src/addrman.cpp index fe4a79b2e66..cdfd079fcdf 100644 --- a/src/addrman.cpp +++ b/src/addrman.cpp @@ -714,28 +714,41 @@ void AddrManImpl::Attempt_(const CService& addr, bool fCountFailure, NodeSeconds } } -std::pair AddrManImpl::Select_(bool new_only) const +std::pair AddrManImpl::Select_(bool new_only, std::optional network) const { AssertLockHeld(cs); if (vRandom.empty()) return {}; - if (new_only && nNew == 0) return {}; - // Decide if we are going to search the new or tried table - bool search_tried; - int bucket_count; + size_t new_count = nNew; + size_t tried_count = nTried; - // Use a 50% chance for choosing between tried and new table entries. - if (!new_only && - (nTried > 0 && - (nNew == 0 || insecure_rand.randbool() == 0))) { - search_tried = true; - bucket_count = ADDRMAN_TRIED_BUCKET_COUNT; - } else { - search_tried = false; - bucket_count = ADDRMAN_NEW_BUCKET_COUNT; + 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; + if (new_only || tried_count == 0) { + search_tried = false; + } else if (new_count == 0) { + search_tried = true; + } else { + 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 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 AddrManImpl::SelectTriedCollision() return ret; } -std::pair AddrManImpl::Select(bool new_only) const +std::pair AddrManImpl::Select(bool new_only, std::optional 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 AddrMan::SelectTriedCollision() return m_impl->SelectTriedCollision(); } -std::pair AddrMan::Select(bool new_only) const +std::pair AddrMan::Select(bool new_only, std::optional network) const { - return m_impl->Select(new_only); + return m_impl->Select(new_only, network); } std::vector AddrMan::GetAddr(size_t max_addresses, size_t max_pct, std::optional network) const diff --git a/src/addrman.h b/src/addrman.h index 374996e501f..c8e31fe2833 100644 --- a/src/addrman.h +++ b/src/addrman.h @@ -146,11 +146,12 @@ public: /** * 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 Select(bool new_only = false) const; + std::pair Select(bool new_only = false, std::optional network = std::nullopt) const; /** * Return all or many randomly selected addresses, optionally by network. diff --git a/src/addrman_impl.h b/src/addrman_impl.h index 3cf3b838d68..7aead2812ba 100644 --- a/src/addrman_impl.h +++ b/src/addrman_impl.h @@ -127,7 +127,7 @@ public: std::pair SelectTriedCollision() EXCLUSIVE_LOCKS_REQUIRED(!cs); - std::pair Select(bool new_only) const + std::pair Select(bool new_only, std::optional network) const EXCLUSIVE_LOCKS_REQUIRED(!cs); std::vector GetAddr(size_t max_addresses, size_t max_pct, std::optional network) const @@ -251,7 +251,7 @@ private: void Attempt_(const CService& addr, bool fCountFailure, NodeSeconds time) EXCLUSIVE_LOCKS_REQUIRED(cs); - std::pair Select_(bool new_only) const EXCLUSIVE_LOCKS_REQUIRED(cs); + std::pair Select_(bool new_only, std::optional network) const EXCLUSIVE_LOCKS_REQUIRED(cs); /** Helper to generalize looking up an addrman entry from either table. *