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

1 files changed, 57 insertions, 18 deletions

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