From 7c9076a2d2e5dd117c31ca871e81c9170aa0e371 Mon Sep 17 00:00:00 2001
From: furszy <matiasfurszyfer@protonmail.com>
Date: Tue, 13 Feb 2024 12:34:16 -0300
Subject: [PATCH] wallet: migration, consolidate main wallet db writes

Perform a single db write operation for the entire
migration procedure.
---
 src/wallet/wallet.cpp | 29 +++++++++++++----------------
 src/wallet/wallet.h   |  2 +-
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index e00537e9204..9c3b05b444d 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -4068,7 +4068,7 @@ std::optional<MigrationData> CWallet::GetDescriptorsForLegacy(bilingual_str& err
     return res;
 }
 
-util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
+util::Result<void> CWallet::ApplyMigrationData(WalletBatch& local_wallet_batch, MigrationData& data)
 {
     AssertLockHeld(cs_wallet);
 
@@ -4096,7 +4096,7 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
     }
 
     // Remove the LegacyScriptPubKeyMan from disk
-    if (!legacy_spkm->DeleteRecords()) {
+    if (!legacy_spkm->DeleteRecordsWithDB(local_wallet_batch)) {
         return util::Error{_("Error: cannot remove legacy wallet records")};
     }
 
@@ -4106,9 +4106,8 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
     m_internal_spk_managers.clear();
 
     // Setup new descriptors
-    SetWalletFlag(WALLET_FLAG_DESCRIPTORS);
+    SetWalletFlagWithDB(local_wallet_batch, WALLET_FLAG_DESCRIPTORS);
     if (!IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS)) {
-        WalletBatch local_wallet_batch(GetDatabase());
         // Use the existing master key if we have it
         if (data.master_key.key.IsValid()) {
             SetupDescriptorScriptPubKeyMans(local_wallet_batch, data.master_key);
@@ -4120,7 +4119,7 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
 
     // Get best block locator so that we can copy it to the watchonly and solvables
     CBlockLocator best_block_locator;
-    if (!WalletBatch(GetDatabase()).ReadBestBlock(best_block_locator)) {
+    if (!local_wallet_batch.ReadBestBlock(best_block_locator)) {
         return util::Error{_("Error: Unable to read wallet's best block locator record")};
     }
 
@@ -4179,7 +4178,7 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
     watchonly_batch.reset(); // Flush
     // Do the removes
     if (txids_to_delete.size() > 0) {
-        if (auto res = RemoveTxs(txids_to_delete); !res) {
+        if (auto res = RemoveTxs(local_wallet_batch, txids_to_delete); !res) {
             return util::Error{_("Error: Could not delete watchonly transactions. ") + util::ErrorString(res)};
         }
     }
@@ -4253,8 +4252,6 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
     }
 
     // Remove the things to delete in this wallet
-    WalletBatch local_wallet_batch(GetDatabase());
-    local_wallet_batch.TxnBegin();
     if (dests_to_delete.size() > 0) {
         for (const auto& dest : dests_to_delete) {
             if (!DelAddressBookWithDB(local_wallet_batch, dest)) {
@@ -4262,9 +4259,6 @@ util::Result<void> CWallet::ApplyMigrationData(MigrationData& data)
             }
         }
     }
-    local_wallet_batch.TxnCommit();
-
-    WalletLogPrintf("Wallet migration complete.\n");
 
     return {}; // all good
 }
@@ -4377,11 +4371,14 @@ bool DoMigration(CWallet& wallet, WalletContext& context, bilingual_str& error,
     }
 
     // Add the descriptors to wallet, remove LegacyScriptPubKeyMan, and cleanup txs and address book data
-    if (auto res_migration = wallet.ApplyMigrationData(*data); !res_migration) {
-        error = util::ErrorString(res_migration);
-        return false;
-    }
-    return true;
+    return RunWithinTxn(wallet.GetDatabase(), /*process_desc=*/"apply migration process", [&](WalletBatch& batch) EXCLUSIVE_LOCKS_REQUIRED(wallet.cs_wallet){
+        if (auto res_migration = wallet.ApplyMigrationData(batch, *data); !res_migration) {
+            error = util::ErrorString(res_migration);
+            return false;
+        }
+        wallet.WalletLogPrintf("Wallet migration complete.\n");
+        return true;
+    });
 }
 
 util::Result<MigrationResult> MigrateLegacyToDescriptor(const std::string& wallet_name, const SecureString& passphrase, WalletContext& context)
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 1d350fb210b..d869f031bbf 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -1051,7 +1051,7 @@ public:
 
     //! Adds the ScriptPubKeyMans given in MigrationData to this wallet, removes LegacyScriptPubKeyMan,
     //! and where needed, moves tx and address book entries to watchonly_wallet or solvable_wallet
-    util::Result<void> ApplyMigrationData(MigrationData& data) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    util::Result<void> ApplyMigrationData(WalletBatch& local_wallet_batch, MigrationData& data) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     //! Whether the (external) signer performs R-value signature grinding
     bool CanGrindR() const;