From 0cd9ad208c327127cc4616ccdc37557fad3cf381 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Barbosa?= <joao.paulo.barbosa@gmail.com>
Date: Wed, 12 Dec 2018 23:21:19 +0000
Subject: rpc: Make unloadwallet wait for complete wallet unload

Github-Pull: #14941
Rebased-From: c37851d
---
 src/init.cpp             |  2 +-
 src/wallet/init.cpp      |  8 ++++++--
 src/wallet/rpcwallet.cpp | 10 +---------
 src/wallet/wallet.cpp    | 39 +++++++++++++++++++++++++++++++++++++++
 src/wallet/wallet.h      |  9 +++++++++
 5 files changed, 56 insertions(+), 12 deletions(-)

(limited to 'src')

diff --git a/src/init.cpp b/src/init.cpp
index c16009a0ba..075dc42099 100644
--- a/src/init.cpp
+++ b/src/init.cpp
@@ -282,10 +282,10 @@ void Shutdown()
         LogPrintf("%s: Unable to remove pidfile: %s\n", __func__, e.what());
     }
 #endif
+    g_wallet_init_interface.Close();
     UnregisterAllValidationInterfaces();
     GetMainSignals().UnregisterBackgroundSignalScheduler();
     GetMainSignals().UnregisterWithMempoolSignals(mempool);
-    g_wallet_init_interface.Close();
     globalVerifyHandle.reset();
     ECC_Stop();
     LogPrintf("%s: done\n", __func__);
diff --git a/src/wallet/init.cpp b/src/wallet/init.cpp
index b9f267210e..7dc6c30645 100644
--- a/src/wallet/init.cpp
+++ b/src/wallet/init.cpp
@@ -264,7 +264,11 @@ void WalletInit::Stop() const
 
 void WalletInit::Close() const
 {
-    for (const std::shared_ptr<CWallet>& pwallet : GetWallets()) {
-        RemoveWallet(pwallet);
+    auto wallets = GetWallets();
+    while (!wallets.empty()) {
+        auto wallet = wallets.back();
+        wallets.pop_back();
+        RemoveWallet(wallet);
+        UnloadWallet(std::move(wallet));
     }
 }
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
index e02171ec40..cc429f5e56 100644
--- a/src/wallet/rpcwallet.cpp
+++ b/src/wallet/rpcwallet.cpp
@@ -3239,16 +3239,8 @@ static UniValue unloadwallet(const JSONRPCRequest& request)
     if (!RemoveWallet(wallet)) {
         throw JSONRPCError(RPC_MISC_ERROR, "Requested wallet already unloaded");
     }
-    UnregisterValidationInterface(wallet.get());
 
-    // The wallet can be in use so it's not possible to explicitly unload here.
-    // Just notify the unload intent so that all shared pointers are released.
-    // The wallet will be destroyed once the last shared pointer is released.
-    wallet->NotifyUnload();
-
-    // There's no point in waiting for the wallet to unload.
-    // At this point this method should never fail. The unloading could only
-    // fail due to an unexpected error which would cause a process termination.
+    UnloadWallet(std::move(wallet));
 
     return NullUniValue;
 }
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 060325921f..e86aa04d18 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -81,13 +81,52 @@ std::shared_ptr<CWallet> GetWallet(const std::string& name)
     return nullptr;
 }
 
+static std::mutex g_wallet_release_mutex;
+static std::condition_variable g_wallet_release_cv;
+static std::set<CWallet*> g_unloading_wallet_set;
+
 // Custom deleter for shared_ptr<CWallet>.
 static void ReleaseWallet(CWallet* wallet)
 {
+    // Unregister and delete the wallet right after BlockUntilSyncedToCurrentChain
+    // so that it's in sync with the current chainstate.
     wallet->WalletLogPrintf("Releasing wallet\n");
     wallet->BlockUntilSyncedToCurrentChain();
     wallet->Flush();
+    UnregisterValidationInterface(wallet);
     delete wallet;
+    // Wallet is now released, notify UnloadWallet, if any.
+    {
+        std::lock_guard<std::mutex> lock(g_wallet_release_mutex);
+        if (g_unloading_wallet_set.erase(wallet) == 0) {
+            // UnloadWallet was not called for this wallet, all done.
+            return;
+        }
+    }
+    g_wallet_release_cv.notify_all();
+}
+
+void UnloadWallet(std::shared_ptr<CWallet>&& wallet)
+{
+    // Mark wallet for unloading.
+    CWallet* pwallet = wallet.get();
+    {
+        std::lock_guard<std::mutex> lock(g_wallet_release_mutex);
+        auto it = g_unloading_wallet_set.insert(pwallet);
+        assert(it.second);
+    }
+    // The wallet can be in use so it's not possible to explicitly unload here.
+    // Notify the unload intent so that all remaining shared pointers are
+    // released.
+    pwallet->NotifyUnload();
+    // Time to ditch our shared_ptr and wait for ReleaseWallet call.
+    wallet.reset();
+    {
+        std::unique_lock<std::mutex> lock(g_wallet_release_mutex);
+        while (g_unloading_wallet_set.count(pwallet) == 1) {
+            g_wallet_release_cv.wait(lock);
+        }
+    }
 }
 
 const uint32_t BIP32_HARDENED_KEY_LIMIT = 0x80000000;
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 7b791ae416..1f20834843 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -33,6 +33,13 @@
 #include <utility>
 #include <vector>
 
+//! Explicitly unload and delete the wallet.
+//  Blocks the current thread after signaling the unload intent so that all
+//  wallet clients release the wallet.
+//  Note that, when blocking is not required, the wallet is implicitly unloaded
+//  by the shared pointer deleter.
+void UnloadWallet(std::shared_ptr<CWallet>&& wallet);
+
 bool AddWallet(const std::shared_ptr<CWallet>& wallet);
 bool RemoveWallet(const std::shared_ptr<CWallet>& wallet);
 bool HasWallets();
@@ -817,6 +824,8 @@ public:
 
     ~CWallet()
     {
+        // Should not have slots connected at this point.
+        assert(NotifyUnload.empty());
         delete encrypted_batch;
         encrypted_batch = nullptr;
     }
-- 
cgit v1.2.3