Merge bitcoin#28894: wallet: batch all individual spkms setup db writes in a single db txn

f053024 wallet: batch external signer descriptor import (Sjors Provoost)
1f65241 wallet: descriptors setup, batch db operations (furszy)
3eb769f wallet: batch legacy spkm TopUp (furszy)
075aa44 wallet: batch descriptor spkm TopUp (furszy)
bb4554c bench: add benchmark for wallet creation procedure (furszy)

Pull request description:

  Work decoupled from bitcoin#28574.

  Instead of performing multiple single write operations per spkm
  setup call, this PR batches them all within a single atomic db txn.

  Speeding up the process and preventing the wallet from entering
  an inconsistent state if any of the intermediate transactions fail
  (which shouldn't happen but.. if it does, it is better to not store
  any spkm rather than storing them partially).

  To compare the changes, added benchmark in the first commit.

ACKs for top commit:
  Sjors:
    re-utACK f053024
  achow101:
    ACK f053024
  BrandonOdiwuor:
    ACK f053024
  theStack:
    Code-review ACK f053024

Tree-SHA512: aead8548473e17d4d53e8e7039bbaf5e8bf2fe83f33b33f81cdedefe8a31b7003ceb6d5379b1bad1ca2692e909492009a21284ec8338eede078df3d19046ab5a

Author: fanquake

src/Makefile.bench.include

@@ -84,6 +84,7 @@ endif
 if ENABLE_WALLET
 bench_bench_bitcoin_SOURCES += bench/coin_selection.cpp
 bench_bench_bitcoin_SOURCES += bench/wallet_balance.cpp
+bench_bench_bitcoin_SOURCES += bench/wallet_create.cpp
 bench_bench_bitcoin_SOURCES += bench/wallet_loading.cpp
 bench_bench_bitcoin_SOURCES += bench/wallet_create_tx.cpp
 bench_bench_bitcoin_LDADD += $(BDB_LIBS) $(SQLITE_LIBS)

src/bench/wallet_create.cpp

@@ -0,0 +1,55 @@
+// Copyright (c) 2023-present The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or https://www.opensource.org/licenses/mit-license.php.
+
+#include <bench/bench.h>
+#include <node/context.h>
+#include <random.h>
+#include <test/util/setup_common.h>
+#include <wallet/context.h>
+#include <wallet/wallet.h>
+
+namespace wallet {
+static void WalletCreate(benchmark::Bench& bench, bool encrypted)
+{
+    auto test_setup = MakeNoLogFileContext<TestingSetup>();
+    FastRandomContext random;
+
+    WalletContext context;
+    context.args = &test_setup->m_args;
+    context.chain = test_setup->m_node.chain.get();
+
+    DatabaseOptions options;
+    options.require_format = DatabaseFormat::SQLITE;
+    options.require_create = true;
+    options.create_flags = WALLET_FLAG_DESCRIPTORS;
+
+    if (encrypted) {
+        options.create_passphrase = random.rand256().ToString();
+    }
+
+    DatabaseStatus status;
+    bilingual_str error_string;
+    std::vector<bilingual_str> warnings;
+
+    fs::path wallet_path = test_setup->m_path_root / strprintf("test_wallet_%d", random.rand32()).c_str();
+    bench.run([&] {
+        auto wallet = CreateWallet(context, wallet_path.u8string(), /*load_on_start=*/std::nullopt, options, status, error_string, warnings);
+        assert(status == DatabaseStatus::SUCCESS);
+        assert(wallet != nullptr);
+
+        // Cleanup
+        wallet.reset();
+        fs::remove_all(wallet_path);
+    });
+}
+
+static void WalletCreatePlain(benchmark::Bench& bench) { WalletCreate(bench, /*encrypted=*/false); }
+static void WalletCreateEncrypted(benchmark::Bench& bench) { WalletCreate(bench, /*encrypted=*/true); }
+
+#ifdef USE_SQLITE
+BENCHMARK(WalletCreatePlain, benchmark::PriorityLevel::LOW);
+BENCHMARK(WalletCreateEncrypted, benchmark::PriorityLevel::LOW);
+#endif
+
+} // namespace wallet

src/wallet/external_signer_scriptpubkeyman.cpp

@@ -16,7 +16,7 @@
 #include <vector>
 
 namespace wallet {
-bool ExternalSignerScriptPubKeyMan::SetupDescriptor(std::unique_ptr<Descriptor> desc)
+bool ExternalSignerScriptPubKeyMan::SetupDescriptor(WalletBatch& batch, std::unique_ptr<Descriptor> desc)
 {
     LOCK(cs_desc_man);
     assert(m_storage.IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS));
@@ -29,13 +29,12 @@ bool ExternalSignerScriptPubKeyMan::SetupDescriptor(std::unique_ptr<Descriptor>
     m_wallet_descriptor = w_desc;
 
     // Store the descriptor
-    WalletBatch batch(m_storage.GetDatabase());
     if (!batch.WriteDescriptor(GetID(), m_wallet_descriptor)) {
         throw std::runtime_error(std::string(__func__) + ": writing descriptor failed");
     }
 
     // TopUp
-    TopUp();
+    TopUpWithDB(batch);
 
     m_storage.UnsetBlankWalletFlag(batch);
     return true;

src/wallet/external_signer_scriptpubkeyman.h

@@ -23,7 +23,7 @@ class ExternalSignerScriptPubKeyMan : public DescriptorScriptPubKeyMan
   /** Provide a descriptor at setup time
   * Returns false if already setup or setup fails, true if setup is successful
   */
-  bool SetupDescriptor(std::unique_ptr<Descriptor>desc);
+  bool SetupDescriptor(WalletBatch& batch, std::unique_ptr<Descriptor>desc);
 
   static ExternalSigner GetExternalSigner();
 

src/wallet/scriptpubkeyman.cpp

@@ -333,7 +333,8 @@ bool LegacyScriptPubKeyMan::TopUpInactiveHDChain(const CKeyID seed_id, int64_t i
         chain.m_next_external_index = std::max(chain.m_next_external_index, index + 1);
     }
 
-    TopUpChain(chain, 0);
+    WalletBatch batch(m_storage.GetDatabase());
+    TopUpChain(batch, chain, 0);
 
     return true;
 }
@@ -1274,19 +1275,22 @@ bool LegacyScriptPubKeyMan::TopUp(unsigned int kpSize)
         return false;
     }
 
-    if (!TopUpChain(m_hd_chain, kpSize)) {
+    WalletBatch batch(m_storage.GetDatabase());
+    if (!batch.TxnBegin()) return false;
+    if (!TopUpChain(batch, m_hd_chain, kpSize)) {
         return false;
     }
     for (auto& [chain_id, chain] : m_inactive_hd_chains) {
-        if (!TopUpChain(chain, kpSize)) {
+        if (!TopUpChain(batch, chain, kpSize)) {
             return false;
         }
     }
+    if (!batch.TxnCommit()) throw std::runtime_error(strprintf("Error during keypool top up. Cannot commit changes for wallet %s", m_storage.GetDisplayName()));
     NotifyCanGetAddressesChanged();
     return true;
 }
 
-bool LegacyScriptPubKeyMan::TopUpChain(CHDChain& chain, unsigned int kpSize)
+bool LegacyScriptPubKeyMan::TopUpChain(WalletBatch& batch, CHDChain& chain, unsigned int kpSize)
 {
     LOCK(cs_KeyStore);
 
@@ -1318,7 +1322,6 @@ bool LegacyScriptPubKeyMan::TopUpChain(CHDChain& chain, unsigned int kpSize)
         missingInternal = 0;
     }
     bool internal = false;
-    WalletBatch batch(m_storage.GetDatabase());
     for (int64_t i = missingInternal + missingExternal; i--;) {
         if (i < missingInternal) {
             internal = true;
@@ -2135,6 +2138,15 @@ std::map<CKeyID, CKey> DescriptorScriptPubKeyMan::GetKeys() const
 }
 
 bool DescriptorScriptPubKeyMan::TopUp(unsigned int size)
+{
+    WalletBatch batch(m_storage.GetDatabase());
+    if (!batch.TxnBegin()) return false;
+    bool res = TopUpWithDB(batch, size);
+    if (!batch.TxnCommit()) throw std::runtime_error(strprintf("Error during descriptors keypool top up. Cannot commit changes for wallet %s", m_storage.GetDisplayName()));
+    return res;
+}
+
+bool DescriptorScriptPubKeyMan::TopUpWithDB(WalletBatch& batch, unsigned int size)
 {
     LOCK(cs_desc_man);
     unsigned int target_size;
@@ -2157,7 +2169,6 @@ bool DescriptorScriptPubKeyMan::TopUp(unsigned int size)
     FlatSigningProvider provider;
     provider.keys = GetKeys();
 
-    WalletBatch batch(m_storage.GetDatabase());
     uint256 id = GetID();
     for (int32_t i = m_max_cached_index + 1; i < new_range_end; ++i) {
         FlatSigningProvider out_keys;
@@ -2264,7 +2275,7 @@ bool DescriptorScriptPubKeyMan::AddDescriptorKeyWithDB(WalletBatch& batch, const
     }
 }
 
-bool DescriptorScriptPubKeyMan::SetupDescriptorGeneration(const CExtKey& master_key, OutputType addr_type, bool internal)
+bool DescriptorScriptPubKeyMan::SetupDescriptorGeneration(WalletBatch& batch, const CExtKey& master_key, OutputType addr_type, bool internal)
 {
     LOCK(cs_desc_man);
     assert(m_storage.IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS));
@@ -2325,7 +2336,6 @@ bool DescriptorScriptPubKeyMan::SetupDescriptorGeneration(const CExtKey& master_
     m_wallet_descriptor = w_desc;
 
     // Store the master private key, and descriptor
-    WalletBatch batch(m_storage.GetDatabase());
     if (!AddDescriptorKeyWithDB(batch, master_key.key, master_key.key.GetPubKey())) {
         throw std::runtime_error(std::string(__func__) + ": writing descriptor master private key failed");
     }
@@ -2334,7 +2344,7 @@ bool DescriptorScriptPubKeyMan::SetupDescriptorGeneration(const CExtKey& master_
     }
 
     // TopUp
-    TopUp();
+    TopUpWithDB(batch);
 
     m_storage.UnsetBlankWalletFlag(batch);
     return true;

src/wallet/scriptpubkeyman.h

@@ -366,7 +366,7 @@ class LegacyScriptPubKeyMan : public ScriptPubKeyMan, public FillableSigningProv
      */
     bool TopUpInactiveHDChain(const CKeyID seed_id, int64_t index, bool internal);
 
-    bool TopUpChain(CHDChain& chain, unsigned int size);
+    bool TopUpChain(WalletBatch& batch, CHDChain& chain, unsigned int size);
 public:
     LegacyScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) : ScriptPubKeyMan(storage), m_keypool_size(keypool_size) {}
 
@@ -583,7 +583,10 @@ class DescriptorScriptPubKeyMan : public ScriptPubKeyMan
     std::unique_ptr<FlatSigningProvider> GetSigningProvider(int32_t index, bool include_private = false) const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man);
 
 protected:
-  WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man);
+    WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man);
+
+    //! Same as 'TopUp' but designed for use within a batch transaction context
+    bool TopUpWithDB(WalletBatch& batch, unsigned int size = 0);
 
 public:
     DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size)
@@ -618,12 +621,7 @@ class DescriptorScriptPubKeyMan : public ScriptPubKeyMan
     bool IsHDEnabled() const override;
 
     //! Setup descriptors based on the given CExtkey
-    bool SetupDescriptorGeneration(const CExtKey& master_key, OutputType addr_type, bool internal);
-
-    /** Provide a descriptor at setup time
-    * Returns false if already setup or setup fails, true if setup is successful
-    */
-    bool SetupDescriptor(std::unique_ptr<Descriptor>desc);
+    bool SetupDescriptorGeneration(WalletBatch& batch, const CExtKey& master_key, OutputType addr_type, bool internal);
 
     bool HavePrivateKeys() const override;
 

src/wallet/wallet.cpp

@@ -3535,23 +3535,30 @@ void CWallet::SetupDescriptorScriptPubKeyMans(const CExtKey& master_key)
 {
     AssertLockHeld(cs_wallet);
 
+    // Create single batch txn
+    WalletBatch batch(GetDatabase());
+    if (!batch.TxnBegin()) throw std::runtime_error("Error: cannot create db transaction for descriptors setup");
+
     for (bool internal : {false, true}) {
         for (OutputType t : OUTPUT_TYPES) {
             auto spk_manager = std::unique_ptr<DescriptorScriptPubKeyMan>(new DescriptorScriptPubKeyMan(*this, m_keypool_size));
             if (IsCrypted()) {
                 if (IsLocked()) {
                     throw std::runtime_error(std::string(__func__) + ": Wallet is locked, cannot setup new descriptors");
                 }
-                if (!spk_manager->CheckDecryptionKey(vMasterKey) && !spk_manager->Encrypt(vMasterKey, nullptr)) {
+                if (!spk_manager->CheckDecryptionKey(vMasterKey) && !spk_manager->Encrypt(vMasterKey, &batch)) {
                     throw std::runtime_error(std::string(__func__) + ": Could not encrypt new descriptors");
                 }
             }
-            spk_manager->SetupDescriptorGeneration(master_key, t, internal);
+            spk_manager->SetupDescriptorGeneration(batch, master_key, t, internal);
             uint256 id = spk_manager->GetID();
             AddScriptPubKeyMan(id, std::move(spk_manager));
-            AddActiveScriptPubKeyMan(id, t, internal);
+            AddActiveScriptPubKeyManWithDb(batch, id, t, internal);
         }
     }
+
+    // Ensure information is committed to disk
+    if (!batch.TxnCommit()) throw std::runtime_error("Error: cannot commit db transaction for descriptors setup");
 }
 
 void CWallet::SetupDescriptorScriptPubKeyMans()
@@ -3578,6 +3585,10 @@ void CWallet::SetupDescriptorScriptPubKeyMans()
         UniValue signer_res = signer.GetDescriptors(account);
 
         if (!signer_res.isObject()) throw std::runtime_error(std::string(__func__) + ": Unexpected result");
+
+        WalletBatch batch(GetDatabase());
+        if (!batch.TxnBegin()) throw std::runtime_error("Error: cannot create db transaction for descriptors import");
+
         for (bool internal : {false, true}) {
             const UniValue& descriptor_vals = signer_res.find_value(internal ? "internal" : "receive");
             if (!descriptor_vals.isArray()) throw std::runtime_error(std::string(__func__) + ": Unexpected result");
@@ -3594,18 +3605,26 @@ void CWallet::SetupDescriptorScriptPubKeyMans()
                 }
                 OutputType t =  *desc->GetOutputType();
                 auto spk_manager = std::unique_ptr<ExternalSignerScriptPubKeyMan>(new ExternalSignerScriptPubKeyMan(*this, m_keypool_size));
-                spk_manager->SetupDescriptor(std::move(desc));
+                spk_manager->SetupDescriptor(batch, std::move(desc));
                 uint256 id = spk_manager->GetID();
                 AddScriptPubKeyMan(id, std::move(spk_manager));
-                AddActiveScriptPubKeyMan(id, t, internal);
+                AddActiveScriptPubKeyManWithDb(batch, id, t, internal);
             }
         }
+
+        // Ensure imported descriptors are committed to disk
+        if (!batch.TxnCommit()) throw std::runtime_error("Error: cannot commit db transaction for descriptors import");
     }
 }
 
 void CWallet::AddActiveScriptPubKeyMan(uint256 id, OutputType type, bool internal)
 {
     WalletBatch batch(GetDatabase());
+    return AddActiveScriptPubKeyManWithDb(batch, id, type, internal);
+}
+
+void CWallet::AddActiveScriptPubKeyManWithDb(WalletBatch& batch, uint256 id, OutputType type, bool internal)
+{
     if (!batch.WriteActiveScriptPubKeyMan(static_cast<uint8_t>(type), id, internal)) {
         throw std::runtime_error(std::string(__func__) + ": writing active ScriptPubKeyMan id failed");
     }

src/wallet/wallet.h

@@ -419,6 +419,9 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
     // Must be the only method adding data to it.
     void AddScriptPubKeyMan(const uint256& id, std::unique_ptr<ScriptPubKeyMan> spkm_man);
 
+    // Same as 'AddActiveScriptPubKeyMan' but designed for use within a batch transaction context
+    void AddActiveScriptPubKeyManWithDb(WalletBatch& batch, uint256 id, OutputType type, bool internal);
+
     /**
      * Catch wallet up to current chain, scanning new blocks, updating the best
      * block locator and m_last_block_processed, and registering for