Add RAII wallet rescan reserver

3 files changed, 55 insertions, 1 deletions

diff --git a/src/wallet/rpcdump.cpp b/src/wallet/rpcdump.cpp
index 1711b400e3..a0b57ecd22 100644
--- a/src/wallet/rpcdump.cpp
+++ b/src/wallet/rpcdump.cpp
@@ -262,6 +262,11 @@ UniValue importaddress(const JSONRPCRequest& request)
     if (fRescan && fPruneMode)
         throw JSONRPCError(RPC_WALLET_ERROR, "Rescan is disabled in pruned mode");
 
+    WalletRescanReserver reserver(pwallet);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     // Whether to import a p2sh version, too
     bool fP2SH = false;
     if (!request.params[3].isNull())
@@ -429,6 +434,11 @@ UniValue importpubkey(const JSONRPCRequest& request)
     if (fRescan && fPruneMode)
         throw JSONRPCError(RPC_WALLET_ERROR, "Rescan is disabled in pruned mode");
 
+    WalletRescanReserver reserver(pwallet);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     if (!IsHex(request.params[0].get_str()))
         throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey must be a hex string");
     std::vector<unsigned char> data(ParseHex(request.params[0].get_str()));
@@ -480,6 +490,11 @@ UniValue importwallet(const JSONRPCRequest& request)
     if (fPruneMode)
         throw JSONRPCError(RPC_WALLET_ERROR, "Importing wallets is disabled in pruned mode");
 
+    WalletRescanReserver reserver(pwallet);
+    if (!reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     int64_t nTimeBegin = 0;
     bool fGood = true;
     {
@@ -1138,6 +1153,11 @@ UniValue importmulti(const JSONRPCRequest& mainRequest)
         }
     }
 
+    WalletRescanReserver reserver(pwallet);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     int64_t now = 0;
     bool fRunScan = false;
     int64_t nLowestTimestamp = 0;
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
index bc0d544aab..50642e6b8a 100644
--- a/src/wallet/rpcwallet.cpp
+++ b/src/wallet/rpcwallet.cpp
@@ -3398,7 +3398,8 @@ UniValue rescanblockchain(const JSONRPCRequest& request)
             );
     }
 
-    if (pwallet->IsScanning()) {
+    WalletRescanReserver reserver(pwallet);
+    if (!reserver.reserve()) {
         throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
     }
 
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 53a2c6b9d5..5d782faa60 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -669,6 +669,9 @@ private:
     static std::atomic<bool> fFlushScheduled;
     std::atomic<bool> fAbortRescan;
     std::atomic<bool> fScanningWallet;
+    std::mutex mutexScanning;
+    friend class WalletRescanReserver;
+
 
     /**
      * Select a set of coins such that nValueRet >= nTargetValue and at least
@@ -1263,4 +1266,34 @@ CTxDestination GetDestinationForKey(const CPubKey& key, OutputType);
 /** Get all destinations (potentially) supported by the wallet for the given key. */
 std::vector<CTxDestination> GetAllDestinationsForKey(const CPubKey& key);
 
+/** RAII object to check and reserve a wallet rescan */
+class WalletRescanReserver
+{
+private:
+    CWalletRef m_wallet;
+    bool m_could_reserve;
+public:
+    explicit WalletRescanReserver(CWalletRef w) : m_wallet(w), m_could_reserve(false) {}
+
+    bool reserve()
+    {
+        assert(!m_could_reserve);
+        std::lock_guard<std::mutex> lock(m_wallet->mutexScanning);
+        if (m_wallet->fScanningWallet) {
+            return false;
+        }
+        m_wallet->fScanningWallet = true;
+        m_could_reserve = true;
+        return true;
+    }
+
+    ~WalletRescanReserver()
+    {
+        std::lock_guard<std::mutex> lock(m_wallet->mutexScanning);
+        if (m_could_reserve) {
+            m_wallet->fScanningWallet = false;
+        }
+    }
+};
+
 #endif // BITCOIN_WALLET_WALLET_H