Test multiple chain resets (#2897)

* Try simulating a chain growth

* Adjust the transaction expiry height

The mempool evicts expired transactions. When working with mocked data,
appending a new block typically clears the mempool because transactions become
expired. For this reason, the expiry height of each transactions is adjusted so
that it is greater than the new chain tip's height.

* Refactor the code so that it works with `VerifiedUnminedTx`

* Fix a typo

* Fix clippy warnings

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

zebra-state/src/service/chain_tip.rs

@@ -55,7 +55,7 @@ pub struct ChainTipBlock {
     ///
     /// If the best chain fork has changed, or some blocks have been skipped,
     /// this hash will be different to the last returned `ChainTipBlock.hash`.
-    pub(crate) previous_block_hash: block::Hash,
+    pub previous_block_hash: block::Hash,
 }
 
 impl From<ContextuallyValidBlock> for ChainTipBlock {

zebrad/src/components/mempool/tests/prop.rs

@@ -1,10 +1,13 @@
 //! Randomised property tests for the mempool.
 
+use proptest::collection::vec;
 use proptest::prelude::*;
+use proptest_derive::Arbitrary;
+
 use tokio::time;
 use tower::{buffer::Buffer, util::BoxService};
 
-use zebra_chain::{parameters::Network, transaction::VerifiedUnminedTx};
+use zebra_chain::{block, parameters::Network, transaction::VerifiedUnminedTx};
 use zebra_consensus::{error::TransactionError, transaction as tx};
 use zebra_network as zn;
 use zebra_state::{self as zs, ChainTipBlock, ChainTipSender};
@@ -24,6 +27,8 @@ type MockState = MockService<zs::Request, zs::Response, PropTestAssertion>;
 /// A [`MockService`] representing the Zebra transaction verifier service.
 type MockTxVerifier = MockService<tx::Request, tx::Response, PropTestAssertion, TransactionError>;
 
+const CHAIN_LENGTH: usize = 10;
+
 proptest! {
     /// Test if the mempool storage is cleared on a chain reset.
     #[test]
@@ -79,6 +84,94 @@ proptest! {
         })?;
     }
 
+    /// Test if the mempool storage is cleared on multiple chain resets.
+    #[test]
+    fn storage_is_cleared_on_chain_resets(
+        network in any::<Network>(),
+        mut previous_chain_tip in any::<ChainTipBlock>(),
+        mut transactions in vec(any::<VerifiedUnminedTx>(), 0..CHAIN_LENGTH),
+        fake_chain_tips in vec(any::<FakeChainTip>(), 0..CHAIN_LENGTH),
+    ) {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("Failed to create Tokio runtime");
+        let _guard = runtime.enter();
+
+        runtime.block_on(async move {
+            let (
+                mut mempool,
+                mut peer_set,
+                mut state_service,
+                mut tx_verifier,
+                mut recent_syncs,
+                mut chain_tip_sender,
+            ) = setup(network);
+
+            time::pause();
+
+            mempool.enable(&mut recent_syncs).await;
+
+            // Set the initial chain tip.
+            chain_tip_sender.set_best_non_finalized_tip(previous_chain_tip.clone());
+
+            // Call the mempool so that it is aware of the initial chain tip.
+            mempool.dummy_call().await;
+
+            for (fake_chain_tip, transaction) in fake_chain_tips.iter().zip(transactions.iter_mut()) {
+                // Obtain a new chain tip based on the previous one.
+                let chain_tip = fake_chain_tip.to_chain_tip_block(&previous_chain_tip);
+
+                // Adjust the transaction expiry height based on the new chain
+                // tip height so that the mempool does not evict the transaction
+                // when there is a chain growth.
+                if let Some(expiry_height) = transaction.transaction.transaction.expiry_height() {
+                    if chain_tip.height >= expiry_height {
+                        let mut tmp_tx = (*transaction.transaction.transaction).clone();
+
+                        // Set a new expiry height that is greater than the
+                        // height of the current chain tip.
+                        *tmp_tx.expiry_height_mut() = block::Height(chain_tip.height.0 + 1);
+                        transaction.transaction = tmp_tx.into();
+                    }
+                }
+
+                // Insert the dummy transaction into the mempool.
+                mempool
+                    .storage()
+                    .insert(transaction.clone())
+                    .expect("Inserting a transaction should succeed");
+
+                // Set the new chain tip.
+                chain_tip_sender.set_best_non_finalized_tip(chain_tip.clone());
+
+                // Call the mempool so that it is aware of the new chain tip.
+                mempool.dummy_call().await;
+
+                match fake_chain_tip {
+                    FakeChainTip::Grow(_) => {
+                        // The mempool should not be empty because we had a regular chain growth.
+                        prop_assert_ne!(mempool.storage().transaction_count(), 0);
+                    }
+
+                    FakeChainTip::Reset(_) => {
+                        // The mempool should be empty because we had a chain tip reset.
+                        prop_assert_eq!(mempool.storage().transaction_count(), 0);
+                    },
+                }
+
+                // Remember the current chain tip so that the next one can refer to it.
+                previous_chain_tip = chain_tip;
+            }
+
+            peer_set.expect_no_requests().await?;
+            state_service.expect_no_requests().await?;
+            tx_verifier.expect_no_requests().await?;
+
+            Ok(())
+        })?;
+    }
+
     /// Test if the mempool storage is cleared if the syncer falls behind and starts to catch up.
     #[test]
     fn storage_is_cleared_if_syncer_falls_behind(
@@ -173,3 +266,28 @@ fn setup(
         chain_tip_sender,
     )
 }
+
+/// A helper enum for simulating either a chain reset or growth.
+#[derive(Arbitrary, Debug, Clone)]
+enum FakeChainTip {
+    Grow(ChainTipBlock),
+    Reset(ChainTipBlock),
+}
+
+impl FakeChainTip {
+    /// Returns a new [`ChainTipBlock`] placed on top of the previous block if
+    /// the chain is supposed to grow. Otherwise returns a [`ChainTipBlock`]
+    /// that does not reference the previous one.
+    fn to_chain_tip_block(&self, previous: &ChainTipBlock) -> ChainTipBlock {
+        match self {
+            Self::Grow(chain_tip_block) => ChainTipBlock {
+                hash: chain_tip_block.hash,
+                height: block::Height(previous.height.0 + 1),
+                transaction_hashes: chain_tip_block.transaction_hashes.clone(),
+                previous_block_hash: previous.hash,
+            },
+
+            Self::Reset(chain_tip_block) => chain_tip_block.clone(),
+        }
+    }
+}