From 9d8322a3418ca0fbdc31cb684e392af86663927f Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Fri, 20 Sep 2019 21:41:49 +0100
Subject: [PATCH] Enable RPC tests to wait on mempool notifications

---
 qa/rpc-tests/test_framework/util.py | 10 +++++++++-
 src/rpc/blockchain.cpp              |  4 ++++
 src/txmempool.cpp                   | 16 ++++++++++++++++
 src/txmempool.h                     |  4 ++++
 4 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/qa/rpc-tests/test_framework/util.py b/qa/rpc-tests/test_framework/util.py
index ef933f546..7c804b32c 100644
--- a/qa/rpc-tests/test_framework/util.py
+++ b/qa/rpc-tests/test_framework/util.py
@@ -57,7 +57,7 @@ def sync_blocks(rpc_connections, wait=1):
 def sync_mempools(rpc_connections, wait=1):
     """
     Wait until everybody has the same transactions in their memory
-    pools
+    pools, and has notified all internal listeners of them
     """
     while True:
         pool = set(rpc_connections[0].getrawmempool())
@@ -69,6 +69,14 @@ def sync_mempools(rpc_connections, wait=1):
             break
         time.sleep(wait)
 
+    # Now that the mempools are in sync, wait for the internal
+    # notifications to finish
+    while True:
+        notified = [ x.getmempoolinfo()['fullyNotified'] for x in rpc_connections ]
+        if notified == [ True ] * len(notified):
+            break
+        time.sleep(wait)
+
 bitcoind_processes = {}
 
 def initialize_datadir(dirname, n):
diff --git a/src/rpc/blockchain.cpp b/src/rpc/blockchain.cpp
index a140d42a4..7a581f807 100644
--- a/src/rpc/blockchain.cpp
+++ b/src/rpc/blockchain.cpp
@@ -1185,6 +1185,10 @@ UniValue mempoolInfoToJSON()
     ret.push_back(Pair("bytes", (int64_t) mempool.GetTotalTxSize()));
     ret.push_back(Pair("usage", (int64_t) mempool.DynamicMemoryUsage()));
 
+    if (Params().NetworkIDString() == "regtest") {
+        ret.push_back(Pair("fullyNotified", mempool.IsFullyNotified()));
+    }
+
     return ret;
 }
 
diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index f484cdc9c..325109fe4 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -105,6 +105,7 @@ bool CTxMemPool::addUnchecked(const uint256& hash, const CTxMemPoolEntry &entry,
     mapTx.insert(entry);
     const CTransaction& tx = mapTx.find(hash)->GetTx();
     mapRecentlyAddedTx[tx.GetHash()] = &tx;
+    nRecentlyAddedSequence += 1;
     for (unsigned int i = 0; i < tx.vin.size(); i++)
         mapNextTx[tx.vin[i].prevout] = CInPoint(&tx, i);
     BOOST_FOREACH(const JSDescription &joinsplit, tx.vJoinSplit) {
@@ -729,9 +730,11 @@ bool CTxMemPool::nullifierExists(const uint256& nullifier, ShieldedType type) co
 
 void CTxMemPool::NotifyRecentlyAdded()
 {
+    uint64_t recentlyAddedSequence;
     std::vector<CTransaction> txs;
     {
         LOCK(cs);
+        recentlyAddedSequence = nRecentlyAddedSequence;
         for (const auto& kv : mapRecentlyAddedTx) {
             txs.push_back(*(kv.second));
         }
@@ -753,6 +756,19 @@ void CTxMemPool::NotifyRecentlyAdded()
             PrintExceptionContinue(NULL, "CTxMemPool::NotifyRecentlyAdded()");
         }
     }
+
+    // Update the notified sequence number. We only need this in regtest mode,
+    // and should not lock on cs after calling SyncWithWallets otherwise.
+    if (Params().NetworkIDString() == "regtest") {
+        LOCK(cs);
+        nNotifiedSequence = recentlyAddedSequence;
+    }
+}
+
+bool CTxMemPool::IsFullyNotified() {
+    assert(Params().NetworkIDString() == "regtest");
+    LOCK(cs);
+    return nRecentlyAddedSequence == nNotifiedSequence;
 }
 
 CCoinsViewMemPool::CCoinsViewMemPool(CCoinsView *baseIn, CTxMemPool &mempoolIn) : CCoinsViewBacked(baseIn), mempool(mempoolIn) { }
diff --git a/src/txmempool.h b/src/txmempool.h
index caa1c9f1b..23f864adc 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -134,6 +134,9 @@ private:
     uint64_t cachedInnerUsage; //!< sum of dynamic memory usage of all the map elements (NOT the maps themselves)
 
     std::map<uint256, const CTransaction*> mapRecentlyAddedTx;
+    uint64_t nRecentlyAddedSequence = 0;
+    uint64_t nNotifiedSequence = 0;
+
     std::map<uint256, const CTransaction*> mapSproutNullifiers;
     std::map<uint256, const CTransaction*> mapSaplingNullifiers;
 
@@ -219,6 +222,7 @@ public:
     bool nullifierExists(const uint256& nullifier, ShieldedType type) const;
 
     void NotifyRecentlyAdded();
+    bool IsFullyNotified();
 
     unsigned long size()
     {