diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -19,7 +19,8 @@ #include using PeerId = uint32_t; -static constexpr PeerId NO_PEER = ~uint32_t(0); +static constexpr PeerId NO_PEER = -1; +static constexpr NodeId NO_NODE = -1; struct Slot { private: @@ -90,18 +91,34 @@ uint32_t index; Peer(uint32_t score_, uint32_t index_) : score(score_), index(index_) {} + + bool addNode(NodeId nodeid, CPubKey pubkey); + NodeId getSuitableNodeToQuery() const; }; PeerId nextPeerId = 0; std::unordered_map peers; static constexpr int SELECT_PEER_MAX_RETRY = 3; + static constexpr int SELECT_NODE_MAX_RETRY = 3; public: + /** + * Peer API. + */ PeerId addPeer(uint32_t score) { return addPeer(nextPeerId++, score); } bool removePeer(PeerId p); bool rescorePeer(PeerId p, uint32_t score); + /** + * Node API. + */ + bool addNodeToPeer(PeerId peerid, NodeId nodeid, CPubKey pubkey); + NodeId getSuitableNodeToQuery(); + + /** + * Exposed for tests. + */ PeerId selectPeer() const; /** diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -82,6 +82,57 @@ return true; } +bool PeerManager::Peer::addNode(NodeId nodeid, CPubKey pubkey) { + return nodes + .insert({nodeid, std::chrono::steady_clock::now(), std::move(pubkey)}) + .second; +} + +bool PeerManager::addNodeToPeer(PeerId peerid, NodeId nodeid, CPubKey pubkey) { + auto it = peers.find(peerid); + if (it == peers.end()) { + return false; + } + + return it->second.addNode(nodeid, std::move(pubkey)); +} + +NodeId PeerManager::Peer::getSuitableNodeToQuery() const { + auto it = nodes.get().begin(); + if (it == nodes.get().end()) { + return NO_NODE; + } + + if (it->nextRequestTime <= std::chrono::steady_clock::now()) { + return it->nodeid; + } + + return NO_NODE; +} + +NodeId PeerManager::getSuitableNodeToQuery() { + for (int retry = 0; retry < SELECT_NODE_MAX_RETRY; retry++) { + PeerId p = selectPeer(); + + // If we cannot find a peer, it may be due to the fact that it is + // unlikely due to high fragmentation, so compact and retry. + if (p == NO_PEER) { + compact(); + continue; + } + + // See if that peer has an available node. + auto it = peers.find(p); + assert(it != peers.end()); + NodeId n = it->second.getSuitableNodeToQuery(); + if (n != NO_NODE) { + return n; + } + } + + return NO_NODE; +} + PeerId PeerManager::selectPeer() const { if (slots.empty() || slotCount == 0) { return NO_PEER; @@ -99,6 +150,11 @@ } uint64_t PeerManager::compact() { + // There is nothing to compact. + if (fragmentation == 0) { + return 0; + } + // Shrink the vector to the expected size. while (slots.size() > peers.size()) { slots.pop_back(); 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 @@ -350,4 +350,22 @@ BOOST_CHECK_EQUAL(pm.getFragmentation(), 0); } +BOOST_AUTO_TEST_CASE(add_and_get_node) { + PeerManager pm; + + // Create one peer. + PeerId peerid = pm.addPeer(100); + BOOST_CHECK_EQUAL(pm.getSuitableNodeToQuery(), NO_NODE); + + // Add 4 nodes. + for (int i = 0; i < 4; i++) { + BOOST_CHECK(pm.addNodeToPeer(peerid, i, CPubKey())); + } + + for (int i = 0; i < 100; i++) { + NodeId n = pm.getSuitableNodeToQuery(); + BOOST_CHECK((n >= 0 && n < 4) || n == NO_PEER); + } +} + BOOST_AUTO_TEST_SUITE_END()