Fix mempool package tracking edge case

CalculateMemPoolAncestors was always looping over a transaction's inputs
to find in-mempool parents.  When adding a new transaction, this is the
correct behavior, but when removing a transaction, we want to use the
ancestor set that would be calculated by walking mapLinks (which should
in general be the same set, except during a reorg when the mempool is
in an inconsistent state, and the mapLinks-based calculation would be the
correct one).

2 files changed, 39 insertions, 16 deletions

diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index 2f603e3c9f..1370cab0c0 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -159,26 +159,30 @@ void CTxMemPool::UpdateTransactionsFromBlock(const std::vector<uint256> &vHashes
     }
 }
 
-bool CTxMemPool::CalculateMemPoolAncestors(const CTxMemPoolEntry &entry, setEntries &setAncestors, uint64_t limitAncestorCount, uint64_t limitAncestorSize, uint64_t limitDescendantCount, uint64_t limitDescendantSize, std::string &errString)
+bool CTxMemPool::CalculateMemPoolAncestors(const CTxMemPoolEntry &entry, setEntries &setAncestors, uint64_t limitAncestorCount, uint64_t limitAncestorSize, uint64_t limitDescendantCount, uint64_t limitDescendantSize, std::string &errString, bool fSearchForParents /* = true */)
 {
     setEntries parentHashes;
     const CTransaction &tx = entry.GetTx();
 
-    // Get parents of this transaction that are in the mempool
-    // Entry may or may not already be in the mempool, and GetMemPoolParents()
-    // is only valid for entries in the mempool, so we iterate mapTx to find
-    // parents.
-    // TODO: optimize this so that we only check limits and walk
-    // tx.vin when called on entries not already in the mempool.
-    for (unsigned int i = 0; i < tx.vin.size(); i++) {
-        txiter piter = mapTx.find(tx.vin[i].prevout.hash);
-        if (piter != mapTx.end()) {
-            parentHashes.insert(piter);
-            if (parentHashes.size() + 1 > limitAncestorCount) {
-                errString = strprintf("too many unconfirmed parents [limit: %u]", limitAncestorCount);
-                return false;
+    if (fSearchForParents) {
+        // Get parents of this transaction that are in the mempool
+        // GetMemPoolParents() is only valid for entries in the mempool, so we
+        // iterate mapTx to find parents.
+        for (unsigned int i = 0; i < tx.vin.size(); i++) {
+            txiter piter = mapTx.find(tx.vin[i].prevout.hash);
+            if (piter != mapTx.end()) {
+                parentHashes.insert(piter);
+                if (parentHashes.size() + 1 > limitAncestorCount) {
+                    errString = strprintf("too many unconfirmed parents [limit: %u]", limitAncestorCount);
+                    return false;
+                }
             }
         }
+    } else {
+        // If we're not searching for parents, we require this to be an
+        // entry in the mempool already.
+        txiter it = mapTx.iterator_to(entry);
+        parentHashes = GetMemPoolParents(it);
     }
 
     size_t totalSizeWithAncestors = entry.GetTxSize();
@@ -249,7 +253,24 @@ void CTxMemPool::UpdateForRemoveFromMempool(const setEntries &entriesToRemove)
         setEntries setAncestors;
         const CTxMemPoolEntry &entry = *removeIt;
         std::string dummy;
-        CalculateMemPoolAncestors(entry, setAncestors, nNoLimit, nNoLimit, nNoLimit, nNoLimit, dummy);
+        // Since this is a tx that is already in the mempool, we can call CMPA
+        // with fSearchForParents = false.  If the mempool is in a consistent
+        // state, then using true or false should both be correct, though false
+        // should be a bit faster.
+        // However, if we happen to be in the middle of processing a reorg, then
+        // the mempool can be in an inconsistent state.  In this case, the set
+        // of ancestors reachable via mapLinks will be the same as the set of 
+        // ancestors whose packages include this transaction, because when we
+        // add a new transaction to the mempool in addUnchecked(), we assume it
+        // has no children, and in the case of a reorg where that assumption is
+        // false, the in-mempool children aren't linked to the in-block tx's
+        // until UpdateTransactionsFromBlock() is called.
+        // So if we're being called during a reorg, ie before
+        // UpdateTransactionsFromBlock() has been called, then mapLinks[] will
+        // differ from the set of mempool parents we'd calculate by searching,
+        // and it's important that we use the mapLinks[] notion of ancestor
+        // transactions as the set of things to update for removal.
+        CalculateMemPoolAncestors(entry, setAncestors, nNoLimit, nNoLimit, nNoLimit, nNoLimit, dummy, false);
         // Note that UpdateAncestorsOf severs the child links that point to
         // removeIt in the entries for the parents of removeIt.  This is
         // fine since we don't need to use the mempool children of any entries
diff --git a/src/txmempool.h b/src/txmempool.h
index f0c3f7e0f1..c0eef0dd22 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -392,8 +392,10 @@ public:
      *  limitDescendantCount = max number of descendants any ancestor can have
      *  limitDescendantSize = max size of descendants any ancestor can have
      *  errString = populated with error reason if any limits are hit
+     *  fSearchForParents = whether to search a tx's vin for in-mempool parents, or
+     *    look up parents from mapLinks. Must be true for entries not in the mempool
      */
-    bool CalculateMemPoolAncestors(const CTxMemPoolEntry &entry, setEntries &setAncestors, uint64_t limitAncestorCount, uint64_t limitAncestorSize, uint64_t limitDescendantCount, uint64_t limitDescendantSize, std::string &errString);
+    bool CalculateMemPoolAncestors(const CTxMemPoolEntry &entry, setEntries &setAncestors, uint64_t limitAncestorCount, uint64_t limitAncestorSize, uint64_t limitDescendantCount, uint64_t limitDescendantSize, std::string &errString, bool fSearchForParents = true);
 
     unsigned long size()
     {