4 files changed, 123 insertions, 33 deletions

diff --git a/src/wallet/db.cpp b/src/wallet/db.cpp
index a2e795056a..ae40553268 100644
--- a/src/wallet/db.cpp
+++ b/src/wallet/db.cpp
@@ -48,7 +48,7 @@ void CheckUniqueFileid(const BerkeleyEnvironment& env, const std::string& filena
 }
 
 CCriticalSection cs_db;
-std::map<std::string, BerkeleyEnvironment> g_dbenvs GUARDED_BY(cs_db); //!< Map from directory name to open db environment.
+std::map<std::string, std::weak_ptr<BerkeleyEnvironment>> g_dbenvs GUARDED_BY(cs_db); //!< Map from directory name to db environment.
 } // namespace
 
 bool WalletDatabaseFileId::operator==(const WalletDatabaseFileId& rhs) const
@@ -80,19 +80,29 @@ bool IsWalletLoaded(const fs::path& wallet_path)
     LOCK(cs_db);
     auto env = g_dbenvs.find(env_directory.string());
     if (env == g_dbenvs.end()) return false;
-    return env->second.IsDatabaseLoaded(database_filename);
+    auto database = env->second.lock();
+    return database && database->IsDatabaseLoaded(database_filename);
 }
 
-BerkeleyEnvironment* GetWalletEnv(const fs::path& wallet_path, std::string& database_filename)
+/**
+ * @param[in] wallet_path Path to wallet directory. Or (for backwards compatibility only) a path to a berkeley btree data file inside a wallet directory.
+ * @param[out] database_filename Filename of berkeley btree data file inside the wallet directory.
+ * @return A shared pointer to the BerkeleyEnvironment object for the wallet directory, never empty because ~BerkeleyEnvironment
+ * erases the weak pointer from the g_dbenvs map.
+ * @post A new BerkeleyEnvironment weak pointer is inserted into g_dbenvs if the directory path key was not already in the map.
+ */
+std::shared_ptr<BerkeleyEnvironment> GetWalletEnv(const fs::path& wallet_path, std::string& database_filename)
 {
     fs::path env_directory;
     SplitWalletPath(wallet_path, env_directory, database_filename);
     LOCK(cs_db);
-    // Note: An unused temporary BerkeleyEnvironment object may be created inside the
-    // emplace function if the key already exists. This is a little inefficient,
-    // but not a big concern since the map will be changed in the future to hold
-    // pointers instead of objects, anyway.
-    return &g_dbenvs.emplace(std::piecewise_construct, std::forward_as_tuple(env_directory.string()), std::forward_as_tuple(env_directory)).first->second;
+    auto inserted = g_dbenvs.emplace(env_directory.string(), std::weak_ptr<BerkeleyEnvironment>());
+    if (inserted.second) {
+        auto env = std::make_shared<BerkeleyEnvironment>(env_directory.string());
+        inserted.first->second = env;
+        return env;
+    }
+    return inserted.first->second.lock();
 }
 
 //
@@ -137,6 +147,7 @@ BerkeleyEnvironment::BerkeleyEnvironment(const fs::path& dir_path) : strPath(dir
 
 BerkeleyEnvironment::~BerkeleyEnvironment()
 {
+    g_dbenvs.erase(strPath);
     Close();
 }
 
@@ -214,10 +225,10 @@ bool BerkeleyEnvironment::Open(bool retry)
     return true;
 }
 
-void BerkeleyEnvironment::MakeMock()
+//! Construct an in-memory mock Berkeley environment for testing and as a place-holder for g_dbenvs emplace
+BerkeleyEnvironment::BerkeleyEnvironment()
 {
-    if (fDbEnvInit)
-        throw std::runtime_error("BerkeleyEnvironment::MakeMock: Already initialized");
+    Reset();
 
     boost::this_thread::interruption_point();
 
@@ -305,7 +316,7 @@ BerkeleyBatch::SafeDbt::operator Dbt*()
 bool BerkeleyBatch::Recover(const fs::path& file_path, void *callbackDataIn, bool (*recoverKVcallback)(void* callbackData, CDataStream ssKey, CDataStream ssValue), std::string& newFilename)
 {
     std::string filename;
-    BerkeleyEnvironment* env = GetWalletEnv(file_path, filename);
+    std::shared_ptr<BerkeleyEnvironment> env = GetWalletEnv(file_path, filename);
 
     // Recovery procedure:
     // move wallet file to walletfilename.timestamp.bak
@@ -374,7 +385,7 @@ bool BerkeleyBatch::Recover(const fs::path& file_path, void *callbackDataIn, boo
 bool BerkeleyBatch::VerifyEnvironment(const fs::path& file_path, std::string& errorStr)
 {
     std::string walletFile;
-    BerkeleyEnvironment* env = GetWalletEnv(file_path, walletFile);
+    std::shared_ptr<BerkeleyEnvironment> env = GetWalletEnv(file_path, walletFile);
     fs::path walletDir = env->Directory();
 
     LogPrintf("Using BerkeleyDB version %s\n", DbEnv::version(nullptr, nullptr, nullptr));
@@ -398,7 +409,7 @@ bool BerkeleyBatch::VerifyEnvironment(const fs::path& file_path, std::string& er
 bool BerkeleyBatch::VerifyDatabaseFile(const fs::path& file_path, std::string& warningStr, std::string& errorStr, BerkeleyEnvironment::recoverFunc_type recoverFunc)
 {
     std::string walletFile;
-    BerkeleyEnvironment* env = GetWalletEnv(file_path, walletFile);
+    std::shared_ptr<BerkeleyEnvironment> env = GetWalletEnv(file_path, walletFile);
     fs::path walletDir = env->Directory();
 
     if (fs::exists(walletDir / walletFile))
@@ -502,7 +513,7 @@ BerkeleyBatch::BerkeleyBatch(BerkeleyDatabase& database, const char* pszMode, bo
 {
     fReadOnly = (!strchr(pszMode, '+') && !strchr(pszMode, 'w'));
     fFlushOnClose = fFlushOnCloseIn;
-    env = database.env;
+    env = database.env.get();
     if (database.IsDummy()) {
         return;
     }
@@ -559,7 +570,7 @@ BerkeleyBatch::BerkeleyBatch(BerkeleyDatabase& database, const char* pszMode, bo
             // versions of BDB have an set_lk_exclusive method for this
             // purpose, but the older version we use does not.)
             for (const auto& env : g_dbenvs) {
-                CheckUniqueFileid(env.second, strFilename, *pdb_temp, this->env->m_fileids[strFilename]);
+                CheckUniqueFileid(*env.second.lock().get(), strFilename, *pdb_temp, this->env->m_fileids[strFilename]);
             }
 
             pdb = pdb_temp.release();
@@ -660,7 +671,7 @@ bool BerkeleyBatch::Rewrite(BerkeleyDatabase& database, const char* pszSkip)
     if (database.IsDummy()) {
         return true;
     }
-    BerkeleyEnvironment *env = database.env;
+    BerkeleyEnvironment *env = database.env.get();
     const std::string& strFile = database.strFile;
     while (true) {
         {
@@ -791,7 +802,7 @@ bool BerkeleyBatch::PeriodicFlush(BerkeleyDatabase& database)
         return true;
     }
     bool ret = false;
-    BerkeleyEnvironment *env = database.env;
+    BerkeleyEnvironment *env = database.env.get();
     const std::string& strFile = database.strFile;
     TRY_LOCK(cs_db, lockDb);
     if (lockDb)
diff --git a/src/wallet/db.h b/src/wallet/db.h
index 9dc373c89b..9df965305a 100644
--- a/src/wallet/db.h
+++ b/src/wallet/db.h
@@ -50,10 +50,10 @@ public:
     std::condition_variable_any m_db_in_use;
 
     BerkeleyEnvironment(const fs::path& env_directory);
+    BerkeleyEnvironment();
     ~BerkeleyEnvironment();
     void Reset();
 
-    void MakeMock();
     bool IsMock() const { return fMockDb; }
     bool IsInitialized() const { return fDbEnvInit; }
     bool IsDatabaseLoaded(const std::string& db_filename) const { return m_databases.find(db_filename) != m_databases.end(); }
@@ -102,7 +102,7 @@ public:
 bool IsWalletLoaded(const fs::path& wallet_path);
 
 /** Get BerkeleyEnvironment and database filename given a wallet path. */
-BerkeleyEnvironment* GetWalletEnv(const fs::path& wallet_path, std::string& database_filename);
+std::shared_ptr<BerkeleyEnvironment> GetWalletEnv(const fs::path& wallet_path, std::string& database_filename);
 
 /** An instance of this class represents one database.
  * For BerkeleyDB this is just a (env, strFile) tuple.
@@ -117,17 +117,11 @@ public:
     }
 
     /** Create DB handle to real database */
-    BerkeleyDatabase(const fs::path& wallet_path, bool mock = false) :
-        nUpdateCounter(0), nLastSeen(0), nLastFlushed(0), nLastWalletUpdate(0)
+    BerkeleyDatabase(std::shared_ptr<BerkeleyEnvironment> env, std::string filename) :
+        nUpdateCounter(0), nLastSeen(0), nLastFlushed(0), nLastWalletUpdate(0), env(std::move(env)), strFile(std::move(filename))
     {
-        env = GetWalletEnv(wallet_path, strFile);
-        auto inserted = env->m_databases.emplace(strFile, std::ref(*this));
+        auto inserted = this->env->m_databases.emplace(strFile, std::ref(*this));
         assert(inserted.second);
-        if (mock) {
-            env->Close();
-            env->Reset();
-            env->MakeMock();
-        }
     }
 
     ~BerkeleyDatabase() {
@@ -140,7 +134,8 @@ public:
     /** Return object for accessing database at specified path. */
     static std::unique_ptr<BerkeleyDatabase> Create(const fs::path& path)
     {
-        return MakeUnique<BerkeleyDatabase>(path);
+        std::string filename;
+        return MakeUnique<BerkeleyDatabase>(GetWalletEnv(path, filename), std::move(filename));
     }
 
     /** Return object for accessing dummy database with no read/write capabilities. */
@@ -152,7 +147,7 @@ public:
     /** Return object for accessing temporary in-memory database. */
     static std::unique_ptr<BerkeleyDatabase> CreateMock()
     {
-        return MakeUnique<BerkeleyDatabase>("", true /* mock */);
+        return MakeUnique<BerkeleyDatabase>(std::make_shared<BerkeleyEnvironment>(), "");
     }
 
     /** Rewrite the entire database on disk, with the exception of key pszSkip if non-zero
@@ -176,12 +171,21 @@ public:
     unsigned int nLastFlushed;
     int64_t nLastWalletUpdate;
 
+    /**
+     * Pointer to shared database environment.
+     *
+     * Normally there is only one BerkeleyDatabase object per
+     * BerkeleyEnvivonment, but in the special, backwards compatible case where
+     * multiple wallet BDB data files are loaded from the same directory, this
+     * will point to a shared instance that gets freed when the last data file
+     * is closed.
+     */
+    std::shared_ptr<BerkeleyEnvironment> env;
+
     /** Database pointer. This is initialized lazily and reset during flushes, so it can be null. */
     std::unique_ptr<Db> m_db;
 
 private:
-    /** BerkeleyDB specific */
-    BerkeleyEnvironment *env;
     std::string strFile;
 
     /** Return whether this database handle is a dummy for testing.
diff --git a/src/wallet/test/db_tests.cpp b/src/wallet/test/db_tests.cpp
new file mode 100644
index 0000000000..2a64749379
--- /dev/null
+++ b/src/wallet/test/db_tests.cpp
@@ -0,0 +1,72 @@
+// Copyright (c) 2018 The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#include <memory>
+
+#include <boost/test/unit_test.hpp>
+
+#include <fs.h>
+#include <test/test_bitcoin.h>
+#include <wallet/db.h>
+
+
+BOOST_FIXTURE_TEST_SUITE(db_tests, BasicTestingSetup)
+
+BOOST_AUTO_TEST_CASE(getwalletenv_file)
+{
+    std::string test_name = "test_name.dat";
+    fs::path datadir = SetDataDir("tempdir");
+    fs::path file_path = datadir / test_name;
+    std::ofstream f(file_path.BOOST_FILESYSTEM_C_STR);
+    f.close();
+
+    std::string filename;
+    std::shared_ptr<BerkeleyEnvironment> env = GetWalletEnv(file_path, filename);
+    BOOST_CHECK(filename == test_name);
+    BOOST_CHECK(env->Directory() == datadir);
+}
+
+BOOST_AUTO_TEST_CASE(getwalletenv_directory)
+{
+    std::string expected_name = "wallet.dat";
+    fs::path datadir = SetDataDir("tempdir");
+
+    std::string filename;
+    std::shared_ptr<BerkeleyEnvironment> env = GetWalletEnv(datadir, filename);
+    BOOST_CHECK(filename == expected_name);
+    BOOST_CHECK(env->Directory() == datadir);
+}
+
+BOOST_AUTO_TEST_CASE(getwalletenv_g_dbenvs_multiple)
+{
+    fs::path datadir = SetDataDir("tempdir");
+    fs::path datadir_2 = SetDataDir("tempdir_2");
+    std::string filename;
+
+    std::shared_ptr<BerkeleyEnvironment> env_1 = GetWalletEnv(datadir, filename);
+    std::shared_ptr<BerkeleyEnvironment> env_2 = GetWalletEnv(datadir, filename);
+    std::shared_ptr<BerkeleyEnvironment> env_3 = GetWalletEnv(datadir_2, filename);
+
+    BOOST_CHECK(env_1 == env_2);
+    BOOST_CHECK(env_2 != env_3);
+}
+
+BOOST_AUTO_TEST_CASE(getwalletenv_g_dbenvs_free_instance)
+{
+    fs::path datadir = SetDataDir("tempdir");
+    fs::path datadir_2 = SetDataDir("tempdir_2");
+    std::string filename;
+
+    std::shared_ptr <BerkeleyEnvironment> env_1_a = GetWalletEnv(datadir, filename);
+    std::shared_ptr <BerkeleyEnvironment> env_2_a = GetWalletEnv(datadir_2, filename);
+    env_1_a.reset();
+
+    std::shared_ptr<BerkeleyEnvironment> env_1_b = GetWalletEnv(datadir, filename);
+    std::shared_ptr<BerkeleyEnvironment> env_2_b = GetWalletEnv(datadir_2, filename);
+
+    BOOST_CHECK(env_1_a != env_1_b);
+    BOOST_CHECK(env_2_a == env_2_b);
+}
+
+BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 879ff8fb54..5c379aacd6 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -3922,6 +3922,9 @@ bool CWallet::Verify(interfaces::Chain& chain, const WalletLocation& location, b
         return false;
     }
 
+    // Keep same database environment instance across Verify/Recover calls below.
+    std::unique_ptr<WalletDatabase> database = WalletDatabase::Create(wallet_path);
+
     try {
         if (!WalletBatch::VerifyEnvironment(wallet_path, error_string)) {
             return false;