From 83e967ae063eebe91dce4afa8c14be358ddc683a Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@anza.xyz>
Date: Tue, 30 Jul 2024 09:53:48 -0500
Subject: [PATCH] startup shrink all slots has to avoid the snapshot slot
 (#2339)

---
 accounts-db/src/accounts_db.rs      | 19 +++++++++++++++----
 runtime/src/bank.rs                 |  2 ++
 runtime/src/serde_snapshot/tests.rs |  4 ++--
 runtime/tests/accounts.rs           |  2 +-
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index 3ffd87efc6fbdc..c06fb5408466be 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -4829,17 +4829,28 @@ impl AccountsDb {
         num_candidates
     }
 
+    /// This is only called at startup from bank when we are being extra careful such as when we downloaded a snapshot.
+    /// Also called from tests.
+    /// `newest_slot_skip_shrink_inclusive` is used to avoid shrinking the slot we are loading a snapshot from. If we shrink that slot, we affect
+    /// the bank hash calculation verification at startup.
     pub fn shrink_all_slots(
         &self,
         is_startup: bool,
         last_full_snapshot_slot: Option<Slot>,
         epoch_schedule: &EpochSchedule,
+        newest_slot_skip_shrink_inclusive: Option<Slot>,
     ) {
         let _guard = self.active_stats.activate(ActiveStatItem::Shrink);
         const DIRTY_STORES_CLEANING_THRESHOLD: usize = 10_000;
         const OUTER_CHUNK_SIZE: usize = 2000;
+        let mut slots = self.all_slots_in_storage();
+        if let Some(newest_slot_skip_shrink_inclusive) = newest_slot_skip_shrink_inclusive {
+            // at startup, we cannot shrink the slot that we're about to replay and recalculate bank hash for.
+            // That storage's contents are used to verify the bank hash (and accounts delta hash) of the startup slot.
+            slots.retain(|slot| slot < &newest_slot_skip_shrink_inclusive);
+        }
+
         if is_startup {
-            let slots = self.all_slots_in_storage();
             let threads = num_cpus::get();
             let inner_chunk_size = std::cmp::max(OUTER_CHUNK_SIZE / threads, 1);
             slots.chunks(OUTER_CHUNK_SIZE).for_each(|chunk| {
@@ -4853,7 +4864,7 @@ impl AccountsDb {
                 }
             });
         } else {
-            for slot in self.all_slots_in_storage() {
+            for slot in slots {
                 self.shrink_slot_forced(slot);
                 if self.dirty_stores.len() > DIRTY_STORES_CLEANING_THRESHOLD {
                     self.clean_accounts(None, is_startup, last_full_snapshot_slot, epoch_schedule);
@@ -12052,7 +12063,7 @@ pub mod tests {
                 accounts.shrink_candidate_slots(&epoch_schedule);
             }
 
-            accounts.shrink_all_slots(*startup, None, &EpochSchedule::default());
+            accounts.shrink_all_slots(*startup, None, &EpochSchedule::default(), None);
         }
     }
 
@@ -12110,7 +12121,7 @@ pub mod tests {
         );
 
         // Now, do full-shrink.
-        accounts.shrink_all_slots(false, None, &EpochSchedule::default());
+        accounts.shrink_all_slots(false, None, &EpochSchedule::default(), None);
         assert_eq!(
             pubkey_count_after_shrink,
             accounts.all_account_count_in_accounts_file(shrink_slot)
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index ca62027030f39b..e34f9d0a470d57 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -5903,6 +5903,8 @@ impl Bank {
                     true,
                     Some(last_full_snapshot_slot),
                     self.epoch_schedule(),
+                    // we cannot allow the snapshot slot to be shrunk
+                    Some(self.slot()),
                 );
                 info!("Shrinking... Done.");
             } else {
diff --git a/runtime/src/serde_snapshot/tests.rs b/runtime/src/serde_snapshot/tests.rs
index 13dd0243141542..81616510d54491 100644
--- a/runtime/src/serde_snapshot/tests.rs
+++ b/runtime/src/serde_snapshot/tests.rs
@@ -823,7 +823,7 @@ mod serde_snapshot_tests {
                 pubkey_count,
                 accounts.all_account_count_in_accounts_file(shrink_slot)
             );
-            accounts.shrink_all_slots(*startup, None, &EpochSchedule::default());
+            accounts.shrink_all_slots(*startup, None, &EpochSchedule::default(), None);
             assert_eq!(
                 pubkey_count_after_shrink,
                 accounts.all_account_count_in_accounts_file(shrink_slot)
@@ -851,7 +851,7 @@ mod serde_snapshot_tests {
                 .unwrap();
 
             // repeating should be no-op
-            accounts.shrink_all_slots(*startup, None, &epoch_schedule);
+            accounts.shrink_all_slots(*startup, None, &epoch_schedule, None);
             assert_eq!(
                 pubkey_count_after_shrink,
                 accounts.all_account_count_in_accounts_file(shrink_slot)
diff --git a/runtime/tests/accounts.rs b/runtime/tests/accounts.rs
index bcb6f7746ac905..e1ff4c854a8483 100644
--- a/runtime/tests/accounts.rs
+++ b/runtime/tests/accounts.rs
@@ -40,7 +40,7 @@ fn test_shrink_and_clean() {
             if exit_for_shrink.load(Ordering::Relaxed) {
                 break;
             }
-            accounts_for_shrink.shrink_all_slots(false, None, &EpochSchedule::default());
+            accounts_for_shrink.shrink_all_slots(false, None, &EpochSchedule::default(), None);
         });
 
         let mut alive_accounts = vec![];