diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -221,6 +221,8 @@ func(p); } } + /** For a given peer, return the nodes associated with it */ + std::vector nodesForPeer(const PeerId &) const; /** * Update the peer set when a new block is connected. diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -527,4 +527,13 @@ m_unbroadcast_proofids.erase(proofid); } +std::vector PeerManager::nodesForPeer(const PeerId &peerId) const { + std::vector nodeIds; + for (auto [it, last] = nodes.get().equal_range(peerId); + it != last; ++it) { + nodeIds.emplace_back(it->nodeid); + } + return nodeIds; +} + } // namespace avalanche 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 @@ -788,4 +788,40 @@ !pm.registerProof(std::make_shared(std::move(badProof)))); } +BOOST_AUTO_TEST_CASE(nodesForPeer) { + avalanche::PeerManager pm; + + const auto firstProof = getRandomProofPtr(MIN_VALID_PROOF_SCORE); + const auto firstPeerId = pm.getPeerId(firstProof); + + // firstPeerId has no nodes + BOOST_CHECK(pm.nodesForPeer(firstPeerId).empty()); + + const auto theFuture(std::chrono::steady_clock::now() + + std::chrono::hours(24)); + + // Add nodes to this peer and update their request time far in the future + int i = 0; + for (; i < 10; i++) { + BOOST_CHECK(pm.addNode(i, firstProof->getId())); + BOOST_CHECK(pm.updateNextRequestTime(i, theFuture)); + } + + // firstPeerId now should have 10 nodes associated with it + BOOST_CHECK(pm.nodesForPeer(firstPeerId).size() == 10); + + // add another peer + const auto secondProof = getRandomProofPtr(MIN_VALID_PROOF_SCORE); + const auto secondPeerId = pm.getPeerId(secondProof); + + // now add a new set of nodes to secondPeerId as well + for (; i < 30; i++) { + BOOST_CHECK(pm.addNode(i, secondProof->getId())); + BOOST_CHECK(pm.updateNextRequestTime(i, theFuture)); + } + + // secondPeerId should have 20 nodes + BOOST_CHECK(pm.nodesForPeer(secondPeerId).size() == 20); +} + BOOST_AUTO_TEST_SUITE_END()