diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -495,6 +495,7 @@
addrdb.cpp
addrman.cpp
avalanche/avalanche.cpp
+ avalanche/peermanager.cpp
banman.cpp
blockencodings.cpp
blockfilter.cpp
diff --git a/src/Makefile.am b/src/Makefile.am
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -115,6 +115,7 @@
addrman.h \
attributes.h \
avalanche/avalanche.h \
+ avalanche/peermanager.h \
banman.h \
base58.h \
bloom.h \
@@ -294,6 +295,7 @@
addrdb.cpp \
addrman.cpp \
avalanche/avalanche.cpp \
+ avalanche/peermanager.cpp \
banman.cpp \
blockencodings.cpp \
blockfilter.cpp \
diff --git a/src/Makefile.test.include b/src/Makefile.test.include
--- a/src/Makefile.test.include
+++ b/src/Makefile.test.include
@@ -80,6 +80,7 @@
# test_bitcoin binary #
BITCOIN_TESTS =\
avalanche/test/avalanche_tests.cpp \
+ avalanche/test/peermanager_tests.cpp \
test/scriptnum10.h \
test/activation_tests.cpp \
test/addrman_tests.cpp \
diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h
new file mode 100644
--- /dev/null
+++ b/src/avalanche/peermanager.h
@@ -0,0 +1,37 @@
+// Copyright (c) 2020 The Bitcoin developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#ifndef BITCOIN_AVALANCHE_PEERMANAGER_H
+#define BITCOIN_AVALANCHE_PEERMANAGER_H
+
+#include <atomic>
+#include <cstdint>
+#include <vector>
+
+static constexpr size_t NO_PEER = ~size_t(0);
+
+struct Slot {
+ uint64_t start;
+ uint64_t stop;
+
+ Slot(uint64_t startIn, uint64_t stopIn) : start(startIn), stop(stopIn) {}
+};
+
+class PeerManager {
+ std::vector<Slot> slots;
+ uint64_t slotCount = 0;
+
+public:
+ void addPeer(uint64_t score);
+
+ size_t selectPeer() const;
+};
+
+/**
+ * This is an internal method that is exposed for testing purposes.
+ */
+size_t selectPeerImpl(const std::vector<Slot> &slots, const uint64_t slot,
+ const uint64_t max);
+
+#endif // BITCOIN_AVALANCHE_PEERMANAGER_H
diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp
new file mode 100644
--- /dev/null
+++ b/src/avalanche/peermanager.cpp
@@ -0,0 +1,80 @@
+// Copyright (c) 2020 The Bitcoin developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#include <avalanche/peermanager.h>
+
+#include <random.h>
+
+void PeerManager::addPeer(uint64_t score) {
+ const uint64_t start = slotCount;
+ const uint64_t stop = start + score;
+ slots.emplace_back(start, stop);
+ slotCount = stop;
+}
+
+size_t PeerManager::selectPeer() const {
+ if (slots.empty()) {
+ return NO_PEER;
+ }
+
+ const uint64_t max = slotCount;
+ while (true) {
+ size_t i = selectPeerImpl(slots, GetRand(max), max);
+ if (i != NO_PEER) {
+ return i;
+ }
+ }
+}
+
+size_t selectPeerImpl(const std::vector<Slot> &slots, const uint64_t slot,
+ const uint64_t max) {
+ assert(slot <= max);
+
+ size_t begin = 0, end = slots.size();
+ uint64_t bottom = 0, top = max;
+
+ // Try to find the slot using dichotomic search.
+ while ((end - begin) > 8) {
+ // The slot we picked in not allocated.
+ if (slot < bottom || slot >= top) {
+ return NO_PEER;
+ }
+
+ // Guesstimate the position of the slot.
+ size_t i = begin + ((slot - bottom) * (end - begin) / (top - bottom));
+
+ // We have a match.
+ if (slots[i].start <= slot && slot < slots[i].stop) {
+ return i;
+ }
+
+ // We undershooted.
+ if (slot >= slots[i].stop) {
+ begin = i + 1;
+ bottom = slots[begin].start;
+ continue;
+ }
+
+ // We overshooted.
+ if (slots[i].start > slot) {
+ end = i;
+ top = slots[end].start;
+ continue;
+ }
+
+ // We have an unalocated slot.
+ return NO_PEER;
+ }
+
+ // Enough of that nonsense, let fallback to linear search.
+ for (size_t i = begin; i < end; i++) {
+ // We have a match.
+ if (slots[i].start <= slot && slot < slots[i].stop) {
+ return i;
+ }
+ }
+
+ // We failed to find a slot, retry.
+ return NO_PEER;
+}
diff --git a/src/avalanche/test/CMakeLists.txt b/src/avalanche/test/CMakeLists.txt
--- a/src/avalanche/test/CMakeLists.txt
+++ b/src/avalanche/test/CMakeLists.txt
@@ -8,11 +8,12 @@
create_test_suite(avalanche)
add_dependencies(check check-avalanche)
-add_boost_unit_tests_to_suite(avalanche test_avalanche
+add_boost_unit_tests_to_suite(avalanche test-avalanche
fixture.cpp
TESTS
avalanche_tests.cpp
+ peermanager_tests.cpp
)
-target_link_libraries(test_avalanche server testutil)
+target_link_libraries(test-avalanche server testutil)
diff --git a/src/avalanche/test/peermanager_tests.cpp b/src/avalanche/test/peermanager_tests.cpp
new file mode 100644
--- /dev/null
+++ b/src/avalanche/test/peermanager_tests.cpp
@@ -0,0 +1,180 @@
+// Copyright (c) 2020 The Bitcoin developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#include <avalanche/peermanager.h>
+
+#include <test/util/setup_common.h>
+
+#include <boost/test/unit_test.hpp>
+
+BOOST_FIXTURE_TEST_SUITE(peermanager_tests, BasicTestingSetup)
+
+BOOST_AUTO_TEST_CASE(select_peer_linear) {
+ // No peers.
+ BOOST_CHECK_EQUAL(selectPeerImpl({}, 0, 0), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl({}, 1, 3), NO_PEER);
+
+ // One peer
+ const std::vector<Slot> oneslot = {{100, 200}};
+
+ // Undershoot
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 0, 300), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 42, 300), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 99, 300), NO_PEER);
+
+ // Nailed it
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 100, 300), 0);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 142, 300), 0);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 199, 300), 0);
+
+ // Overshoot
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 200, 300), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 242, 300), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 299, 300), NO_PEER);
+
+ // Two peers
+ const std::vector<Slot> twoslots = {{100, 200}, {300, 400}};
+
+ // Undershoot
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 0, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 42, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 99, 500), NO_PEER);
+
+ // First entry
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 100, 500), 0);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 142, 500), 0);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 199, 500), 0);
+
+ // In betwenn
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 200, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 242, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 299, 500), NO_PEER);
+
+ // Second entry
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 300, 500), 1);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 342, 500), 1);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 399, 500), 1);
+
+ // Overshoot
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 400, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 442, 500), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 499, 500), NO_PEER);
+}
+
+BOOST_AUTO_TEST_CASE(select_peer_dichotomic) {
+ std::vector<Slot> slots;
+
+ // 100 peers of size 1 with 1 empty element apart.
+ uint64_t max = 1;
+ for (int i = 0; i < 100; i++) {
+ slots.emplace_back(max, max + 1);
+ max += 2;
+ }
+
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 4, max), NO_PEER);
+
+ // Check that we get what we expect.
+ for (int i = 0; i < 100; i++) {
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 2 * i, max), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 2 * i + 1, max), i);
+ }
+
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, max, max), NO_PEER);
+
+ // Update the slots to be heavily skewed toward the last element.
+ slots.rbegin()->stop = 300;
+ max = 300;
+
+ for (int i = 0; i < 100; i++) {
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 2 * i, max), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 2 * i + 1, max), i);
+ }
+
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 200, max), 99);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 256, max), 99);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 299, max), 99);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 300, max), NO_PEER);
+
+ // Update the slots to be heavily skewed toward the first element.
+ for (int i = 0; i < 100; i++) {
+ slots[i].start += 100;
+ slots[i].stop += 100;
+ }
+
+ slots.begin()->start = 1;
+ slots.rbegin()->stop = 300;
+
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 0, max), NO_PEER);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 1, max), 0);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 42, max), 0);
+
+ for (int i = 0; i < 100; i++) {
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 100 + 2 * i + 1, max), i);
+ BOOST_CHECK_EQUAL(selectPeerImpl(slots, 100 + 2 * i + 2, max), NO_PEER);
+ }
+}
+
+BOOST_AUTO_TEST_CASE(select_peer_random) {
+ for (int c = 0; c < 1000; c++) {
+ size_t size = InsecureRandBits(10) + 1;
+ std::vector<Slot> slots;
+ slots.reserve(size);
+
+ uint64_t max = InsecureRandBits(3);
+ auto next = [&]() {
+ uint64_t r = max;
+ max += InsecureRandBits(3);
+ return r;
+ };
+
+ for (size_t i = 0; i < size; i++) {
+ uint64_t start = next();
+ uint64_t stop = next();
+ slots.emplace_back(start, stop);
+ }
+
+ for (int k = 0; k < 100; k++) {
+ uint64_t s = InsecureRandRange(max);
+ auto i = selectPeerImpl(slots, s, max);
+ BOOST_CHECK(i == NO_PEER ||
+ (slots[i].start <= s && s < slots[i].stop));
+ }
+ }
+}
+
+BOOST_AUTO_TEST_CASE(add_peer) {
+ // No peers.
+ PeerManager pm;
+ BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER);
+
+ // One peer, we always return it.
+ pm.addPeer(100);
+ BOOST_CHECK_EQUAL(pm.selectPeer(), 0);
+
+ // Two peers, verify ratio.
+ pm.addPeer(200);
+
+ std::array<int, 3> results = {};
+ for (int i = 0; i < 10000; i++) {
+ size_t p = pm.selectPeer();
+ BOOST_CHECK(p <= 1);
+ results[p]++;
+ }
+
+ BOOST_CHECK(abs(2 * results[0] - results[1]) < 500);
+
+ // Three peers, verify ratio.
+ pm.addPeer(100);
+
+ results = {};
+ for (int i = 0; i < 10000; i++) {
+ size_t p = pm.selectPeer();
+ BOOST_CHECK(p <= 2);
+ results[p]++;
+ }
+
+ BOOST_CHECK(abs(results[0] - results[1] + results[2]) < 500);
+}
+
+BOOST_AUTO_TEST_SUITE_END()