From 51ac1abf6fdb6340110397a959681d52e29f73f7 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 | 82 ++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 src/bench/wallet_encrypt.cpp

diff --git a/src/bench/CMakeLists.txt b/src/bench/CMakeLists.txt
index 50b29a14667..e1d5e4097ff 100644
--- a/src/bench/CMakeLists.txt
+++ b/src/bench/CMakeLists.txt
@@ -77,6 +77,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..3d73081c7cc
--- /dev/null
+++ b/src/bench/wallet_encrypt.cpp
@@ -0,0 +1,82 @@
+// 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 <key_io.h>
+#include <outputtype.h>
+#include <random.h>
+#include <support/allocators/secure.h>
+#include <test/util/setup_common.h>
+#include <util/time.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, unsigned int key_count)
+{
+    auto test_setup = MakeNoLogFileContext<TestingSetup>();
+    FastRandomContext rand{/*fDeterministic=*/true};
+
+    auto password{rand.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 = WALLET_FLAG_DESCRIPTORS;
+    auto database = CreateMockableWalletDatabase();
+    auto wallet = TestCreateWallet(std::move(database), context, create_flags);
+
+    {
+        LOCK(wallet->cs_wallet);
+        for (unsigned int i = 0; i < key_count; i++) {
+            CKey key = GenerateRandomKey();
+            FlatSigningProvider keys;
+            std::string error;
+            std::vector<std::unique_ptr<Descriptor>> desc = Parse("combo(" + EncodeSecret(key) + ")", keys, error, /*require_checksum=*/false);
+            WalletDescriptor w_desc(std::move(desc.at(0)), /*creation_time=*/0, /*range_start=*/0, /*range_end=*/0, /*next_index=*/0);
+            Assert(wallet->AddWalletDescriptor(w_desc, keys, /*label=*/"", /*internal=*/false));
+        }
+    }
+
+    database = DuplicateMockDatabase(wallet->GetDatabase());
+
+    // reload the wallet for the actual benchmark
+    TestUnloadWallet(std::move(wallet));
+
+    // Setting a mock time is necessary to force default derive iteration count during
+    // wallet encryption.
+    SetMockTime(1);
+
+    // This benchmark has a lot of overhead, this should be good enough to catch
+    // any regressions, but for an accurate measurement of how long wallet
+    // encryption takes, this should be reworked after something like
+    // https://github.com/bitcoin/bitcoin/pull/34208 is merged.
+    bench.batch(key_count).unit("key").run([&] {
+        wallet = TestLoadWallet(std::move(database), context);
+
+        // Save a copy of the db before encrypting
+        database = DuplicateMockDatabase(wallet->GetDatabase());
+
+        wallet->EncryptWallet(secure_pass);
+
+        for (const auto& [_, key] : wallet->mapMasterKeys){
+            assert(key.nDeriveIterations == CMasterKey::DEFAULT_DERIVE_ITERATIONS);
+        }
+
+        TestUnloadWallet(std::move(wallet));
+    });
+}
+
+constexpr unsigned int KEY_COUNT = 2000;
+
+static void WalletEncryptDescriptors(benchmark::Bench& bench) { WalletEncrypt(bench, /*key_count=*/KEY_COUNT); }
+BENCHMARK(WalletEncryptDescriptors);
+
+} // namespace wallet