1 files changed, 39 insertions, 8 deletions

diff --git a/src/bench/wallet_loading.cpp b/src/bench/wallet_loading.cpp
index 8fbd814d29..0bbafee76e 100644
--- a/src/bench/wallet_loading.cpp
+++ b/src/bench/wallet_loading.cpp
@@ -17,21 +17,19 @@
 #include <optional>
 
 using wallet::CWallet;
+using wallet::DatabaseFormat;
 using wallet::DatabaseOptions;
-using wallet::DatabaseStatus;
 using wallet::ISMINE_SPENDABLE;
 using wallet::MakeWalletDatabase;
 using wallet::TxStateInactive;
 using wallet::WALLET_FLAG_DESCRIPTORS;
 using wallet::WalletContext;
+using wallet::WalletDatabase;
 
-static const std::shared_ptr<CWallet> BenchLoadWallet(WalletContext& context, DatabaseOptions& options)
+static const std::shared_ptr<CWallet> BenchLoadWallet(std::unique_ptr<WalletDatabase> database, WalletContext& context, DatabaseOptions& options)
 {
-    DatabaseStatus status;
     bilingual_str error;
     std::vector<bilingual_str> warnings;
-    auto database = MakeWalletDatabase("", options, status, error);
-    assert(database);
     auto wallet = CWallet::Create(context, "", std::move(database), options.create_flags, error, warnings);
     NotifyWalletLoaded(context, wallet);
     if (context.chain) {
@@ -60,6 +58,30 @@ static void AddTx(CWallet& wallet)
     wallet.AddToWallet(MakeTransactionRef(mtx), TxStateInactive{});
 }
 
+static std::unique_ptr<WalletDatabase> DuplicateMockDatabase(WalletDatabase& database, DatabaseOptions& options)
+{
+    auto new_database = CreateMockWalletDatabase(options);
+
+    // Get a cursor to the original database
+    auto batch = database.MakeBatch();
+    batch->StartCursor();
+
+    // Get a batch for the new database
+    auto new_batch = new_database->MakeBatch();
+
+    // Read all records from the original database and write them to the new one
+    while (true) {
+        CDataStream key(SER_DISK, CLIENT_VERSION);
+        CDataStream value(SER_DISK, CLIENT_VERSION);
+        bool complete;
+        batch->ReadAtCursor(key, value, complete);
+        if (complete) break;
+        new_batch->Write(key, value);
+    }
+
+    return new_database;
+}
+
 static void WalletLoading(benchmark::Bench& bench, bool legacy_wallet)
 {
     const auto test_setup = MakeNoLogFileContext<TestingSetup>();
@@ -72,21 +94,30 @@ static void WalletLoading(benchmark::Bench& bench, bool legacy_wallet)
     // Setup the wallet
     // Loading the wallet will also create it
     DatabaseOptions options;
-    if (!legacy_wallet) options.create_flags = WALLET_FLAG_DESCRIPTORS;
-    auto wallet = BenchLoadWallet(context, options);
+    if (legacy_wallet) {
+        options.require_format = DatabaseFormat::BERKELEY;
+    } else {
+        options.create_flags = WALLET_FLAG_DESCRIPTORS;
+        options.require_format = DatabaseFormat::SQLITE;
+    }
+    auto database = CreateMockWalletDatabase(options);
+    auto wallet = BenchLoadWallet(std::move(database), context, options);
 
     // Generate a bunch of transactions and addresses to put into the wallet
     for (int i = 0; i < 1000; ++i) {
         AddTx(*wallet);
     }
 
+    database = DuplicateMockDatabase(wallet->GetDatabase(), options);
+
     // reload the wallet for the actual benchmark
     BenchUnloadWallet(std::move(wallet));
 
     bench.epochs(5).run([&] {
-        wallet = BenchLoadWallet(context, options);
+        wallet = BenchLoadWallet(std::move(database), context, options);
 
         // Cleanup
+        database = DuplicateMockDatabase(wallet->GetDatabase(), options);
         BenchUnloadWallet(std::move(wallet));
     });
 }