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 @@ -130,7 +131,7 @@ PeerId nextPeerId = 0; PeerSet peers; - std::unordered_map utxos; + ProofPool pool; using NodeSet = boost::multi_index_container< Node, @@ -174,6 +175,9 @@ std::unordered_set m_unbroadcast_proofids; public: + // Limit the pool to a depth to a single proof per utxo + PeerManager() : pool(1) {} + /** * Node API. */ @@ -273,6 +277,8 @@ bool addNodeToPeer(const PeerSet::iterator &it); bool removeNodeFromPeer(const PeerSet::iterator &it, uint32_t count = 1); + bool addProofToPool(const ProofRef &proof); + friend struct ::avalanche::TestPeerManager; }; diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -147,6 +147,32 @@ state.GetResult() == ProofValidationResult::HEIGHT_MISMATCH; } +bool PeerManager::addProofToPool(const ProofRef &proof) { + std::set shiftedProofs; + if (!pool.addProof(proof, shiftedProofs)) { + return false; + } + + // The proof is attached with no conflict + if (shiftedProofs.empty()) { + return true; + } + + // For now, if there is a conflict, just cleanup the mess. + pool.removeProof(proof); + + // Restore the previous state. Since we removed the conflicting proof, this + // can never fail. + for (const ProofRef &p : shiftedProofs) { + std::set inconsistent; + const bool attached = pool.addProof(p, inconsistent); + assert(attached); + assert(inconsistent.empty()); + } + + return false; +} + bool PeerManager::registerProof(const ProofRef &proof) { if (exists(proof->getId())) { // The proof is already registered, or orphaned. @@ -170,6 +196,12 @@ return false; } + if (!addProofToPool(proof)) { + // Rejected due to conflicting proofs + orphanProofs.addProof(proof); + return false; + } + return createPeer(proof); } @@ -265,35 +297,6 @@ // New peer means new peerid! const PeerId peerid = nextPeerId++; - // Attach UTXOs to this proof. - std::unordered_set conflicting_proofs; - for (const auto &s : proof->getStakes()) { - auto p = utxos.emplace(s.getStake().getUTXO(), proof); - if (!p.second) { - // We have a collision with an existing proof. - conflicting_proofs.insert(p.first->second); - } - } - - // For now, if there is a conflict, just cleanup the mess. - if (conflicting_proofs.size() > 0) { - for (const auto &s : proof->getStakes()) { - auto it = utxos.find(s.getStake().getUTXO()); - assert(it != utxos.end()); - - // We need to delete that one. - if (it->second->getId() == proofid) { - utxos.erase(it); - } - } - - // Orphan the proof so it can be pulled back if the conflicting ones are - // invalidated. - orphanProofs.addProof(proof); - - return false; - } - // We have no peer for this proof, time to create it. auto inserted = peers.emplace(peerid, proof); assert(inserted.second); @@ -341,11 +344,8 @@ nview.upper_bound(boost::make_tuple( peerid, std::chrono::steady_clock::now()))); - // Release UTXOs attached to this proof. - for (const auto &s : it->proof->getStakes()) { - bool deleted = utxos.erase(s.getStake().getUTXO()) > 0; - assert(deleted); - } + bool removed = pool.removeProof(it->proof); + assert(removed); m_unbroadcast_proofids.erase(it->proof->getId()); @@ -432,18 +432,17 @@ return false; } - // Check utxos consistency - for (const auto &ss : p.proof->getStakes()) { - auto it = utxos.find(ss.getStake().getUTXO()); - - if (it == utxos.end()) { - return false; - } - if (it->second->getId() != p.getProofId()) { + // Make sure all the proof utxos are in the pool + for (const SignedStake &ss : p.proof->getStakes()) { + const COutPoint &outpoint = ss.getStake().getUTXO(); + if (!pool.forUtxo(outpoint, [&](const auto &proofs) { + return proofs.find(p.proof) != proofs.end(); + })) { return false; } - if (!peersUtxos.emplace(it->first).second) { + // Duplicated utxo + if (!peersUtxos.emplace(outpoint).second) { return false; } } @@ -484,14 +483,14 @@ } } - // Check there is no dangling utxo - for (const auto &[outpoint, proof] : utxos) { - if (!peersUtxos.count(outpoint)) { - return false; - } - } + // Check there is no dangling utxo in the pool + bool danglingUtxo = false; + pool.forEachUtxo([&](const COutPoint &outpoint, const auto &proof) { + danglingUtxo = + danglingUtxo || peersUtxos.find(outpoint) == peersUtxos.end(); + }); - return true; + return !danglingUtxo; } PeerId selectPeerImpl(const std::vector &slots, const uint64_t slot, diff --git a/src/avalanche/proofpool.h b/src/avalanche/proofpool.h --- a/src/avalanche/proofpool.h +++ b/src/avalanche/proofpool.h @@ -43,6 +43,12 @@ auto it = utxos.find(outpoint); return it != utxos.end() && fun(it->second); } + + template void forEachUtxo(Callable fun) const { + for (const auto &[outpoint, proofs] : utxos) { + fun(outpoint, proofs); + } + } }; } // namespace avalanche