test: add reset_chainstate parameter for snapshot unittests

This CreateAndActivateUTXOSnapshot parameter is necessary once we
perform snapshot completion within ABC, since the existing UpdateTip
test will fail because the IBD chain that has generated the snapshot
will exceed the base of the snapshot.

Being able to test snapshots being loaded into a mostly-uninitialized
datadir allows for more realistic unittest scenarios.

3 files changed, 52 insertions, 10 deletions

diff --git a/src/test/util/chainstate.h b/src/test/util/chainstate.h
index 2f0021b114..79bc882215 100644
--- a/src/test/util/chainstate.h
+++ b/src/test/util/chainstate.h
@@ -11,6 +11,7 @@
 #include <node/context.h>
 #include <node/utxo_snapshot.h>
 #include <rpc/blockchain.h>
+#include <test/util/setup_common.h>
 #include <validation.h>
 
 #include <univalue.h>
@@ -20,11 +21,20 @@ const auto NoMalleation = [](AutoFile& file, node::SnapshotMetadata& meta){};
 /**
  * Create and activate a UTXO snapshot, optionally providing a function to
  * malleate the snapshot.
+ *
+ * If `reset_chainstate` is true, reset the original chainstate back to the genesis
+ * block. This allows us to simulate more realistic conditions in which a snapshot is
+ * loaded into an otherwise mostly-uninitialized datadir. It also allows us to test
+ * conditions that would otherwise cause shutdowns based on the IBD chainstate going
+ * past the snapshot it generated.
  */
 template<typename F = decltype(NoMalleation)>
 static bool
-CreateAndActivateUTXOSnapshot(node::NodeContext& node, const fs::path root, F malleation = NoMalleation)
+CreateAndActivateUTXOSnapshot(TestingSetup* fixture, F malleation = NoMalleation, bool reset_chainstate = false)
 {
+    node::NodeContext& node = fixture->m_node;
+    fs::path root = fixture->m_path_root;
+
     // Write out a snapshot to the test's tempdir.
     //
     int height;
@@ -47,6 +57,37 @@ CreateAndActivateUTXOSnapshot(node::NodeContext& node, const fs::path root, F ma
 
     malleation(auto_infile, metadata);
 
+    if (reset_chainstate) {
+        {
+            // What follows is code to selectively reset chainstate data without
+            // disturbing the existing BlockManager instance, which is needed to
+            // recognize the headers chain previously generated by the chainstate we're
+            // removing. Without those headers, we can't activate the snapshot below.
+            //
+            // This is a stripped-down version of node::LoadChainstate which
+            // preserves the block index.
+            LOCK(::cs_main);
+            uint256 gen_hash = node.chainman->ActiveChainstate().m_chain[0]->GetBlockHash();
+            node.chainman->ResetChainstates();
+            node.chainman->InitializeChainstate(node.mempool.get());
+            Chainstate& chain = node.chainman->ActiveChainstate();
+            Assert(chain.LoadGenesisBlock());
+            // These cache values will be corrected shortly in `MaybeRebalanceCaches`.
+            chain.InitCoinsDB(1 << 20, true, false, "");
+            chain.InitCoinsCache(1 << 20);
+            chain.CoinsTip().SetBestBlock(gen_hash);
+            chain.setBlockIndexCandidates.insert(node.chainman->m_blockman.LookupBlockIndex(gen_hash));
+            chain.LoadChainTip();
+            node.chainman->MaybeRebalanceCaches();
+        }
+        BlockValidationState state;
+        if (!node.chainman->ActiveChainstate().ActivateBestChain(state)) {
+            throw std::runtime_error(strprintf("ActivateBestChain failed. (%s)", state.ToString()));
+        }
+        Assert(
+            0 == WITH_LOCK(node.chainman->GetMutex(), return node.chainman->ActiveHeight()));
+    }
+
     return node.chainman->ActivateSnapshot(auto_infile, metadata, /*in_memory=*/true);
 }
 
diff --git a/src/test/validation_chainstate_tests.cpp b/src/test/validation_chainstate_tests.cpp
index 347a967b33..f868c0d4e6 100644
--- a/src/test/validation_chainstate_tests.cpp
+++ b/src/test/validation_chainstate_tests.cpp
@@ -89,7 +89,8 @@ BOOST_FIXTURE_TEST_CASE(chainstate_update_tip, TestChain100Setup)
     // After adding some blocks to the tip, best block should have changed.
     BOOST_CHECK(::g_best_block != curr_tip);
 
-    BOOST_REQUIRE(CreateAndActivateUTXOSnapshot(m_node, m_path_root));
+    BOOST_REQUIRE(CreateAndActivateUTXOSnapshot(
+        this, NoMalleation, /*reset_chainstate=*/ true));
 
     // Ensure our active chain is the snapshot chainstate.
     BOOST_CHECK(WITH_LOCK(::cs_main, return chainman.IsSnapshotActive()));
diff --git a/src/test/validation_chainstatemanager_tests.cpp b/src/test/validation_chainstatemanager_tests.cpp
index 70e4e28c5f..04195e7a45 100644
--- a/src/test/validation_chainstatemanager_tests.cpp
+++ b/src/test/validation_chainstatemanager_tests.cpp
@@ -185,7 +185,7 @@ struct SnapshotTestSetup : TestChain100Setup {
         Chainstate& validation_chainstate = chainman.ActiveChainstate();
 
         // Snapshot should refuse to load at this height.
-        BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(m_node, m_path_root));
+        BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(this));
         BOOST_CHECK(!chainman.ActiveChainstate().m_from_snapshot_blockhash);
         BOOST_CHECK(!chainman.SnapshotBlockhash());
 
@@ -198,7 +198,7 @@ struct SnapshotTestSetup : TestChain100Setup {
 
         // Should not load malleated snapshots
         BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(
-            m_node, m_path_root, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
+            this, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
                 // A UTXO is missing but count is correct
                 metadata.m_coins_count -= 1;
 
@@ -209,27 +209,27 @@ struct SnapshotTestSetup : TestChain100Setup {
                 auto_infile >> coin;
         }));
         BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(
-            m_node, m_path_root, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
+            this, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
                 // Coins count is larger than coins in file
                 metadata.m_coins_count += 1;
         }));
         BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(
-            m_node, m_path_root, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
+            this, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
                 // Coins count is smaller than coins in file
                 metadata.m_coins_count -= 1;
         }));
         BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(
-            m_node, m_path_root, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
+            this, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
                 // Wrong hash
                 metadata.m_base_blockhash = uint256::ZERO;
         }));
         BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(
-            m_node, m_path_root, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
+            this, [](AutoFile& auto_infile, SnapshotMetadata& metadata) {
                 // Wrong hash
                 metadata.m_base_blockhash = uint256::ONE;
         }));
 
-        BOOST_REQUIRE(CreateAndActivateUTXOSnapshot(m_node, m_path_root));
+        BOOST_REQUIRE(CreateAndActivateUTXOSnapshot(this));
 
         // Ensure our active chain is the snapshot chainstate.
         BOOST_CHECK(!chainman.ActiveChainstate().m_from_snapshot_blockhash->IsNull());
@@ -320,7 +320,7 @@ struct SnapshotTestSetup : TestChain100Setup {
         }
 
         // Snapshot should refuse to load after one has already loaded.
-        BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(m_node, m_path_root));
+        BOOST_REQUIRE(!CreateAndActivateUTXOSnapshot(this));
 
         // Snapshot blockhash should be unchanged.
         BOOST_CHECK_EQUAL(