Adds tests for EAH and snapshot interactions (#28304)

core/tests/epoch_accounts_hash.rs

@@ -11,13 +11,18 @@ use {
             AbsRequestHandlers, AbsRequestSender, AccountsBackgroundService, DroppedSlotsReceiver,
             PrunedBanksRequestHandler, SnapshotRequestHandler,
         },
+        accounts_db::AccountShrinkThreshold,
         accounts_hash::CalcAccountsHashConfig,
+        accounts_index::AccountSecondaryIndexes,
         bank::Bank,
         bank_forks::BankForks,
         epoch_accounts_hash::{self, EpochAccountsHash},
         genesis_utils::{self, GenesisConfigInfo},
+        runtime_config::RuntimeConfig,
+        snapshot_archive_info::SnapshotArchiveInfoGetter,
         snapshot_config::SnapshotConfig,
         snapshot_package::{PendingAccountsPackage, PendingSnapshotPackage},
+        snapshot_utils,
     },
     solana_sdk::{
         clock::Slot,
@@ -39,6 +44,7 @@ use {
         time::Duration,
     },
     tempfile::TempDir,
+    test_case::test_case,
 };
 
 struct TestEnvironment {
@@ -49,7 +55,7 @@ struct TestEnvironment {
     bank_forks: Arc<RwLock<BankForks>>,
     background_services: BackgroundServices,
     genesis_config_info: GenesisConfigInfo,
-    _snapshot_config: SnapshotConfig,
+    snapshot_config: SnapshotConfig,
     _bank_snapshots_dir: TempDir,
     _full_snapshot_archives_dir: TempDir,
     _incremental_snapshot_archives_dir: TempDir,
@@ -64,6 +70,24 @@ impl TestEnvironment {
 
     #[must_use]
     fn new() -> TestEnvironment {
+        Self::_new(SnapshotConfig::new_load_only())
+    }
+
+    #[must_use]
+    fn new_with_snapshots(
+        full_snapshot_archive_interval_slots: Slot,
+        incremental_snapshot_archive_interval_slots: Slot,
+    ) -> TestEnvironment {
+        let snapshot_config = SnapshotConfig {
+            full_snapshot_archive_interval_slots,
+            incremental_snapshot_archive_interval_slots,
+            ..SnapshotConfig::default()
+        };
+        Self::_new(snapshot_config)
+    }
+
+    #[must_use]
+    fn _new(snapshot_config: SnapshotConfig) -> TestEnvironment {
         let bank_snapshots_dir = TempDir::new().unwrap();
         let full_snapshot_archives_dir = TempDir::new().unwrap();
         let incremental_snapshot_archives_dir = TempDir::new().unwrap();
@@ -80,7 +104,7 @@ impl TestEnvironment {
                 .path()
                 .to_path_buf(),
             bank_snapshots_dir: bank_snapshots_dir.path().to_path_buf(),
-            ..SnapshotConfig::new_load_only()
+            ..snapshot_config
         };
 
         let mut bank_forks =
@@ -124,7 +148,7 @@ impl TestEnvironment {
             _bank_snapshots_dir: bank_snapshots_dir,
             _full_snapshot_archives_dir: full_snapshot_archives_dir,
             _incremental_snapshot_archives_dir: incremental_snapshot_archives_dir,
-            _snapshot_config: snapshot_config,
+            snapshot_config,
             background_services,
         }
     }
@@ -222,20 +246,21 @@ impl Drop for BackgroundServices {
     }
 }
 
-/// Run through a few epochs and ensure the Epoch Accounts Hash is calculated correctly
-#[test]
-fn test_epoch_accounts_hash() {
+/// Ensure that EAHs are requested, calculated, and awaited correctly.
+/// Test both with and without snapshots to make sure they don't interfere with EAH.
+#[test_case(TestEnvironment::new()                      ; "without snapshots")]
+#[test_case(TestEnvironment::new_with_snapshots(20, 10) ; "with snapshots")]
+fn test_epoch_accounts_hash_basic(test_environment: TestEnvironment) {
     solana_logger::setup();
 
     const NUM_EPOCHS_TO_TEST: u64 = 2;
     const SET_ROOT_INTERVAL: Slot = 3;
 
-    let test_config = TestEnvironment::new();
-    let bank_forks = &test_config.bank_forks;
+    let bank_forks = &test_environment.bank_forks;
 
     let mut expected_epoch_accounts_hash = None;
 
-    let slots_per_epoch = test_config
+    let slots_per_epoch = test_environment
         .genesis_config_info
         .genesis_config
         .epoch_schedule
@@ -250,7 +275,7 @@ fn test_epoch_accounts_hash() {
             ));
 
             let transaction = system_transaction::transfer(
-                &test_config.genesis_config_info.mint_keypair,
+                &test_environment.genesis_config_info.mint_keypair,
                 &Pubkey::new_unique(),
                 1,
                 bank.last_blockhash(),
@@ -267,7 +292,7 @@ fn test_epoch_accounts_hash() {
             trace!("rooting bank {}", bank.slot());
             bank_forks.write().unwrap().set_root(
                 bank.slot(),
-                &test_config
+                &test_environment
                     .background_services
                     .accounts_background_request_sender,
                 None,
@@ -322,3 +347,132 @@ fn test_epoch_accounts_hash() {
         std::thread::yield_now();
     }
 }
+
+/// Ensure that snapshots always have the expected EAH
+///
+/// Generate snapshots:
+/// - Before EAH start
+/// - After EAH start but before EAH stop
+/// - After EAH stop
+///
+/// In Epoch 0, this will correspond to all three EAH states (invalid, in-flight, and valid). In
+/// Epoch 1, this will correspond to a normal running cluster, where EAH will only be either
+/// in-flight or valid.
+#[test]
+fn test_snapshots_have_expected_epoch_accounts_hash() {
+    solana_logger::setup();
+
+    const NUM_EPOCHS_TO_TEST: u64 = 2;
+
+    // Since slots-per-epoch is 100, EAH start will be slots 25 and 125, and EAH stop will be slots
+    // 75 and 175.  Pick a full snapshot interval that triggers in the three scenarios outlined in
+    // the test's description.
+    const FULL_SNAPSHOT_INTERVAL: Slot = 20;
+
+    let test_environment =
+        TestEnvironment::new_with_snapshots(FULL_SNAPSHOT_INTERVAL, FULL_SNAPSHOT_INTERVAL);
+    let bank_forks = &test_environment.bank_forks;
+
+    let slots_per_epoch = test_environment
+        .genesis_config_info
+        .genesis_config
+        .epoch_schedule
+        .slots_per_epoch;
+    for _ in 0..slots_per_epoch * NUM_EPOCHS_TO_TEST {
+        let bank = {
+            let parent = bank_forks.read().unwrap().working_bank();
+            let bank = bank_forks.write().unwrap().insert(Bank::new_from_parent(
+                &parent,
+                &Pubkey::default(),
+                parent.slot() + 1,
+            ));
+
+            let transaction = system_transaction::transfer(
+                &test_environment.genesis_config_info.mint_keypair,
+                &Pubkey::new_unique(),
+                1,
+                bank.last_blockhash(),
+            );
+            bank.process_transaction(&transaction).unwrap();
+            bank.fill_bank_with_ticks_for_tests();
+
+            bank
+        };
+        trace!("new bank {}", bank.slot());
+
+        // Root every bank.  This is what a normal validator does as well.
+        // `set_root()` is also what requests snapshots and EAH calculations.
+        bank_forks.write().unwrap().set_root(
+            bank.slot(),
+            &test_environment
+                .background_services
+                .accounts_background_request_sender,
+            None,
+        );
+
+        // After submitting an EAH calculation request, wait until it gets handled by ABS so that
+        // subsequent snapshot requests are not swallowed.
+        if bank.slot() == epoch_accounts_hash::calculation_start(&bank) {
+            while dbg!(bank.epoch_accounts_hash()).is_none() {
+                std::thread::sleep(Duration::from_secs(1));
+            }
+        }
+
+        // After submitting a snapshot request...
+        // - Wait until the snapshot archive has been generated
+        // - Deserialize the bank from the snapshot archive
+        // - Ensure the EAHs match
+        if bank.slot() % FULL_SNAPSHOT_INTERVAL == 0 {
+            let snapshot_config = &test_environment.snapshot_config;
+            let full_snapshot_archive_info = loop {
+                if let Some(full_snapshot_archive_info) =
+                    snapshot_utils::get_highest_full_snapshot_archive_info(
+                        &snapshot_config.full_snapshot_archives_dir,
+                    )
+                {
+                    if full_snapshot_archive_info.slot() == bank.slot() {
+                        break full_snapshot_archive_info;
+                    }
+                }
+
+                _ = dbg!(snapshot_utils::get_full_snapshot_archives(
+                    &snapshot_config.full_snapshot_archives_dir
+                ));
+                std::thread::sleep(Duration::from_millis(1000));
+            };
+
+            let accounts_dir = TempDir::new().unwrap();
+            let deserialized_bank = snapshot_utils::bank_from_snapshot_archives(
+                &[accounts_dir.into_path()],
+                &snapshot_config.bank_snapshots_dir,
+                &full_snapshot_archive_info,
+                None,
+                &test_environment.genesis_config_info.genesis_config,
+                &RuntimeConfig::default(),
+                None,
+                None,
+                AccountSecondaryIndexes::default(),
+                false,
+                None,
+                AccountShrinkThreshold::default(),
+                true,
+                true,
+                true,
+                None,
+                None,
+                &Arc::new(AtomicBool::new(false)),
+            )
+            .unwrap()
+            .0;
+
+            assert_eq!(&deserialized_bank, bank.as_ref());
+            assert_eq!(
+                deserialized_bank.epoch_accounts_hash(),
+                bank.epoch_accounts_hash(),
+            );
+        }
+
+        // Give the background services a chance to run
+        std::thread::yield_now();
+    }
+}