diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -175,10 +175,18 @@ /** * Proof and Peer related API. */ + bool registerProof(const std::shared_ptr &proof); bool exists(const ProofId &proofid) const { return getProof(proofid) != nullptr; } + template + bool forPeer(const ProofId &proofid, Callable &&func) const { + auto &pview = peers.get(); + auto it = pview.find(proofid); + return it != pview.end() && func(*it); + } + template void forEachPeer(Callable &&func) const { for (const auto &p : peers) { func(p); @@ -232,7 +240,6 @@ uint64_t getFragmentation() const { return fragmentation; } std::shared_ptr getProof(const ProofId &proofid) const; - std::chrono::seconds getProofRegistrationTime(const ProofId &proofid) const; bool isOrphan(const ProofId &id) const; std::shared_ptr getOrphan(const ProofId &id) const; diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -125,6 +125,10 @@ return nodes.modify(it, [&](Node &n) { n.nextRequestTime = timeout; }); } +bool PeerManager::registerProof(const std::shared_ptr &proof) { + return !getProof(proof->getId()) && getPeerId(proof) != NO_PEER; +} + NodeId PeerManager::selectNode() { for (int retry = 0; retry < SELECT_NODE_MAX_RETRY; retry++) { const PeerId p = selectPeer(); @@ -189,17 +193,14 @@ } std::shared_ptr PeerManager::getProof(const ProofId &proofid) const { - auto &pview = peers.get(); - auto it = pview.find(proofid); - return it == pview.end() ? nullptr : it->proof; -} + std::shared_ptr proof = nullptr; -std::chrono::seconds -PeerManager::getProofRegistrationTime(const ProofId &proofid) const { - auto &pview = peers.get(); - auto it = pview.find(proofid); - return it == pview.end() ? std::chrono::seconds::max() - : it->registration_time; + forPeer(proofid, [&](const Peer &p) { + proof = p.proof; + return true; + }); + + return proof; } PeerManager::PeerSet::iterator diff --git a/src/avalanche/processor.h b/src/avalanche/processor.h --- a/src/avalanche/processor.h +++ b/src/avalanche/processor.h @@ -294,7 +294,6 @@ bool addProof(const std::shared_ptr &proof); std::shared_ptr getProof(const ProofId &proofid) const; std::shared_ptr getLocalProof() const; - std::chrono::seconds getProofRegistrationTime(const ProofId &proofid) const; std::shared_ptr getOrphan(const ProofId &proofid) const; /* diff --git a/src/avalanche/processor.cpp b/src/avalanche/processor.cpp --- a/src/avalanche/processor.cpp +++ b/src/avalanche/processor.cpp @@ -571,8 +571,7 @@ bool Processor::addProof(const std::shared_ptr &proof) { LOCK(cs_peerManager); - return !peerManager->getProof(proof->getId()) && - peerManager->getPeerId(proof) != NO_PEER; + return peerManager->registerProof(proof); } std::shared_ptr Processor::getProof(const ProofId &proofid) const { @@ -584,12 +583,6 @@ return peerData ? peerData->proof : nullptr; } -std::chrono::seconds -Processor::getProofRegistrationTime(const ProofId &proofid) const { - LOCK(cs_peerManager); - return peerManager->getProofRegistrationTime(proofid); -} - std::shared_ptr Processor::getOrphan(const ProofId &proofid) const { LOCK(cs_peerManager); return peerManager->getOrphan(proofid); diff --git a/src/avalanche/test/processor_tests.cpp b/src/avalanche/test/processor_tests.cpp --- a/src/avalanche/test/processor_tests.cpp +++ b/src/avalanche/test/processor_tests.cpp @@ -973,33 +973,21 @@ proofs.push_back(GetProof()); } - const auto checkpoint = GetTime(); - - for (int i = 0; i < numProofs; i++) { - BOOST_CHECK(m_processor->addProof(proofs[i])); - // Fail to add an existing proof - BOOST_CHECK(!m_processor->addProof(proofs[i])); - - for (int added = 0; added <= i; added++) { - auto proof = m_processor->getProof(proofs[added]->getId()); - BOOST_CHECK(proof != nullptr); - - const ProofId &proofid = proof->getId(); - BOOST_CHECK_EQUAL(proofid, proofs[added]->getId()); - - const auto proofTime = - m_processor->getProofRegistrationTime(proofid); - BOOST_CHECK(proofTime != std::chrono::seconds::max()); - BOOST_CHECK(proofTime >= checkpoint); - } - - for (int missing = i + 1; missing < numProofs; missing++) { - const ProofId &proofid = proofs[missing]->getId(); - BOOST_CHECK(!m_processor->getProof(proofid)); - BOOST_CHECK(m_processor->getProofRegistrationTime(proofid) == - std::chrono::seconds::max()); + m_processor->withPeerManager([&](avalanche::PeerManager &pm) { + for (int i = 0; i < numProofs; i++) { + BOOST_CHECK(pm.registerProof(proofs[i])); + // Fail to add an existing proof + BOOST_CHECK(!pm.registerProof(proofs[i])); + + for (int added = 0; added <= i; added++) { + auto proof = pm.getProof(proofs[added]->getId()); + BOOST_CHECK(proof != nullptr); + + const ProofId &proofid = proof->getId(); + BOOST_CHECK_EQUAL(proofid, proofs[added]->getId()); + } } - } + }); // No stake, copied from proof_tests.cpp const std::string badProofHex( diff --git a/src/net_processing.cpp b/src/net_processing.cpp --- a/src/net_processing.cpp +++ b/src/net_processing.cpp @@ -2074,27 +2074,33 @@ static std::shared_ptr FindProofForGetData(const CNode &peer, const avalanche::ProofId &proofid, const std::chrono::seconds now) { - auto proof = g_avalanche->getProof(proofid); + std::shared_ptr proof = nullptr; + + bool send_unconditionally = + g_avalanche->withPeerManager([&](const avalanche::PeerManager &pm) { + return pm.forPeer(proofid, [&](const avalanche::Peer &peer) { + proof = peer.proof; + + // If we know that proof for long enough, allow for requesting + // it. + return peer.registration_time <= + now - UNCONDITIONAL_RELAY_DELAY; + }); + }); // We don't have this proof if (!proof) { return nullptr; } - auto proofRegistrationTime = g_avalanche->getProofRegistrationTime(proofid); - - // If we know that proof for long enough, allow for requesting it - if (proofRegistrationTime <= now - UNCONDITIONAL_RELAY_DELAY) { + if (send_unconditionally) { return proof; } - { - LOCK(cs_main); - // Otherwise, the proofs must have been announced recently. - if (State(peer.GetId()) - ->m_recently_announced_proofs.contains(proofid)) { - return proof; - } + // Otherwise, the proofs must have been announced recently. + LOCK(cs_main); + if (State(peer.GetId())->m_recently_announced_proofs.contains(proofid)) { + return proof; } return nullptr;