diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -5,25 +5,34 @@ #ifndef BITCOIN_AVALANCHE_PEERMANAGER_H #define BITCOIN_AVALANCHE_PEERMANAGER_H -#include #include +#include #include -static constexpr size_t NO_PEER = ~size_t(0); +using PeerId = uint32_t; +static constexpr PeerId NO_PEER = ~uint32_t(0); struct Slot { private: uint64_t start; uint32_t score; + PeerId peerid; public: - Slot(uint64_t startIn, uint32_t scoreIn) : start(startIn), score(scoreIn) {} + Slot(uint64_t startIn, uint32_t scoreIn, PeerId peeridIn) + : start(startIn), score(scoreIn), peerid(peeridIn) {} - Slot withScore(uint64_t scoreIn) { return Slot(start, scoreIn); } + Slot withScore(uint64_t scoreIn) const { + return Slot(start, scoreIn, peerid); + } + Slot withPeerId(PeerId peeridIn) const { + return Slot(start, score, peeridIn); + } uint64_t getStart() const { return start; } uint64_t getStop() const { return start + score; } uint32_t getScore() const { return score; } + uint32_t getPeerId() const { return peerid; } bool contains(uint64_t slot) const { return getStart() <= slot && slot < getStop(); @@ -37,22 +46,28 @@ uint64_t slotCount = 0; uint64_t fragmentation = 0; + PeerId nextPeerId = 0; + std::unordered_map peerIndices; + public: - void addPeer(uint32_t score); - void rescorePeer(size_t i, uint32_t score); - void removePeer(size_t i) { rescorePeer(i, 0); } + PeerId addPeer(uint32_t score) { return addPeer(nextPeerId++, score); } + bool removePeer(PeerId p); + bool rescorePeer(PeerId p, uint32_t score); - size_t selectPeer() const; + PeerId selectPeer() const; // Accssors. uint64_t getSlotCount() const { return slotCount; } uint64_t getFragmentation() const { return fragmentation; } + +private: + PeerId addPeer(PeerId peerid, uint32_t score); }; /** * This is an internal method that is exposed for testing purposes. */ -size_t selectPeerImpl(const std::vector &slots, const uint64_t slot, +PeerId selectPeerImpl(const std::vector &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 --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -6,13 +6,46 @@ #include -void PeerManager::addPeer(uint32_t score) { +#include + +PeerId PeerManager::addPeer(PeerId p, uint32_t score) { + auto inserted = peerIndices.emplace(p, slots.size()); + assert(inserted.second); + const uint64_t start = slotCount; - slots.emplace_back(start, score); + slots.emplace_back(start, score, p); slotCount = start + score; + return p; +} + +bool PeerManager::removePeer(PeerId p) { + auto it = peerIndices.find(p); + if (it == peerIndices.end()) { + return false; + } + + size_t i = it->second; + assert(i < slots.size()); + + if (i + 1 == slots.size()) { + slots.pop_back(); + slotCount = slots.empty() ? 0 : slots.rbegin()->getStop(); + } else { + fragmentation += slots[i].getScore(); + slots[i] = slots[i].withPeerId(NO_PEER); + } + + peerIndices.erase(it); + return true; } -void PeerManager::rescorePeer(size_t i, uint32_t score) { +bool PeerManager::rescorePeer(PeerId p, uint32_t score) { + auto it = peerIndices.find(p); + if (it == peerIndices.end()) { + return false; + } + + size_t i = it->second; assert(i < slots.size()); const uint64_t start = slots[i].getStart(); @@ -21,7 +54,7 @@ if (i + 1 == slots.size()) { slots[i] = slots[i].withScore(score); slotCount = slots[i].getStop(); - return; + return true; } const uint64_t stop = start + score; @@ -31,15 +64,21 @@ if (stop <= nextStart) { fragmentation += (slots[i].getStop() - stop); slots[i] = slots[i].withScore(score); - return; + return true; } - // So we remove and insert a new entry. - addPeer(score); - removePeer(i); + // So we need to insert a new entry. + fragmentation += slots[i].getScore(); + slots[i] = slots[i].withPeerId(NO_PEER); + it->second = slots.size(); + const uint64_t newStart = slotCount; + slots.emplace_back(newStart, score, p); + slotCount = newStart + score; + + return true; } -size_t PeerManager::selectPeer() const { +PeerId PeerManager::selectPeer() const { if (slots.empty()) { return NO_PEER; } @@ -53,7 +92,7 @@ } } -size_t selectPeerImpl(const std::vector &slots, const uint64_t slot, +PeerId selectPeerImpl(const std::vector &slots, const uint64_t slot, const uint64_t max) { assert(slot <= max); @@ -73,7 +112,7 @@ // We have a match. if (slots[i].contains(slot)) { - return i; + return slots[i].getPeerId(); } // We undershooted. @@ -102,7 +141,7 @@ for (size_t i = begin; i < end; i++) { // We have a match. if (slots[i].contains(slot)) { - return i; + return slots[i].getPeerId(); } } 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 @@ -16,7 +16,7 @@ BOOST_CHECK_EQUAL(selectPeerImpl({}, 1, 3), NO_PEER); // One peer - const std::vector oneslot = {{100, 100}}; + const std::vector oneslot = {{100, 100, 23}}; // Undershoot BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 0, 300), NO_PEER); @@ -24,9 +24,9 @@ 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); + BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 100, 300), 23); + BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 142, 300), 23); + BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 199, 300), 23); // Overshoot BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 200, 300), NO_PEER); @@ -34,7 +34,7 @@ BOOST_CHECK_EQUAL(selectPeerImpl(oneslot, 299, 300), NO_PEER); // Two peers - const std::vector twoslots = {{100, 100}, {300, 100}}; + const std::vector twoslots = {{100, 100, 69}, {300, 100, 42}}; // Undershoot BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 0, 500), NO_PEER); @@ -42,9 +42,9 @@ 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); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 100, 500), 69); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 142, 500), 69); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 199, 500), 69); // In betwenn BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 200, 500), NO_PEER); @@ -52,9 +52,9 @@ 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); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 300, 500), 42); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 342, 500), 42); + BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 399, 500), 42); // Overshoot BOOST_CHECK_EQUAL(selectPeerImpl(twoslots, 400, 500), NO_PEER); @@ -68,7 +68,7 @@ // 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, 1); + slots.emplace_back(max, 1, i); max += 2; } @@ -99,10 +99,11 @@ // Update the slots to be heavily skewed toward the first element. for (int i = 0; i < 100; i++) { - slots[i] = Slot(slots[i].getStart() + 100, slots[i].getScore()); + slots[i] = Slot(slots[i].getStart() + 100, slots[i].getScore(), + slots[i].getPeerId()); } - slots[0] = Slot(1, slots[0].getStop() - 1); + slots[0] = Slot(1, slots[0].getStop() - 1, slots[0].getPeerId()); slots[99] = slots[99].withScore(1); max = slots[99].getStop(); BOOST_CHECK_EQUAL(max, 300); @@ -134,12 +135,14 @@ const uint64_t start = next(); const uint32_t score = InsecureRandBits(3); max += score; - slots.emplace_back(start, score); + slots.emplace_back(start, score, i); } for (int k = 0; k < 100; k++) { uint64_t s = InsecureRandRange(max); auto i = selectPeerImpl(slots, s, max); + // /!\ Because of the way we construct the vector, the peer id is + // always the index. This might not be the case in practice. BOOST_CHECK(i == NO_PEER || slots[i].contains(s)); } } @@ -151,28 +154,28 @@ BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); // One peer, we always return it. - pm.addPeer(100); - BOOST_CHECK_EQUAL(pm.selectPeer(), 0); + PeerId peer0 = pm.addPeer(100); + BOOST_CHECK_EQUAL(pm.selectPeer(), peer0); // Two peers, verify ratio. - pm.addPeer(200); + PeerId peer1 = pm.addPeer(200); - std::array results = {}; + std::unordered_map results = {}; for (int i = 0; i < 10000; i++) { size_t p = pm.selectPeer(); - BOOST_CHECK(p <= 1); + BOOST_CHECK(p == peer0 || p == peer1); results[p]++; } BOOST_CHECK(abs(2 * results[0] - results[1]) < 500); // Three peers, verify ratio. - pm.addPeer(100); + PeerId peer2 = pm.addPeer(100); - results = {}; + results.clear(); for (int i = 0; i < 10000; i++) { size_t p = pm.selectPeer(); - BOOST_CHECK(p <= 2); + BOOST_CHECK(p == peer0 || p == peer1 || p == peer2); results[p]++; } @@ -185,37 +188,58 @@ BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); // Add 4 peers. + std::array peerids; for (int i = 0; i < 4; i++) { - pm.addPeer(100); + peerids[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); + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[1] || p == peerids[2] || + p == peerids[3]); } // Remove one peer, it nevers show up now. - pm.removePeer(3); + BOOST_CHECK(pm.removePeer(peerids[2])); + 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 != 3); + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[1] || p == peerids[3]); } // Add 4 more peers. for (int i = 0; i < 4; i++) { - pm.addPeer(100); + peerids[i + 4] = pm.addPeer(100); } - pm.removePeer(0); - pm.removePeer(7); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 800); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 100); + + BOOST_CHECK(pm.removePeer(peerids[0])); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 800); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 200); + + // Removing the last entry do not increase fragmentation. + BOOST_CHECK(pm.removePeer(peerids[7])); + BOOST_CHECK_EQUAL(pm.getSlotCount(), 700); + BOOST_CHECK_EQUAL(pm.getFragmentation(), 200); + 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); + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[1] || p == peerids[3] || p == peerids[4] || + p == peerids[5] || p == peerids[6]); } + + // Removing non existent peers fails. + BOOST_CHECK(!pm.removePeer(peerids[0])); + BOOST_CHECK(!pm.removePeer(peerids[2])); + BOOST_CHECK(!pm.removePeer(peerids[7])); + BOOST_CHECK(!pm.removePeer(NO_PEER)); } BOOST_AUTO_TEST_CASE(rescore_peer, *boost::unit_test::timeout(5)) { @@ -224,52 +248,57 @@ BOOST_CHECK_EQUAL(pm.selectPeer(), NO_PEER); // Add 4 peers. + std::array peerids; for (int i = 0; i < 4; i++) { - pm.addPeer(100); + peerids[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); + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[1] || p == peerids[2] || + p == peerids[3]); } // Set one peer's score to 0, it nevers show up now. - pm.rescorePeer(1, 0); + BOOST_CHECK(pm.rescorePeer(peerids[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); + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[2] || p == peerids[3]); } // "resurrect" the peer. - pm.rescorePeer(1, 100); + BOOST_CHECK(pm.rescorePeer(peerids[1], 100)); BOOST_CHECK_EQUAL(pm.getSlotCount(), 400); BOOST_CHECK_EQUAL(pm.getFragmentation(), 0); while (true) { - size_t p = pm.selectPeer(); - if (p == 1) { + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[1] || p == peerids[2] || + p == peerids[3]); + // Make sure peer 1 reappeared. + if (p == peerids[1]) { break; } } // Grow the peer to a point where it needs to be reallocated. - pm.rescorePeer(1, 200); + BOOST_CHECK(pm.rescorePeer(peerids[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) { + PeerId p = pm.selectPeer(); + BOOST_CHECK(p == peerids[0] || p == peerids[1] || p == peerids[2] || + p == peerids[3]); + // Make sure peer 1 reappeared. + if (p == peerids[1]) { break; } }