diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -151,12 +151,36 @@ bool addNode(NodeId nodeid, const ProofId &proofid); bool removeNode(NodeId nodeid); - bool forNode(NodeId nodeid, std::function func) const; + // Update when a node is to be polled next. bool updateNextRequestTime(NodeId nodeid, TimePoint timeout); // Randomly select a node to poll. NodeId selectNode(); + template + bool forNode(NodeId nodeid, Callable &&func) const { + auto it = nodes.find(nodeid); + return it != nodes.end() && func(*it); + } + + template + void forEachNode(const Peer &peer, Callable &&func) const { + auto &nview = nodes.get(); + auto range = nview.equal_range(peer.peerid); + for (auto it = range.first; it != range.second; ++it) { + func(*it); + } + } + + /** + * Proof and Peer related API. + */ + template void forEachPeer(Callable &&func) const { + for (const auto &p : peers) { + func(p); + } + } + /** * Update the peer set when a new block is connected. */ @@ -203,9 +227,6 @@ uint64_t getSlotCount() const { return slotCount; } uint64_t getFragmentation() const { return fragmentation; } - std::vector getPeers() const; - std::vector getNodeIdsForPeer(PeerId peerId) const; - std::shared_ptr getProof(const ProofId &proofid) const; std::chrono::seconds getProofRegistrationTime(const ProofId &proofid) const; diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -117,12 +117,6 @@ return true; } -bool PeerManager::forNode(NodeId nodeid, - std::function func) const { - auto it = nodes.find(nodeid); - return it != nodes.end() && func(*it); -} - bool PeerManager::updateNextRequestTime(NodeId nodeid, TimePoint timeout) { auto it = nodes.find(nodeid); if (it == nodes.end()) { @@ -472,24 +466,6 @@ return NO_PEER; } -std::vector PeerManager::getPeers() const { - std::vector vpeers; - for (auto &it : peers.get<0>()) { - vpeers.emplace_back(it); - } - return vpeers; -} - -std::vector PeerManager::getNodeIdsForPeer(PeerId peerId) const { - std::vector nodeids; - auto &nview = nodes.get(); - auto nodeRange = nview.equal_range(peerId); - for (auto it = nodeRange.first; it != nodeRange.second; ++it) { - nodeids.emplace_back(it->nodeid); - } - return nodeids; -} - bool PeerManager::isOrphan(const ProofId &id) const { return orphanProofs.getProof(id) != nullptr; } diff --git a/src/avalanche/processor.h b/src/avalanche/processor.h --- a/src/avalanche/processor.h +++ b/src/avalanche/processor.h @@ -6,7 +6,6 @@ #define BITCOIN_AVALANCHE_PROCESSOR_H #include -#include #include #include #include @@ -284,6 +283,11 @@ bool addNode(NodeId nodeid, const ProofId &proofid); bool forNode(NodeId nodeid, std::function func) const; + template auto withPeerManager(Callable &&func) const { + LOCK(cs_peerManager); + return func(*peerManager); + } + CPubKey getSessionPubKey() const; bool sendHello(CNode *pfrom) const; @@ -298,9 +302,6 @@ */ bool isAvalancheServiceAvailable() { return !!peerData; } - std::vector getPeers() const; - std::vector getNodeIdsForPeer(PeerId peerId) const; - bool startEventLoop(CScheduler &scheduler); bool stopEventLoop(); diff --git a/src/avalanche/processor.cpp b/src/avalanche/processor.cpp --- a/src/avalanche/processor.cpp +++ b/src/avalanche/processor.cpp @@ -760,16 +760,6 @@ } while (nodeid != NO_NODE); } -std::vector Processor::getPeers() const { - LOCK(cs_peerManager); - return peerManager->getPeers(); -} - -std::vector Processor::getNodeIdsForPeer(PeerId peerId) const { - LOCK(cs_peerManager); - return peerManager->getNodeIdsForPeer(peerId); -} - void Processor::addUnbroadcastProof(const ProofId &proofid) { LOCK(cs_peerManager); peerManager->addUnbroadcastProof(proofid); 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 @@ -487,12 +487,13 @@ BOOST_CHECK(pm.isOrphan(proof3->getId())); const auto isGoodPeer = [&pm](const std::shared_ptr &p) { - for (const auto &peer : pm.getPeers()) { + bool ret = false; + pm.forEachPeer([&](const Peer &peer) { if (p->getId() == peer.proof->getId()) { - return true; + ret = true; } - } - return false; + }); + return ret; }; BOOST_CHECK(isGoodPeer(proof1)); diff --git a/src/rpc/avalanche.cpp b/src/rpc/avalanche.cpp --- a/src/rpc/avalanche.cpp +++ b/src/rpc/avalanche.cpp @@ -448,24 +448,27 @@ UniValue ret(UniValue::VARR); - for (const auto &peer : g_avalanche->getPeers()) { - UniValue obj(UniValue::VOBJ); + g_avalanche->withPeerManager([&](const avalanche::PeerManager &pm) { + pm.forEachPeer([&](const avalanche::Peer &peer) { + UniValue obj(UniValue::VOBJ); - CDataStream serproof(SER_NETWORK, PROTOCOL_VERSION); - serproof << *peer.proof; + CDataStream serproof(SER_NETWORK, PROTOCOL_VERSION); + serproof << *peer.proof; - obj.pushKV("peerid", uint64_t(peer.peerid)); - obj.pushKV("proof", HexStr(serproof)); + obj.pushKV("peerid", uint64_t(peer.peerid)); + obj.pushKV("proof", HexStr(serproof)); - UniValue nodes(UniValue::VARR); - for (const auto &id : g_avalanche->getNodeIdsForPeer(peer.peerid)) { - nodes.push_back(id); - } - obj.pushKV("nodes", nodes); - obj.pushKV("nodecount", uint64_t(peer.node_count)); + UniValue nodes(UniValue::VARR); + pm.forEachNode(peer, [&](const avalanche::Node &n) { + nodes.push_back(n.nodeid); + }); - ret.push_back(obj); - } + obj.pushKV("nodes", nodes); + obj.pushKV("nodecount", uint64_t(peer.node_count)); + + ret.push_back(obj); + }); + }); return ret; }