Merge bitcoin/bitcoin#26008: wallet: cache IsMine scriptPubKeys to improve performance of descriptor wallets

e041ed9b755468d205ed48b29f6c4b9e9df8bc9f wallet: Retrieve ID from loaded DescSPKM directly (Ava Chow)
39640dd34e980e69d13664ddbc2a7612a1888ab4 wallet: Use scriptPubKeyCache in GetSolvingProvider (Ava Chow)
b410f68791143800968f4c628beda1c9f898b4f6 wallet: Use scriptPubKey cache in GetScriptPubKeyMans (Ava Chow)
edf4e73a163739a64eb9a54cd95413583a0e5c1f wallet: Use scriptPubKey cache in IsMine (Ava Chow)
37232332bd7253422ea92a8c9eeb36b12fc84b56 wallet: Cache scriptPubKeys for all DescriptorSPKMs (Ava Chow)
99a0cddbc04e8bfea3748a6cb4c0107392fab94f wallet: Introduce a callback called after TopUp completes (Ava Chow)
b27682593266c99507c720855cb34f5f7d363dd2 bench: Add a benchmark for ismine (Ava Chow)

Pull request description:

  Wallets that have a ton of non-ranged descriptors (such as a migrated non-HD wallet) perform fairly poorly due to looping through all of the wallet's `ScriptPubKeyMan`s. This is done in various places, such as `IsMine`, and helper functions for fetching a `ScriptPubKeyMan` and a `SolvingProvider`. This also has a bit of a performance impact on standard descriptor wallets, although less noticeable due to the small number of SPKMs.

  As these functions are based on doing `IsMine` for each `ScriptPubKeyMan`, we can improve this performance by caching `IsMine` scriptPubKeys for all descriptors and use that to determine which `ScriptPubKeyMan` to actually use for those things. This cache is used exclusively and we no longer iterate the SPKMs.

  Also added a benchmark for `IsMine`.

ACKs for top commit:
  ryanofsky:
    Code review ACK e041ed9b755468d205ed48b29f6c4b9e9df8bc9f. Just suggested changes since last review
  josibake:
    ACK https://github.com/bitcoin/bitcoin/pull/26008/commits/e041ed9b755468d205ed48b29f6c4b9e9df8bc9f
  furszy:
    Code review ACK e041ed9b

Tree-SHA512: 8e7081991a025e682e9dea838b4543b0d179832d1c47397fb9fe7a97fa01eb699c15a5d5a785634926844fc83a46e6ac07ef753119f39d84423220ef8a548894

7 files changed, 157 insertions, 21 deletions

diff --git a/src/Makefile.bench.include b/src/Makefile.bench.include
index 9e5366f0b4..b24405ce19 100644
--- a/src/Makefile.bench.include
+++ b/src/Makefile.bench.include
@@ -88,6 +88,8 @@ 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_SOURCES += bench/wallet_ismine.cpp
+
 bench_bench_bitcoin_LDADD += $(BDB_LIBS) $(SQLITE_LIBS)
 endif
 
diff --git a/src/bench/wallet_ismine.cpp b/src/bench/wallet_ismine.cpp
new file mode 100644
index 0000000000..261c95c7c8
--- /dev/null
+++ b/src/bench/wallet_ismine.cpp
@@ -0,0 +1,74 @@
+// Copyright (c) 2022 The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#include <bench/bench.h>
+#include <interfaces/chain.h>
+#include <key.h>
+#include <key_io.h>
+#include <node/context.h>
+#include <test/util/setup_common.h>
+#include <util/translation.h>
+#include <validationinterface.h>
+#include <wallet/context.h>
+#include <wallet/wallet.h>
+#include <wallet/walletutil.h>
+#include <wallet/test/util.h>
+
+namespace wallet {
+static void WalletIsMine(benchmark::Bench& bench, bool legacy_wallet, int num_combo = 0)
+{
+    const auto test_setup = MakeNoLogFileContext<TestingSetup>();
+
+    WalletContext context;
+    context.args = &test_setup->m_args;
+    context.chain = test_setup->m_node.chain.get();
+
+    // Setup the wallet
+    // Loading the wallet will also create it
+    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);
+
+    // For a descriptor wallet, fill with num_combo combo descriptors with random keys
+    // This benchmarks a non-HD wallet migrated to descriptors
+    if (!legacy_wallet && num_combo > 0) {
+        LOCK(wallet->cs_wallet);
+        for (int i = 0; i < num_combo; ++i) {
+            CKey key;
+            key.MakeNewKey(/*fCompressed=*/true);
+            FlatSigningProvider keys;
+            std::string error;
+            std::unique_ptr<Descriptor> desc = Parse("combo(" + EncodeSecret(key) + ")", keys, error, /*require_checksum=*/false);
+            WalletDescriptor w_desc(std::move(desc), /*creation_time=*/0, /*range_start=*/0, /*range_end=*/0, /*next_index=*/0);
+            auto spkm = wallet->AddWalletDescriptor(w_desc, keys, /*label=*/"", /*internal=*/false);
+            assert(spkm);
+        }
+    }
+
+    const CScript script = GetScriptForDestination(DecodeDestination(ADDRESS_BCRT1_UNSPENDABLE));
+
+    bench.run([&] {
+        LOCK(wallet->cs_wallet);
+        isminetype mine = wallet->IsMine(script);
+        assert(mine == ISMINE_NO);
+    });
+
+    TestUnloadWallet(std::move(wallet));
+}
+
+#ifdef USE_BDB
+static void WalletIsMineLegacy(benchmark::Bench& bench) { WalletIsMine(bench, /*legacy_wallet=*/true); }
+BENCHMARK(WalletIsMineLegacy, benchmark::PriorityLevel::LOW);
+#endif
+
+#ifdef USE_SQLITE
+static void WalletIsMineDescriptors(benchmark::Bench& bench) { WalletIsMine(bench, /*legacy_wallet=*/false); }
+static void WalletIsMineMigratedDescriptors(benchmark::Bench& bench) { WalletIsMine(bench, /*legacy_wallet=*/false, /*num_combo=*/2000); }
+BENCHMARK(WalletIsMineDescriptors, benchmark::PriorityLevel::LOW);
+BENCHMARK(WalletIsMineMigratedDescriptors, benchmark::PriorityLevel::LOW);
+#endif
+} // namespace wallet
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index e07771b17f..e10a17f003 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -1288,6 +1288,9 @@ bool LegacyScriptPubKeyMan::TopUp(unsigned int kpSize)
     }
     if (!batch.TxnCommit()) throw std::runtime_error(strprintf("Error during keypool top up. Cannot commit changes for wallet %s", m_storage.GetDisplayName()));
     NotifyCanGetAddressesChanged();
+    // Note: Unlike with DescriptorSPKM, LegacySPKM does not need to call
+    // m_storage.TopUpCallback() as we do not know what new scripts the LegacySPKM is
+    // watching for. CWallet's scriptPubKey cache is not used for LegacySPKMs.
     return true;
 }
 
@@ -2152,6 +2155,7 @@ bool DescriptorScriptPubKeyMan::TopUp(unsigned int size)
 bool DescriptorScriptPubKeyMan::TopUpWithDB(WalletBatch& batch, unsigned int size)
 {
     LOCK(cs_desc_man);
+    std::set<CScript> new_spks;
     unsigned int target_size;
     if (size > 0) {
         target_size = size;
@@ -2182,6 +2186,7 @@ bool DescriptorScriptPubKeyMan::TopUpWithDB(WalletBatch& batch, unsigned int siz
             if (!m_wallet_descriptor.descriptor->Expand(i, provider, scripts_temp, out_keys, &temp_cache)) return false;
         }
         // Add all of the scriptPubKeys to the scriptPubKey set
+        new_spks.insert(scripts_temp.begin(), scripts_temp.end());
         for (const CScript& script : scripts_temp) {
             m_map_script_pub_keys[script] = i;
         }
@@ -2207,6 +2212,7 @@ bool DescriptorScriptPubKeyMan::TopUpWithDB(WalletBatch& batch, unsigned int siz
     // By this point, the cache size should be the size of the entire range
     assert(m_wallet_descriptor.range_end - 1 == m_max_cached_index);
 
+    m_storage.TopUpCallback(new_spks, this);
     NotifyCanGetAddressesChanged();
     return true;
 }
@@ -2620,6 +2626,7 @@ uint256 DescriptorScriptPubKeyMan::GetID() const
 void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache)
 {
     LOCK(cs_desc_man);
+    std::set<CScript> new_spks;
     m_wallet_descriptor.cache = cache;
     for (int32_t i = m_wallet_descriptor.range_start; i < m_wallet_descriptor.range_end; ++i) {
         FlatSigningProvider out_keys;
@@ -2628,6 +2635,7 @@ void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache)
             throw std::runtime_error("Error: Unable to expand wallet descriptor from cache");
         }
         // Add all of the scriptPubKeys to the scriptPubKey set
+        new_spks.insert(scripts_temp.begin(), scripts_temp.end());
         for (const CScript& script : scripts_temp) {
             if (m_map_script_pub_keys.count(script) != 0) {
                 throw std::runtime_error(strprintf("Error: Already loaded script at index %d as being at index %d", i, m_map_script_pub_keys[script]));
@@ -2645,6 +2653,8 @@ void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache)
         }
         m_max_cached_index++;
     }
+    // Make sure the wallet knows about our new spks
+    m_storage.TopUpCallback(new_spks, this);
 }
 
 bool DescriptorScriptPubKeyMan::AddKey(const CKeyID& key_id, const CKey& key)
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 1c99e5ffcd..04f7f89d68 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -31,6 +31,7 @@ struct bilingual_str;
 
 namespace wallet {
 struct MigrationData;
+class ScriptPubKeyMan;
 
 // Wallet storage things that ScriptPubKeyMans need in order to be able to store things to the wallet database.
 // It provides access to things that are part of the entire wallet and not specific to a ScriptPubKeyMan such as
@@ -51,6 +52,8 @@ public:
     virtual bool WithEncryptionKey(std::function<bool (const CKeyingMaterial&)> cb) const = 0;
     virtual bool HasEncryptionKeys() const = 0;
     virtual bool IsLocked() const = 0;
+    //! Callback function for after TopUp completes containining any scripts that were added by a SPKMan
+    virtual void TopUpCallback(const std::set<CScript>&, ScriptPubKeyMan*) = 0;
 };
 
 //! Constant representing an unknown spkm creation time
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index db5f54cbb7..26c5256f6f 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -1571,11 +1571,22 @@ isminetype CWallet::IsMine(const CTxDestination& dest) const
 isminetype CWallet::IsMine(const CScript& script) const
 {
     AssertLockHeld(cs_wallet);
-    isminetype result = ISMINE_NO;
-    for (const auto& spk_man_pair : m_spk_managers) {
-        result = std::max(result, spk_man_pair.second->IsMine(script));
+
+    // Search the cache so that IsMine is called only on the relevant SPKMs instead of on everything in m_spk_managers
+    const auto& it = m_cached_spks.find(script);
+    if (it != m_cached_spks.end()) {
+        isminetype res = ISMINE_NO;
+        for (const auto& spkm : it->second) {
+            res = std::max(res, spkm->IsMine(script));
+        }
+        Assume(res == ISMINE_SPENDABLE);
+        return res;
     }
-    return result;
+
+    // Legacy wallet
+    if (IsLegacy()) return GetLegacyScriptPubKeyMan()->IsMine(script);
+
+    return ISMINE_NO;
 }
 
 bool CWallet::IsMine(const CTransaction& tx) const
@@ -3474,12 +3485,18 @@ ScriptPubKeyMan* CWallet::GetScriptPubKeyMan(const OutputType& type, bool intern
 std::set<ScriptPubKeyMan*> CWallet::GetScriptPubKeyMans(const CScript& script) const
 {
     std::set<ScriptPubKeyMan*> spk_mans;
-    SignatureData sigdata;
-    for (const auto& spk_man_pair : m_spk_managers) {
-        if (spk_man_pair.second->CanProvide(script, sigdata)) {
-            spk_mans.insert(spk_man_pair.second.get());
-        }
+
+    // Search the cache for relevant SPKMs instead of iterating m_spk_managers
+    const auto& it = m_cached_spks.find(script);
+    if (it != m_cached_spks.end()) {
+        spk_mans.insert(it->second.begin(), it->second.end());
     }
+    SignatureData sigdata;
+    Assume(std::all_of(spk_mans.begin(), spk_mans.end(), [&script, &sigdata](ScriptPubKeyMan* spkm) { return spkm->CanProvide(script, sigdata); }));
+
+    // Legacy wallet
+    if (IsLegacy() && GetLegacyScriptPubKeyMan()->CanProvide(script, sigdata)) spk_mans.insert(GetLegacyScriptPubKeyMan());
+
     return spk_mans;
 }
 
@@ -3499,11 +3516,17 @@ std::unique_ptr<SigningProvider> CWallet::GetSolvingProvider(const CScript& scri
 
 std::unique_ptr<SigningProvider> CWallet::GetSolvingProvider(const CScript& script, SignatureData& sigdata) const
 {
-    for (const auto& spk_man_pair : m_spk_managers) {
-        if (spk_man_pair.second->CanProvide(script, sigdata)) {
-            return spk_man_pair.second->GetSolvingProvider(script);
-        }
+    // Search the cache for relevant SPKMs instead of iterating m_spk_managers
+    const auto& it = m_cached_spks.find(script);
+    if (it != m_cached_spks.end()) {
+        // All spkms for a given script must already be able to make a SigningProvider for the script, so just return the first one.
+        Assume(it->second.at(0)->CanProvide(script, sigdata));
+        return it->second.at(0)->GetSolvingProvider(script);
     }
+
+    // Legacy wallet
+    if (IsLegacy() && GetLegacyScriptPubKeyMan()->CanProvide(script, sigdata)) return GetLegacyScriptPubKeyMan()->GetSolvingProvider(script);
+
     return nullptr;
 }
 
@@ -3582,15 +3605,16 @@ void CWallet::ConnectScriptPubKeyManNotifiers()
     }
 }
 
-void CWallet::LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc)
+DescriptorScriptPubKeyMan& CWallet::LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc)
 {
+    DescriptorScriptPubKeyMan* spk_manager;
     if (IsWalletFlagSet(WALLET_FLAG_EXTERNAL_SIGNER)) {
-        auto spk_manager = std::unique_ptr<ScriptPubKeyMan>(new ExternalSignerScriptPubKeyMan(*this, desc, m_keypool_size));
-        AddScriptPubKeyMan(id, std::move(spk_manager));
+        spk_manager = new ExternalSignerScriptPubKeyMan(*this, desc, m_keypool_size);
     } else {
-        auto spk_manager = std::unique_ptr<ScriptPubKeyMan>(new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size));
-        AddScriptPubKeyMan(id, std::move(spk_manager));
+        spk_manager = new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size);
     }
+    AddScriptPubKeyMan(id, std::unique_ptr<ScriptPubKeyMan>(spk_manager));
+    return *spk_manager;
 }
 
 void CWallet::SetupDescriptorScriptPubKeyMans(const CExtKey& master_key)
@@ -3945,6 +3969,8 @@ bool CWallet::ApplyMigrationData(MigrationData& data, bilingual_str& error)
         if (ExtractDestination(script, dest)) not_migrated_dests.emplace(dest);
     }
 
+    Assume(!m_cached_spks.empty());
+
     for (auto& desc_spkm : data.desc_spkms) {
         if (m_spk_managers.count(desc_spkm->GetID()) > 0) {
             error = _("Error: Duplicate descriptors created during migration. Your wallet may be corrupted.");
@@ -4428,4 +4454,17 @@ util::Result<MigrationResult> MigrateLegacyToDescriptor(const std::string& walle
     }
     return res;
 }
+
+void CWallet::CacheNewScriptPubKeys(const std::set<CScript>& spks, ScriptPubKeyMan* spkm)
+{
+    for (const auto& script : spks) {
+        m_cached_spks[script].push_back(spkm);
+    }
+}
+
+void CWallet::TopUpCallback(const std::set<CScript>& spks, ScriptPubKeyMan* spkm)
+{
+    // Update scriptPubKey cache
+    CacheNewScriptPubKeys(spks, spkm);
+}
 } // namespace wallet
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 42ceda1071..8b0ee22276 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -422,6 +422,9 @@ private:
     // Same as 'AddActiveScriptPubKeyMan' but designed for use within a batch transaction context
     void AddActiveScriptPubKeyManWithDb(WalletBatch& batch, uint256 id, OutputType type, bool internal);
 
+    //! Cache of descriptor ScriptPubKeys used for IsMine. Maps ScriptPubKey to set of spkms
+    std::unordered_map<CScript, std::vector<ScriptPubKeyMan*>, SaltedSipHasher> m_cached_spks;
+
     /**
      * Catch wallet up to current chain, scanning new blocks, updating the best
      * block locator and m_last_block_processed, and registering for
@@ -994,7 +997,7 @@ public:
     void ConnectScriptPubKeyManNotifiers();
 
     //! Instantiate a descriptor ScriptPubKeyMan from the WalletDescriptor and load it
-    void LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc);
+    DescriptorScriptPubKeyMan& LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc);
 
     //! Adds the active ScriptPubKeyMan for the specified type and internal. Writes it to the wallet file
     //! @param[in] id The unique id for the ScriptPubKeyMan
@@ -1045,6 +1048,11 @@ public:
 
     //! Whether the (external) signer performs R-value signature grinding
     bool CanGrindR() const;
+
+    //! Add scriptPubKeys for this ScriptPubKeyMan into the scriptPubKey cache
+    void CacheNewScriptPubKeys(const std::set<CScript>& spks, ScriptPubKeyMan* spkm);
+
+    void TopUpCallback(const std::set<CScript>& spks, ScriptPubKeyMan* spkm) override;
 };
 
 /**
diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp
index 45c3e90220..3ebe4ee7c9 100644
--- a/src/wallet/walletdb.cpp
+++ b/src/wallet/walletdb.cpp
@@ -804,10 +804,10 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat
             strErr = strprintf("%s\nDetails: %s", strErr, e.what());
             return DBErrors::UNKNOWN_DESCRIPTOR;
         }
-        pwallet->LoadDescriptorScriptPubKeyMan(id, desc);
+        DescriptorScriptPubKeyMan& spkm = pwallet->LoadDescriptorScriptPubKeyMan(id, desc);
 
         // Prior to doing anything with this spkm, verify ID compatibility
-        if (id != pwallet->GetDescriptorScriptPubKeyMan(desc)->GetID()) {
+        if (id != spkm.GetID()) {
             strErr = "The descriptor ID calculated by the wallet differs from the one in DB";
             return DBErrors::CORRUPT;
         }