From 15fa9efd3c5dc42993ab63dd2809aaeeb5403848 Mon Sep 17 00:00:00 2001
From: David Gumberg <davidzgumberg@gmail.com>
Date: Wed, 5 Feb 2025 16:55:33 -0800
Subject: [PATCH] bench: Add wallet encryption benchmark

---
 src/bench/CMakeLists.txt     |   1 +
 src/bench/wallet_encrypt.cpp | 104 +++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 src/bench/wallet_encrypt.cpp

diff --git a/src/bench/CMakeLists.txt b/src/bench/CMakeLists.txt
index c55bbb1e05f..8c1a9582c41 100644
--- a/src/bench/CMakeLists.txt
+++ b/src/bench/CMakeLists.txt
@@ -70,6 +70,7 @@ if(ENABLE_WALLET)
       wallet_balance.cpp
       wallet_create.cpp
       wallet_create_tx.cpp
+      wallet_encrypt.cpp
       wallet_loading.cpp
       wallet_ismine.cpp
       wallet_migration.cpp
diff --git a/src/bench/wallet_encrypt.cpp b/src/bench/wallet_encrypt.cpp
new file mode 100644
index 00000000000..3850e1fd861
--- /dev/null
+++ b/src/bench/wallet_encrypt.cpp
@@ -0,0 +1,104 @@
+// Copyright (c) 2025-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 <outputtype.h>
+#include <random.h>
+#include <support/allocators/secure.h>
+#include <test/util/setup_common.h>
+#include <wallet/context.h>
+#include <wallet/test/util.h>
+#include <wallet/wallet.h>
+#include <wallet/walletutil.h>
+
+#include <cassert>
+
+namespace wallet {
+static void WalletEncrypt(benchmark::Bench& bench, bool legacy_wallet, bool measure_overhead)
+{
+    auto test_setup = MakeNoLogFileContext<TestingSetup>();
+    FastRandomContext random;
+
+    auto password{random.randbytes(20)};
+    SecureString secure_pass{password.begin(), password.end()};
+
+    WalletContext context;
+    context.args = &test_setup->m_args;
+    context.chain = test_setup->m_node.chain.get();
+    uint64_t create_flags{0};
+    if(!legacy_wallet) {
+        create_flags = WALLET_FLAG_DESCRIPTORS;
+    }
+
+    auto database{CreateMockableWalletDatabase()};
+    auto wallet{TestLoadWallet(std::move(database), context, create_flags)};
+
+    int key_count{0};
+
+    if(!legacy_wallet) {
+        // Add destinations
+        for(auto type : OUTPUT_TYPES) {
+            for(int i = 0; i < 10'000; i++) {
+                CMutableTransaction mtx;
+                mtx.vout.emplace_back(COIN, GetScriptForDestination(*Assert(wallet->GetNewDestination(type, ""))));
+                mtx.vin.emplace_back();
+                wallet->AddToWallet(MakeTransactionRef(mtx), TxStateInactive{});
+                key_count++;
+            }
+        }
+    }
+    else {
+        LegacyDataSPKM* legacy_spkm = wallet->GetOrCreateLegacyDataSPKM();
+        /* legacy spkm */
+        for(size_t i = 0; i < 10'000 * OUTPUT_TYPES.size(); i++) {
+            CKey key = GenerateRandomKey();
+            CPubKey pubkey = key.GetPubKey();
+            // Load key, scripts and create address book record
+            Assert(legacy_spkm->LoadKey(key, pubkey));
+            CTxDestination dest{PKHash(pubkey)};
+            Assert(wallet->SetAddressBook(dest, strprintf("legacy_%d", i), /*purpose=*/std::nullopt));
+
+            CMutableTransaction mtx;
+            mtx.vout.emplace_back(COIN, GetScriptForDestination(dest));
+            mtx.vin.emplace_back();
+            wallet->AddToWallet(MakeTransactionRef(mtx), TxStateInactive{});
+            key_count++;
+        }
+    }
+
+    database = DuplicateMockDatabase(wallet->GetDatabase());
+
+    // reload the wallet for the actual benchmark
+    TestUnloadWallet(std::move(wallet));
+
+    bench.batch(key_count).unit("key").run([&] {
+        wallet = TestLoadWallet(std::move(database), context, create_flags);
+
+        // Save a copy of the db before encrypting
+        database = DuplicateMockDatabase(wallet->GetDatabase());
+
+        // Skip actually encrypting wallet on the overhead measuring run, so we
+        // can subtract the overhead from the results.
+        if(!measure_overhead) {
+            wallet->EncryptWallet(secure_pass, 25000);
+        }
+
+        // cleanup
+        TestUnloadWallet(std::move(wallet));
+    });
+}
+
+static void WalletEncryptDescriptors(benchmark::Bench& bench) { WalletEncrypt(bench, /*legacy_wallet=*/false, /*measure_overhead=*/false); }
+static void WalletEncryptLegacy(benchmark::Bench& bench) { WalletEncrypt(bench, /*legacy_wallet=*/true, /*measure_overhead=*/false); }
+
+BENCHMARK(WalletEncryptDescriptors, benchmark::PriorityLevel::HIGH);
+BENCHMARK(WalletEncryptLegacy, benchmark::PriorityLevel::HIGH);
+
+static void WalletEncryptDescriptorsBenchOverhead(benchmark::Bench& bench) { WalletEncrypt(bench, /*legacy_wallet=*/false, /*measure_overhead=*/true); }
+static void WalletEncryptLegacyBenchOverhead(benchmark::Bench& bench) { WalletEncrypt(bench, /*legacy_wallet=*/true, /*measure_overhead=*/true); }
+
+BENCHMARK(WalletEncryptDescriptorsBenchOverhead, benchmark::PriorityLevel::LOW);
+BENCHMARK(WalletEncryptLegacyBenchOverhead, benchmark::PriorityLevel::LOW);
+
+} // namespace wallet