add ChainstateManager::MaybeRebalanceCaches()

Aside from in unittests, this method is unused at the moment. It will be used
in upcoming commits that enable utxo snapshot activation.

4 files changed, 99 insertions, 0 deletions

diff --git a/src/init.cpp b/src/init.cpp
index 4561a10870..9deb6cd656 100644
--- a/src/init.cpp
+++ b/src/init.cpp
@@ -1562,6 +1562,9 @@ bool AppInitMain(const util::Ref& context, NodeContext& node)
             try {
                 LOCK(cs_main);
                 chainman.InitializeChainstate();
+                chainman.m_total_coinstip_cache = nCoinCacheUsage;
+                chainman.m_total_coinsdb_cache = nCoinDBCache;
+
                 UnloadBlockIndex();
 
                 // new CBlockTreeDB tries to delete the existing file, which
diff --git a/src/test/validation_chainstatemanager_tests.cpp b/src/test/validation_chainstatemanager_tests.cpp
index d16633c149..f99191b08f 100644
--- a/src/test/validation_chainstatemanager_tests.cpp
+++ b/src/test/validation_chainstatemanager_tests.cpp
@@ -103,4 +103,58 @@ BOOST_AUTO_TEST_CASE(chainstatemanager)
     WITH_LOCK(::cs_main, manager.Unload());
 }
 
+//! Test rebalancing the caches associated with each chainstate.
+BOOST_AUTO_TEST_CASE(chainstatemanager_rebalance_caches)
+{
+    ChainstateManager manager;
+    size_t max_cache = 10000;
+    manager.m_total_coinsdb_cache = max_cache;
+    manager.m_total_coinstip_cache = max_cache;
+
+    std::vector<CChainState*> chainstates;
+
+    // Create a legacy (IBD) chainstate.
+    //
+    ENTER_CRITICAL_SECTION(cs_main);
+    CChainState& c1 = manager.InitializeChainstate();
+    LEAVE_CRITICAL_SECTION(cs_main);
+    chainstates.push_back(&c1);
+    c1.InitCoinsDB(
+        /* cache_size_bytes */ 1 << 23, /* in_memory */ true, /* should_wipe */ false);
+
+    {
+        LOCK(::cs_main);
+        c1.InitCoinsCache(1 << 23);
+        c1.CoinsTip().SetBestBlock(InsecureRand256());
+        manager.MaybeRebalanceCaches();
+    }
+
+    BOOST_CHECK_EQUAL(c1.m_coinstip_cache_size_bytes, max_cache);
+    BOOST_CHECK_EQUAL(c1.m_coinsdb_cache_size_bytes, max_cache);
+
+    // Create a snapshot-based chainstate.
+    //
+    ENTER_CRITICAL_SECTION(cs_main);
+    CChainState& c2 = manager.InitializeChainstate(GetRandHash());
+    LEAVE_CRITICAL_SECTION(cs_main);
+    chainstates.push_back(&c2);
+    c2.InitCoinsDB(
+        /* cache_size_bytes */ 1 << 23, /* in_memory */ true, /* should_wipe */ false);
+
+    {
+        LOCK(::cs_main);
+        c2.InitCoinsCache(1 << 23);
+        c2.CoinsTip().SetBestBlock(InsecureRand256());
+        manager.MaybeRebalanceCaches();
+    }
+
+    // Since both chainstates are considered to be in initial block download,
+    // the snapshot chainstate should take priority.
+    BOOST_CHECK_CLOSE(c1.m_coinstip_cache_size_bytes, max_cache * 0.05, 1);
+    BOOST_CHECK_CLOSE(c1.m_coinsdb_cache_size_bytes, max_cache * 0.05, 1);
+    BOOST_CHECK_CLOSE(c2.m_coinstip_cache_size_bytes, max_cache * 0.95, 1);
+    BOOST_CHECK_CLOSE(c2.m_coinsdb_cache_size_bytes, max_cache * 0.95, 1);
+
+}
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/validation.cpp b/src/validation.cpp
index 1d7d85b751..906ed943b2 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -5297,3 +5297,33 @@ void ChainstateManager::Reset()
     m_active_chainstate = nullptr;
     m_snapshot_validated = false;
 }
+
+void ChainstateManager::MaybeRebalanceCaches()
+{
+    if (m_ibd_chainstate && !m_snapshot_chainstate) {
+        LogPrintf("[snapshot] allocating all cache to the IBD chainstate\n");
+        // Allocate everything to the IBD chainstate.
+        m_ibd_chainstate->ResizeCoinsCaches(m_total_coinstip_cache, m_total_coinsdb_cache);
+    }
+    else if (m_snapshot_chainstate && !m_ibd_chainstate) {
+        LogPrintf("[snapshot] allocating all cache to the snapshot chainstate\n");
+        // Allocate everything to the snapshot chainstate.
+        m_snapshot_chainstate->ResizeCoinsCaches(m_total_coinstip_cache, m_total_coinsdb_cache);
+    }
+    else if (m_ibd_chainstate && m_snapshot_chainstate) {
+        // If both chainstates exist, determine who needs more cache based on IBD status.
+        //
+        // Note: shrink caches first so that we don't inadvertently overwhelm available memory.
+        if (m_snapshot_chainstate->IsInitialBlockDownload()) {
+            m_ibd_chainstate->ResizeCoinsCaches(
+                m_total_coinstip_cache * 0.05, m_total_coinsdb_cache * 0.05);
+            m_snapshot_chainstate->ResizeCoinsCaches(
+                m_total_coinstip_cache * 0.95, m_total_coinsdb_cache * 0.95);
+        } else {
+            m_snapshot_chainstate->ResizeCoinsCaches(
+                m_total_coinstip_cache * 0.05, m_total_coinsdb_cache * 0.05);
+            m_ibd_chainstate->ResizeCoinsCaches(
+                m_total_coinstip_cache * 0.95, m_total_coinsdb_cache * 0.95);
+        }
+    }
+}
diff --git a/src/validation.h b/src/validation.h
index b201403585..a4948c0187 100644
--- a/src/validation.h
+++ b/src/validation.h
@@ -794,6 +794,14 @@ public:
     //! chainstate to avoid duplicating block metadata.
     BlockManager m_blockman GUARDED_BY(::cs_main);
 
+    //! The total number of bytes available for us to use across all in-memory
+    //! coins caches. This will be split somehow across chainstates.
+    int64_t m_total_coinstip_cache{0};
+    //
+    //! The total number of bytes available for us to use across all leveldb
+    //! coins databases. This will be split somehow across chainstates.
+    int64_t m_total_coinsdb_cache{0};
+
     //! Instantiate a new chainstate and assign it based upon whether it is
     //! from a snapshot.
     //!
@@ -881,6 +889,10 @@ public:
 
     //! Clear (deconstruct) chainstate data.
     void Reset();
+
+    //! Check to see if caches are out of balance and if so, call
+    //! ResizeCoinsCaches() as needed.
+    void MaybeRebalanceCaches() EXCLUSIVE_LOCKS_REQUIRED(::cs_main);
 };
 
 /** DEPRECATED! Please use node.chainman instead. May only be used in validation.cpp internally */