diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -21,11 +21,18 @@ class PeerManager { std::vector slots; uint64_t slotCount = 0; + uint64_t fragmentation = 0; public: void addPeer(uint64_t score); + void rescorePeer(size_t i, uint64_t score); + void removePeer(size_t i) { rescorePeer(i, 0); } size_t selectPeer() const; + + // Accssors. + uint64_t getSlotCount() const { return slotCount; } + uint64_t getFragmentation() const { return fragmentation; } }; /** diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -13,6 +13,33 @@ slotCount = stop; } +void PeerManager::rescorePeer(size_t i, uint64_t score) { + assert(i < slots.size()); + + const uint64_t start = slots[i].start; + const uint64_t stop = start + score; + + // If this is the last element, we can extend/shrink easily. + if (i + 1 == slots.size()) { + slots[i].stop = stop; + slotCount = stop; + return; + } + + const uint64_t nextStart = slots[i + 1].start; + + // We can extend in place. + if (stop <= nextStart) { + fragmentation += (slots[i].stop - stop); + slots[i].stop = stop; + return; + } + + // So we remove and insert a new entry. + addPeer(score); + removePeer(i); +} + size_t PeerManager::selectPeer() const { if (slots.empty()) { return NO_PEER; diff --git a/src/avalanche/test/peermanager_tests.cpp b/src/avalanche/test/peermanager_tests.cpp --- a/src/avalanche/test/peermanager_tests.cpp +++ b/src/avalanche/test/peermanager_tests.cpp @@ -177,4 +177,101 @@ BOOST_CHECK(abs(results[0] - results[1] + results[2]) < 500); } +BOOST_AUTO_TEST_CASE(remove_peer) { + // No peers. + PeerManager pm; + BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); + + // Add 4 peers. + for (int i = 0; i < 4; i++) { + pm.addPeer(100); + } + + for (int i = 0; i < 100; i++) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p <= 4); + } + + // Remove one peer, it nevers show up now. + pm.removePeer(3); + for (int i = 0; i < 100; i++) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p < 4); + BOOST_CHECK(p != 3); + } + + // Add 4 more peers. + for (int i = 0; i < 4; i++) { + pm.addPeer(100); + } + + pm.removePeer(0); + pm.removePeer(7); + for (int i = 0; i < 100; i++) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p < 8); + BOOST_CHECK(p != 0); + BOOST_CHECK(p != 3); + BOOST_CHECK(p != 7); + } +} + +BOOST_AUTO_TEST_CASE(rescore_peer, *boost::unit_test::timeout(5)) { + // No peers. + PeerManager pm; + BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); + + // Add 4 peers. + for (int i = 0; i < 4; i++) { + pm.addPeer(100); + } + + BOOST_CHECK_EQUAL(pm.getSlotCount(), 400); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 0); + + for (int i = 0; i < 100; i++) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p <= 4); + } + + // Set one peer's score to 0, it nevers show up now. + pm.rescorePeer(1, 0); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 400); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 100); + + for (int i = 0; i < 100; i++) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p < 4); + BOOST_CHECK(p != 1); + } + + // "resurrect" the peer. + pm.rescorePeer(1, 100); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 400); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 0); + + while (true) { + size_t p = pm.selectPeer(); + if (p == 1) { + break; + } + } + + // Grow the peer to a point where it needs to be reallocated. + pm.rescorePeer(1, 200); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 600); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 100); + + for (int i = 0; i < 25; i++) { + while (true) { + size_t p = pm.selectPeer(); + BOOST_CHECK(p < 5); + BOOST_CHECK(p != 1); + if (p == 4) { + break; + } + } + } +} + BOOST_AUTO_TEST_SUITE_END()