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

f05302427386fe63f4929a7198652cb1e4ab3bcc wallet: batch external signer descriptor import (Sjors Provoost)
1f65241b733cd1e962c88909ae66816bc6451fd1 wallet: descriptors setup, batch db operations (furszy)
3eb769f15013873755e482707cad341bc1ce8a8c wallet: batch legacy spkm TopUp (furszy)
075aa44ceba41fa82bb3ce2295e2962e5fd0508e wallet: batch descriptor spkm TopUp (furszy)
bb4554c81e0d819d74996f89cbb9c00476aedf8c bench: add benchmark for wallet creation procedure (furszy)

Pull request description:

  Work decoupled from #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 f05302427386fe63f4929a7198652cb1e4ab3bcc
  achow101:
    ACK f05302427386fe63f4929a7198652cb1e4ab3bcc
  BrandonOdiwuor:
    ACK f05302427386fe63f4929a7198652cb1e4ab3bcc
  theStack:
    Code-review ACK f05302427386fe63f4929a7198652cb1e4ab3bcc

Tree-SHA512: aead8548473e17d4d53e8e7039bbaf5e8bf2fe83f33b33f81cdedefe8a31b7003ceb6d5379b1bad1ca2692e909492009a21284ec8338eede078df3d19046ab5a

8 files changed, 111 insertions, 26 deletions

diff --git a/src/Makefile.bench.include b/src/Makefile.bench.include
index 28b779a5a8..217f123b16 100644
--- a/src/Makefile.bench.include
+++ b/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)
diff --git a/src/bench/wallet_create.cpp b/src/bench/wallet_create.cpp
new file mode 100644
index 0000000000..ba3c25d25e
--- /dev/null
+++ b/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
diff --git a/src/wallet/external_signer_scriptpubkeyman.cpp b/src/wallet/external_signer_scriptpubkeyman.cpp
index 6d026d02c1..a71f8f9fbc 100644
--- a/src/wallet/external_signer_scriptpubkeyman.cpp
+++ b/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;
diff --git a/src/wallet/external_signer_scriptpubkeyman.h b/src/wallet/external_signer_scriptpubkeyman.h
index 01dc80b1ca..c052ce6129 100644
--- a/src/wallet/external_signer_scriptpubkeyman.h
+++ b/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();
 
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index d2b2801aa8..0b4800b848 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/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;
@@ -2136,6 +2139,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;
     if (size > 0) {
@@ -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;
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 7c0eca1475..dccbf3ced6 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -366,7 +366,7 @@ private:
      */
     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 @@ private:
     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 @@ public:
     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;
 
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index c34bf96b5e..4d4e23dd4c 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -3535,6 +3535,10 @@ 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));
@@ -3542,16 +3546,19 @@ void CWallet::SetupDescriptorScriptPubKeyMans(const CExtKey& master_key)
                 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");
     }
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index bc45010200..11b7964316 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -419,6 +419,9 @@ private:
     // 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