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.

1 files changed, 42 insertions, 1 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);
 }