From 666ff7bff94ac90a21b9e440741a017f42c692c0 Mon Sep 17 00:00:00 2001
From: Kittywhiskers Van Gogh <63189531+kittywhiskers@users.noreply.github.com>
Date: Sat, 10 Sep 2022 02:15:39 +0300
Subject: [PATCH] merge bitcoin#14193: Add missing mempool locks

Co-authored-by: "UdjinM6 <UdjinM6@users.noreply.github.com>"
---
 src/governance/governance.cpp       |  12 ++-
 src/governance/object.cpp           |   5 +-
 src/net_processing.cpp              |   2 +-
 src/rpc/governance.cpp              |   2 +-
 src/rpc/mining.cpp                  |   3 +-
 src/test/validation_block_tests.cpp | 157 +++++++++++++++++++++++++++-
 src/txmempool.cpp                   |  20 ++--
 src/txmempool.h                     |  36 +++----
 src/validation.cpp                  |  18 ++--
 src/validation.h                    |   7 +-
 10 files changed, 206 insertions(+), 56 deletions(-)

diff --git a/src/governance/governance.cpp b/src/governance/governance.cpp
index 4a351f8257..b023ea707c 100644
--- a/src/governance/governance.cpp
+++ b/src/governance/governance.cpp
@@ -142,7 +142,8 @@ void CGovernanceManager::ProcessMessage(CNode* pfrom, const std::string& msg_typ
             return;
         }
 
-        LOCK2(cs_main, cs);
+        LOCK2(cs_main, ::mempool.cs); // Lock mempool because of GetTransaction deep inside
+        LOCK(cs);
 
         if (mapObjects.count(nHash) || mapPostponedObjects.count(nHash) || mapErasedGovernanceObjects.count(nHash)) {
             // TODO - print error code? what if it's GOVOBJ_ERROR_IMMATURE?
@@ -261,7 +262,8 @@ void CGovernanceManager::AddGovernanceObject(CGovernanceObject& govobj, CConnman
 
     govobj.UpdateSentinelVariables(); //this sets local vars in object
 
-    LOCK2(cs_main, cs);
+    LOCK2(cs_main, ::mempool.cs); // Lock mempool because of GetTransaction deep inside
+    LOCK(cs);
     std::string strError;
 
     // MAKE SURE THIS OBJECT IS OK
@@ -320,7 +322,8 @@ void CGovernanceManager::UpdateCachesAndClean()
 
     std::vector<uint256> vecDirtyHashes = mmetaman.GetAndClearDirtyGovernanceObjectHashes();
 
-    LOCK2(cs_main, cs);
+    LOCK2(cs_main, ::mempool.cs); // Lock mempool because of GetTransaction deep inside
+    LOCK(cs);
 
     for (const uint256& nHash : vecDirtyHashes) {
         auto it = mapObjects.find(nHash);
@@ -835,7 +838,8 @@ void CGovernanceManager::CheckPostponedObjects(CConnman& connman)
 {
     if (!masternodeSync.IsSynced()) return;
 
-    LOCK2(cs_main, cs);
+    LOCK2(cs_main, ::mempool.cs); // Lock mempool because of GetTransaction deep inside
+    LOCK(cs);
 
     // Check postponed proposals
     for (auto it = mapPostponedObjects.begin(); it != mapPostponedObjects.end();) {
diff --git a/src/governance/object.cpp b/src/governance/object.cpp
index d8b3b67782..49ae94ec80 100644
--- a/src/governance/object.cpp
+++ b/src/governance/object.cpp
@@ -438,7 +438,7 @@ UniValue CGovernanceObject::ToJson() const
 
 void CGovernanceObject::UpdateLocalValidity()
 {
-    LOCK(cs_main);
+    AssertLockHeld(cs_main);
     // THIS DOES NOT CHECK COLLATERAL, THIS IS CHECKED UPON ORIGINAL ARRIVAL
     fCachedLocalValidity = IsValidLocally(strLocalValidityError, false);
 }
@@ -526,6 +526,9 @@ CAmount CGovernanceObject::GetMinCollateralFee(bool fork_active) const
 
 bool CGovernanceObject::IsCollateralValid(std::string& strError, bool& fMissingConfirmations) const
 {
+    AssertLockHeld(cs_main);
+    AssertLockHeld(::mempool.cs); // because of GetTransaction
+
     strError = "";
     fMissingConfirmations = false;
     uint256 nExpectedHash = GetHash();
diff --git a/src/net_processing.cpp b/src/net_processing.cpp
index 452cfa2fb0..a9470f2c93 100644
--- a/src/net_processing.cpp
+++ b/src/net_processing.cpp
@@ -4629,7 +4629,7 @@ bool PeerLogicValidation::SendMessages(CNode* pto)
         //
         std::vector<CInv> vInv;
         {
-            LOCK(pto->cs_inventory);
+            LOCK2(mempool.cs, pto->cs_inventory);
 
             size_t reserve = std::min<size_t>(pto->setInventoryTxToSend.size(), INVENTORY_BROADCAST_MAX_PER_1MB_BLOCK * MaxBlockSize() / 1000000);
             reserve = std::max<size_t>(reserve, pto->vInventoryBlockToSend.size());
diff --git a/src/rpc/governance.cpp b/src/rpc/governance.cpp
index 996349e00c..22751687eb 100644
--- a/src/rpc/governance.cpp
+++ b/src/rpc/governance.cpp
@@ -392,7 +392,7 @@ static UniValue gobject_submit(const JSONRPCRequest& request)
             g_txindex->BlockUntilSyncedToCurrentChain();
         }
 
-        LOCK(cs_main);
+        LOCK2(cs_main, ::mempool.cs);
         if (!govobj.IsValidLocally(strError, fMissingConfirmations, true) && !fMissingConfirmations) {
             LogPrintf("gobject(submit) -- Object submission rejected because object is not valid - hash = %s, strError = %s\n", strHash, strError);
             throw JSONRPCError(RPC_INTERNAL_ERROR, "Governance object is not valid - " + strHash + " - " + strError);
diff --git a/src/rpc/mining.cpp b/src/rpc/mining.cpp
index e28aacfe8a..bbad597799 100644
--- a/src/rpc/mining.cpp
+++ b/src/rpc/mining.cpp
@@ -501,7 +501,7 @@ static UniValue getblocktemplate(const JSONRPCRequest& request)
             nTransactionsUpdatedLastLP = nTransactionsUpdatedLast;
         }
 
-        // Release the wallet and main lock while waiting
+        // Release lock while waiting
         LEAVE_CRITICAL_SECTION(cs_main);
         {
             checktxtime = std::chrono::steady_clock::now() + std::chrono::minutes(1);
@@ -512,6 +512,7 @@ static UniValue getblocktemplate(const JSONRPCRequest& request)
                 if (g_best_block_cv.wait_until(lock, checktxtime) == std::cv_status::timeout)
                 {
                     // Timeout: Check transactions for update
+                    // without holding ::mempool.cs to avoid deadlocks
                     if (mempool.GetTransactionsUpdated() != nTransactionsUpdatedLastLP)
                         break;
                     checktxtime += std::chrono::seconds(10);
diff --git a/src/test/validation_block_tests.cpp b/src/test/validation_block_tests.cpp
index 99d3c4beb8..f5d59fccb5 100644
--- a/src/test/validation_block_tests.cpp
+++ b/src/test/validation_block_tests.cpp
@@ -5,11 +5,13 @@
 #include <boost/test/unit_test.hpp>
 
 #include <chainparams.h>
+#include <consensus/consensus.h>
 #include <consensus/merkle.h>
 #include <consensus/validation.h>
 #include <miner.h>
 #include <pow.h>
 #include <random.h>
+#include <script/standard.h>
 #include <test/util/setup_common.h>
 #include <util/time.h>
 #include <validation.h>
@@ -17,6 +19,8 @@
 
 #include <thread>
 
+static const std::vector<unsigned char> V_OP_TRUE{OP_TRUE};
+
 BOOST_FIXTURE_TEST_SUITE(validation_block_tests, RegTestingSetup)
 
 struct TestSubscriber : public CValidationInterface {
@@ -59,6 +63,23 @@ std::shared_ptr<CBlock> Block(const uint256& prev_hash)
     pblock->hashPrevBlock = prev_hash;
     pblock->nTime = ++time;
 
+    pubKey.clear();
+    {
+        pubKey << OP_HASH160 << ToByteVector(CScriptID(CScript() << OP_TRUE))
+               << OP_EQUAL;
+    }
+    // Make the coinbase transaction with two outputs:
+    // One zero-value one that has a unique pubkey to make sure that blocks at
+    // the same height can have a different hash. Another one that has the
+    // coinbase reward in a P2SH with OP_TRUE as scriptPubKey to make it easy to
+    // spend
+    CMutableTransaction txCoinbase(*pblock->vtx[0]);
+    txCoinbase.vout.resize(2);
+    txCoinbase.vout[1].scriptPubKey = pubKey;
+    txCoinbase.vout[1].nValue = txCoinbase.vout[0].nValue;
+    txCoinbase.vout[0].nValue = 0;
+    pblock->vtx[0] = MakeTransactionRef(std::move(txCoinbase));
+
     return pblock;
 }
 
@@ -74,13 +95,13 @@ std::shared_ptr<CBlock> FinalizeBlock(std::shared_ptr<CBlock> pblock)
 }
 
 // construct a valid block
-const std::shared_ptr<const CBlock> GoodBlock(const uint256& prev_hash)
+std::shared_ptr<const CBlock> GoodBlock(const uint256& prev_hash)
 {
     return FinalizeBlock(Block(prev_hash));
 }
 
 // construct an invalid block (but with a valid header)
-const std::shared_ptr<const CBlock> BadBlock(const uint256& prev_hash)
+std::shared_ptr<const CBlock> BadBlock(const uint256& prev_hash)
 {
     auto pblock = Block(prev_hash);
 
@@ -180,4 +201,136 @@ BOOST_AUTO_TEST_CASE(processnewblock_signals_ordering)
     BOOST_CHECK_EQUAL(sub.m_expected_tip, ::ChainActive().Tip()->GetBlockHash());
 }
 
+/**
+ * Test that mempool updates happen atomically with reorgs.
+ *
+ * This prevents RPC clients, among others, from retrieving immediately-out-of-date mempool data
+ * during large reorgs.
+ *
+ * The test verifies this by creating a chain of `num_txs` blocks, matures their coinbases, and then
+ * submits txns spending from their coinbase to the mempool. A fork chain is then processed,
+ * invalidating the txns and evicting them from the mempool.
+ *
+ * We verify that the mempool updates atomically by polling it continuously
+ * from another thread during the reorg and checking that its size only changes
+ * once. The size changing exactly once indicates that the polling thread's
+ * view of the mempool is either consistent with the chain state before reorg,
+ * or consistent with the chain state after the reorg, and not just consistent
+ * with some intermediate state during the reorg.
+ */
+BOOST_AUTO_TEST_CASE(mempool_locks_reorg)
+{
+    bool ignored;
+    auto ProcessBlock = [&ignored](std::shared_ptr<const CBlock> block) -> bool {
+        return ProcessNewBlock(Params(), block, /* fForceProcessing */ true, /* fNewBlock */ &ignored);
+    };
+
+    // Process all mined blocks
+    BOOST_REQUIRE(ProcessBlock(std::make_shared<CBlock>(Params().GenesisBlock())));
+    auto last_mined = GoodBlock(Params().GenesisBlock().GetHash());
+    BOOST_REQUIRE(ProcessBlock(last_mined));
+
+    // Run the test multiple times
+    for (int test_runs = 3; test_runs > 0; --test_runs) {
+        BOOST_CHECK_EQUAL(last_mined->GetHash(), ::ChainActive().Tip()->GetBlockHash());
+
+        // Later on split from here
+        const uint256 split_hash{last_mined->hashPrevBlock};
+
+        // Create a bunch of transactions to spend the miner rewards of the
+        // most recent blocks
+        std::vector<CTransactionRef> txs;
+        for (int num_txs = 22; num_txs > 0; --num_txs) {
+            CMutableTransaction mtx;
+            mtx.vin.push_back(
+                CTxIn(COutPoint(last_mined->vtx[0]->GetHash(), 1),
+                      CScript() << ToByteVector(CScript() << OP_TRUE)));
+            // Two outputs to make sure the transaction is larger than 100 bytes
+            for (int i = 1; i < 3; ++i) {
+                mtx.vout.emplace_back(
+                    CTxOut(50000,
+                           CScript() << OP_DUP << OP_HASH160
+                                     << ToByteVector(CScriptID(CScript() << i))
+                                     << OP_EQUALVERIFY << OP_CHECKSIG));
+            }
+            txs.push_back(MakeTransactionRef(mtx));
+
+            last_mined = GoodBlock(last_mined->GetHash());
+            BOOST_REQUIRE(ProcessBlock(last_mined));
+        }
+
+        // Mature the inputs of the txs
+        for (int j = COINBASE_MATURITY; j > 0; --j) {
+            last_mined = GoodBlock(last_mined->GetHash());
+            BOOST_REQUIRE(ProcessBlock(last_mined));
+        }
+
+        // Mine a reorg (and hold it back) before adding the txs to the mempool
+        const uint256 tip_init{last_mined->GetHash()};
+
+        std::vector<std::shared_ptr<const CBlock>> reorg;
+        last_mined = GoodBlock(split_hash);
+        reorg.push_back(last_mined);
+        for (size_t j = COINBASE_MATURITY + txs.size() + 1; j > 0; --j) {
+            last_mined = GoodBlock(last_mined->GetHash());
+            reorg.push_back(last_mined);
+        }
+
+        // Add the txs to the tx pool
+        {
+            LOCK(cs_main);
+            CValidationState state;
+            for (const auto& tx : txs) {
+                BOOST_REQUIRE(AcceptToMemoryPool(
+                    ::mempool,
+                    state,
+                    tx,
+                    /* pfMissingInputs */ &ignored,
+                    /* bypass_limits */ false,
+                    /* nAbsurdFee */ 0));
+            }
+        }
+
+        // Check that all txs are in the pool
+        {
+            LOCK(::mempool.cs);
+            BOOST_CHECK_EQUAL(::mempool.mapTx.size(), txs.size());
+        }
+
+        // Run a thread that simulates an RPC caller that is polling while
+        // validation is doing a reorg
+        std::thread rpc_thread{[&]() {
+            // This thread is checking that the mempool either contains all of
+            // the transactions invalidated by the reorg, or none of them, and
+            // not some intermediate amount.
+            while (true) {
+                LOCK(::mempool.cs);
+                if (::mempool.mapTx.size() == 0) {
+                    // We are done with the reorg
+                    break;
+                }
+                // Internally, we might be in the middle of the reorg, but
+                // externally the reorg to the most-proof-of-work chain should
+                // be atomic. So the caller assumes that the returned mempool
+                // is consistent. That is, it has all txs that were there
+                // before the reorg.
+                assert(::mempool.mapTx.size() == txs.size());
+                continue;
+            }
+            LOCK(cs_main);
+            // We are done with the reorg, so the tip must have changed
+            assert(tip_init != ::ChainActive().Tip()->GetBlockHash());
+        }};
+
+        // Submit the reorg in this thread to invalidate and remove the txs from the tx pool
+        for (const auto& b : reorg) {
+            ProcessBlock(b);
+        }
+        // Check that the reorg was eventually successful
+        BOOST_CHECK_EQUAL(last_mined->GetHash(), ::ChainActive().Tip()->GetBlockHash());
+
+        // We can join the other thread, which returns when the reorg was successful
+        rpc_thread.join();
+    }
+}
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index 4c0ac5a1f9..ecada5d9e4 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -112,7 +112,7 @@ void CTxMemPool::UpdateForDescendants(txiter updateIt, cacheMap &cachedDescendan
 // for each such descendant, also update the ancestor state to include the parent.
 void CTxMemPool::UpdateTransactionsFromBlock(const std::vector<uint256> &vHashesToUpdate)
 {
-    LOCK(cs);
+    AssertLockHeld(cs);
     // For each entry in vHashesToUpdate, store the set of in-mempool, but not
     // in-vHashesToUpdate transactions, so that we don't have to recalculate
     // descendants when we come across a previously seen entry.
@@ -332,8 +332,8 @@ void CTxMemPoolEntry::UpdateAncestorState(int64_t modifySize, CAmount modifyFee,
     assert(int(nSigOpCountWithAncestors) >= 0);
 }
 
-CTxMemPool::CTxMemPool(CBlockPolicyEstimator* estimator) :
-    nTransactionsUpdated(0), minerPolicyEstimator(estimator), m_epoch(0), m_has_epoch_guard(false)
+CTxMemPool::CTxMemPool(CBlockPolicyEstimator* estimator)
+    : nTransactionsUpdated(0), minerPolicyEstimator(estimator), m_epoch(0), m_has_epoch_guard(false)
 {
     _clear(); //lock free clear
 
@@ -351,13 +351,11 @@ bool CTxMemPool::isSpent(const COutPoint& outpoint) const
 
 unsigned int CTxMemPool::GetTransactionsUpdated() const
 {
-    LOCK(cs);
     return nTransactionsUpdated;
 }
 
 void CTxMemPool::AddTransactionsUpdated(unsigned int n)
 {
-    LOCK(cs);
     nTransactionsUpdated += n;
 }
 
@@ -722,8 +720,7 @@ void CTxMemPool::CalculateDescendants(txiter entryit, setEntries& setDescendants
 void CTxMemPool::removeRecursive(const CTransaction &origTx, MemPoolRemovalReason reason)
 {
     // Remove transaction from memory pool
-    {
-        LOCK(cs);
+    AssertLockHeld(cs);
         setEntries txToRemove;
         txiter origit = mapTx.find(origTx.GetHash());
         if (origit != mapTx.end()) {
@@ -748,13 +745,12 @@ void CTxMemPool::removeRecursive(const CTransaction &origTx, MemPoolRemovalReaso
         }
 
         RemoveStaged(setAllRemoves, false, reason);
-    }
 }
 
 void CTxMemPool::removeForReorg(const CCoinsViewCache *pcoins, unsigned int nMemPoolHeight, int flags) EXCLUSIVE_LOCKS_REQUIRED(cs_main)
 {
     // Remove transactions spending a coinbase which are now immature and no-longer-final transactions
-    LOCK(cs);
+    AssertLockHeld(cs);
     setEntries txToRemove;
     for (indexed_transaction_set::const_iterator it = mapTx.begin(); it != mapTx.end(); it++) {
         const CTransaction& tx = it->GetTx();
@@ -969,7 +965,7 @@ void CTxMemPool::removeProTxConflicts(const CTransaction &tx)
  */
 void CTxMemPool::removeForBlock(const std::vector<CTransactionRef>& vtx, unsigned int nBlockHeight)
 {
-    LOCK(cs);
+    AssertLockHeld(cs);
     std::vector<const CTxMemPoolEntry*> entries;
     for (const auto& tx : vtx)
     {
@@ -1436,7 +1432,7 @@ void CTxMemPool::RemoveStaged(setEntries &stage, bool updateDescendants, MemPool
 }
 
 int CTxMemPool::Expire(int64_t time) {
-    LOCK(cs);
+    AssertLockHeld(cs);
     indexed_transaction_set::index<entry_time>::type::iterator it = mapTx.get<entry_time>().begin();
     setEntries toremove;
     while (it != mapTx.get<entry_time>().end() && it->GetTime() < time) {
@@ -1534,7 +1530,7 @@ void CTxMemPool::trackPackageRemoved(const CFeeRate& rate) {
 }
 
 void CTxMemPool::TrimToSize(size_t sizelimit, std::vector<COutPoint>* pvNoSpendsRemaining) {
-    LOCK(cs);
+    AssertLockHeld(cs);
 
     unsigned nTxnRemoved = 0;
     CFeeRate maxFeeRateRemoved(0);
diff --git a/src/txmempool.h b/src/txmempool.h
index 37772ccfc2..0532498742 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -6,11 +6,12 @@
 #ifndef BITCOIN_TXMEMPOOL_H
 #define BITCOIN_TXMEMPOOL_H
 
-#include <set>
+#include <atomic>
 #include <map>
-#include <vector>
-#include <utility>
+#include <set>
 #include <string>
+#include <utility>
+#include <vector>
 
 #include <addressindex.h>
 #include <spentindex.h>
@@ -449,7 +450,7 @@ class CTxMemPool
 {
 private:
     uint32_t nCheckFrequency GUARDED_BY(cs); //!< Value n means that n times in 2^32 we check.
-    unsigned int nTransactionsUpdated; //!< Used by getblocktemplate to trigger CreateNewBlock() invocation
+    std::atomic<unsigned int> nTransactionsUpdated; //!< Used by getblocktemplate to trigger CreateNewBlock() invocation
     CBlockPolicyEstimator* minerPolicyEstimator;
 
     uint64_t totalTxSize;      //!< sum of all mempool tx' byte sizes
@@ -521,21 +522,12 @@ public:
      * `mempool.cs` whenever adding transactions to the mempool and whenever
      * changing the chain tip. It's necessary to keep both mutexes locked until
      * the mempool is consistent with the new chain tip and fully populated.
-     *
-     * @par Consistency bug
-     *
-     * The second guarantee above is not currently enforced, but
-     * https://github.com/bitcoin/bitcoin/pull/14193 will fix it. No known code
-     * in bitcoin currently depends on second guarantee, but it is important to
-     * fix for third party code that needs be able to frequently poll the
-     * mempool without locking `cs_main` and without encountering missing
-     * transactions during reorgs.
      */
     mutable RecursiveMutex cs;
     indexed_transaction_set mapTx GUARDED_BY(cs);
 
     using txiter = indexed_transaction_set::nth_index<0>::type::const_iterator;
-    std::vector<std::pair<uint256, txiter> > vTxHashes; //!< All tx hashes/entries in mapTx, in random order
+    std::vector<std::pair<uint256, txiter> > vTxHashes GUARDED_BY(cs); //!< All tx hashes/entries in mapTx, in random order
 
     struct CompareIteratorByHash {
         bool operator()(const txiter &a, const txiter &b) const {
@@ -617,16 +609,16 @@ public:
     bool getSpentIndex(CSpentIndexKey &key, CSpentIndexValue &value);
     bool removeSpentIndex(const uint256 txhash);
 
-    void removeRecursive(const CTransaction &tx, MemPoolRemovalReason reason);
-    void removeForReorg(const CCoinsViewCache *pcoins, unsigned int nMemPoolHeight, int flags) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
-    void removeConflicts(const CTransaction &tx) EXCLUSIVE_LOCKS_REQUIRED(cs);
+    void removeRecursive(const CTransaction& tx, MemPoolRemovalReason reason) EXCLUSIVE_LOCKS_REQUIRED(cs);
+    void removeForReorg(const CCoinsViewCache* pcoins, unsigned int nMemPoolHeight, int flags) EXCLUSIVE_LOCKS_REQUIRED(cs, cs_main);
+    void removeConflicts(const CTransaction& tx) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxPubKeyConflicts(const CTransaction &tx, const CKeyID &keyId) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxPubKeyConflicts(const CTransaction &tx, const CBLSPublicKey &pubKey) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxCollateralConflicts(const CTransaction &tx, const COutPoint &collateralOutpoint) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxSpentCollateralConflicts(const CTransaction &tx) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxKeyChangedConflicts(const CTransaction &tx, const uint256& proTxHash, const uint256& newKeyHash) EXCLUSIVE_LOCKS_REQUIRED(cs);
     void removeProTxConflicts(const CTransaction &tx) EXCLUSIVE_LOCKS_REQUIRED(cs);
-    void removeForBlock(const std::vector<CTransactionRef>& vtx, unsigned int nBlockHeight);
+    void removeForBlock(const std::vector<CTransactionRef>& vtx, unsigned int nBlockHeight) EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     void clear();
     void _clear() EXCLUSIVE_LOCKS_REQUIRED(cs); //lock free
@@ -639,7 +631,7 @@ public:
      * Check that none of this transactions inputs are in the mempool, and thus
      * the tx is not dependent on other mempool transactions to be included in a block.
      */
-    bool HasNoInputsOf(const CTransaction& tx) const;
+    bool HasNoInputsOf(const CTransaction& tx) const EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     /** Affect CreateNewBlock prioritisation of transactions */
     void PrioritiseTransaction(const uint256& hash, const CAmount& nFeeDelta);
@@ -673,7 +665,7 @@ public:
      *  for).  Note: vHashesToUpdate should be the set of transactions from the
      *  disconnected block that have been accepted back into the mempool.
      */
-    void UpdateTransactionsFromBlock(const std::vector<uint256>& vHashesToUpdate) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
+    void UpdateTransactionsFromBlock(const std::vector<uint256>& vHashesToUpdate) EXCLUSIVE_LOCKS_REQUIRED(cs, cs_main);
 
     /** Try to calculate all in-mempool ancestors of entry.
      *  (these are all calculated including the tx itself)
@@ -704,10 +696,10 @@ public:
       *  pvNoSpendsRemaining, if set, will be populated with the list of outpoints
       *  which are not in mempool which no longer have any spends in this mempool.
       */
-    void TrimToSize(size_t sizelimit, std::vector<COutPoint>* pvNoSpendsRemaining=nullptr);
+    void TrimToSize(size_t sizelimit, std::vector<COutPoint>* pvNoSpendsRemaining = nullptr) EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     /** Expire all transaction (and their dependencies) in the mempool older than time. Return the number of removed transactions. */
-    int Expire(int64_t time);
+    int Expire(int64_t time) EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     /**
      * Calculate the ancestor and descendant count for the given transaction.
diff --git a/src/validation.cpp b/src/validation.cpp
index cc6a38aa91..d92b344c11 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -394,8 +394,7 @@ bool ContextualCheckTransaction(const CTransaction& tx, CValidationState &state,
 // Returns the script flags which should be checked for a given block
 static unsigned int GetBlockScriptFlags(const CBlockIndex* pindex, const Consensus::Params& chainparams);
 
-static void LimitMempoolSize(CTxMemPool& pool, size_t limit, unsigned long age)
-    EXCLUSIVE_LOCKS_REQUIRED(pool.cs, ::cs_main)
+static void LimitMempoolSize(CTxMemPool& pool, size_t limit, unsigned long age) EXCLUSIVE_LOCKS_REQUIRED(pool.cs, ::cs_main)
 {
     int expired = pool.Expire(GetTime() - age);
     if (expired != 0) {
@@ -433,7 +432,7 @@ static bool IsCurrentForFeeEstimation() EXCLUSIVE_LOCKS_REQUIRED(cs_main)
  * and instead just erase from the mempool as needed.
  */
 
-static void UpdateMempoolForReorg(DisconnectedBlockTransactions &disconnectpool, bool fAddToMempool) EXCLUSIVE_LOCKS_REQUIRED(cs_main)
+static void UpdateMempoolForReorg(DisconnectedBlockTransactions& disconnectpool, bool fAddToMempool) EXCLUSIVE_LOCKS_REQUIRED(cs_main, ::mempool.cs)
 {
     AssertLockHeld(cs_main);
     std::vector<uint256> vHashUpdate;
@@ -469,11 +468,8 @@ static void UpdateMempoolForReorg(DisconnectedBlockTransactions &disconnectpool,
     // We also need to remove any now-immature transactions
     mempool.removeForReorg(&::ChainstateActive().CoinsTip(), ::ChainActive().Tip()->nHeight + 1, STANDARD_LOCKTIME_VERIFY_FLAGS);
 
-    {
-        LOCK(mempool.cs);
-        // Re-limit mempool size, in case we added any transactions
-        LimitMempoolSize(mempool, gArgs.GetArg("-maxmempool", DEFAULT_MAX_MEMPOOL_SIZE) * 1000000, gArgs.GetArg("-mempoolexpiry", DEFAULT_MEMPOOL_EXPIRY) * 60 * 60);
-    }
+    // Re-limit mempool size, in case we added any transactions
+    LimitMempoolSize(mempool, gArgs.GetArg("-maxmempool", DEFAULT_MAX_MEMPOOL_SIZE) * 1000000, gArgs.GetArg("-mempoolexpiry", DEFAULT_MEMPOOL_EXPIRY) * 60 * 60);
 }
 
 // Used to avoid mempool polluting consensus critical paths if CCoinsViewMempool
@@ -3097,7 +3093,7 @@ bool CChainState::ActivateBestChain(CValidationState &state, const CChainParams&
         LimitValidationInterfaceQueue();
 
         {
-            LOCK(cs_main);
+            LOCK2(cs_main, ::mempool.cs); // Lock transaction pool for at least as long as it takes for connectTrace to be consumed
             CBlockIndex* starting_tip = m_chain.Tip();
             bool blocks_connected = false;
             do {
@@ -3260,6 +3256,7 @@ bool CChainState::InvalidateBlock(CValidationState& state, const CChainParams& c
         LimitValidationInterfaceQueue();
 
         LOCK(cs_main);
+        LOCK(::mempool.cs); // Lock for as long as disconnectpool is in scope to make sure UpdateMempoolForReorg is called after DisconnectTip without unlocking in between
         if (!m_chain.Contains(pindex)) break;
         pindex_was_in_chain = true;
         CBlockIndex *invalid_walk_tip = m_chain.Tip();
@@ -3414,6 +3411,8 @@ bool CChainState::MarkConflictingBlock(CValidationState& state, const CChainPara
         pindexBestHeader = pindexBestHeader->pprev;
     }
 
+    {
+    LOCK(::mempool.cs); // Lock for as long as disconnectpool is in scope to make sure UpdateMempoolForReorg is called after DisconnectTip without unlocking in between
     DisconnectedBlockTransactions disconnectpool;
     while (m_chain.Contains(pindex)) {
         const CBlockIndex* pindexOldTip = m_chain.Tip();
@@ -3446,6 +3445,7 @@ bool CChainState::MarkConflictingBlock(CValidationState& state, const CChainPara
     // DisconnectTip will add transactions to disconnectpool; try to add these
     // back to the mempool.
     UpdateMempoolForReorg(disconnectpool, true);
+    } // ::mempool.cs
 
     // The resulting new best tip may not be in setBlockIndexCandidates anymore, so
     // add it again.
diff --git a/src/validation.h b/src/validation.h
index 21090f541d..7bf9e28b59 100644
--- a/src/validation.h
+++ b/src/validation.h
@@ -20,6 +20,7 @@
 #include <script/script_error.h>
 #include <sync.h>
 #include <txdb.h>
+#include <txmempool.h> // For CTxMemPool::cs
 #include <versionbits.h>
 #include <spentindex.h>
 
@@ -671,7 +672,7 @@ public:
     bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pindex, CCoinsViewCache& view, const CChainParams& chainparams, bool fJustCheck = false) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
 
     // Apply the effects of a block disconnection on the UTXO set.
-    bool DisconnectTip(CValidationState& state, const CChainParams& chainparams, DisconnectedBlockTransactions* disconnectpool) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
+    bool DisconnectTip(CValidationState& state, const CChainParams& chainparams, DisconnectedBlockTransactions* disconnectpool) EXCLUSIVE_LOCKS_REQUIRED(cs_main, ::mempool.cs);
 
     // Manual block validity manipulation:
     bool PreciousBlock(CValidationState& state, const CChainParams& params, CBlockIndex* pindex) LOCKS_EXCLUDED(cs_main);
@@ -713,8 +714,8 @@ public:
         size_t max_mempool_size_bytes) EXCLUSIVE_LOCKS_REQUIRED(::cs_main);
 
 private:
-    bool ActivateBestChainStep(CValidationState& state, const CChainParams& chainparams, CBlockIndex* pindexMostWork, const std::shared_ptr<const CBlock>& pblock, bool& fInvalidFound, ConnectTrace& connectTrace) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
-    bool ConnectTip(CValidationState& state, const CChainParams& chainparams, CBlockIndex* pindexNew, const std::shared_ptr<const CBlock>& pblock, ConnectTrace& connectTrace, DisconnectedBlockTransactions& disconnectpool) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
+    bool ActivateBestChainStep(CValidationState& state, const CChainParams& chainparams, CBlockIndex* pindexMostWork, const std::shared_ptr<const CBlock>& pblock, bool& fInvalidFound, ConnectTrace& connectTrace) EXCLUSIVE_LOCKS_REQUIRED(cs_main, ::mempool.cs);
+    bool ConnectTip(CValidationState& state, const CChainParams& chainparams, CBlockIndex* pindexNew, const std::shared_ptr<const CBlock>& pblock, ConnectTrace& connectTrace, DisconnectedBlockTransactions& disconnectpool) EXCLUSIVE_LOCKS_REQUIRED(cs_main, ::mempool.cs);
 
     void InvalidBlockFound(CBlockIndex* pindex, const CValidationState& state) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
     CBlockIndex* FindMostWorkChain() EXCLUSIVE_LOCKS_REQUIRED(cs_main);