diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -292,7 +293,10 @@ bool isOrphan(const ProofId &proofid) const; private: - bool addToPool(ProofPool &pool, const ProofRef &proof); + bool addToPool(ProofPool &pool, const ProofRef &proof, bool allowOverride); + bool addToPool(ProofPool &pool, const ProofRef &proof) { + return addToPool(pool, proof, true /* allowOverride */); + }; bool removeFromPool(ProofPool &pool, const ProofRef &proof); bool createPeer(const ProofRef &proof); diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -170,7 +170,7 @@ return false; } - if (!addToPool(validProofPool, proof)) { + if (!addToPool(validProofPool, proof, false /* allowOverride */)) { // Orphan the proof so it can be pulled back if the conflicting ones are // invalidated. addToPool(orphanProofPool, proof); @@ -270,7 +270,8 @@ return orphanView.find(proofid) != orphanView.end(); } -bool PeerManager::addToPool(ProofPool &pool, const ProofRef &proof) { +bool PeerManager::addToPool(ProofPool &pool, const ProofRef &proof, + bool allowOverride) { // Attach UTXOs to this proof. std::unordered_set conflicting_proofs; for (const auto &s : proof->getStakes()) { @@ -281,8 +282,14 @@ } } - // For now, if there is a conflict, just cleanup the mess. - if (conflicting_proofs.size() > 0) { + if (conflicting_proofs.size() == 0) { + // No conflict + return true; + } + + ConflictingProofComparator compare; + if (!allowOverride || !compare(proof, *conflicting_proofs.begin())) { + // Not the favorite, clean up the mess for (const auto &s : proof->getStakes()) { auto it = pool.find(s.getStake().getUTXO()); assert(it != pool.end()); @@ -292,9 +299,22 @@ pool.erase(it); } } + + return false; + } + + // This is the favorite, erase the other proofs + for (auto &conflictingProof : conflicting_proofs) { + removeFromPool(pool, conflictingProof); + } + + // Either the utxo is already mapped to this proof or it can be emplaced + for (const auto &s : proof->getStakes()) { + auto p = pool.emplace(s.getStake().getUTXO(), proof); + assert(p.first->proof->getId() == proof->getId()); } - return conflicting_proofs.size() == 0; + return true; } bool PeerManager::removeFromPool(ProofPool &pool, const ProofRef &proof) { 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 @@ -968,9 +968,8 @@ BOOST_CHECK(pm.isOrphan(orphan10->getId())); BOOST_CHECK(!pm.registerProof(orphan20)); - BOOST_CHECK(!pm.isOrphan(orphan20->getId())); - BOOST_CHECK(!pm.exists(orphan20->getId())); - BOOST_CHECK(pm.isOrphan(orphan10->getId())); + BOOST_CHECK(pm.isOrphan(orphan20->getId())); + BOOST_CHECK(!pm.exists(orphan10->getId())); } BOOST_AUTO_TEST_SUITE_END()