From bc35e1c5f57e653fe75ca06a2913ff1a37f58b9b Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <wash678@gmail.com>
Date: Tue, 22 Mar 2022 21:27:54 -0500
Subject: [PATCH] snapshot code needs all storages for hash calc (#23840)

---
 runtime/src/snapshot_package.rs | 21 +++++++++-----
 runtime/src/snapshot_utils.rs   | 51 +++++++++------------------------
 2 files changed, 28 insertions(+), 44 deletions(-)

diff --git a/runtime/src/snapshot_package.rs b/runtime/src/snapshot_package.rs
index 0547f554be..6f2bd321c2 100644
--- a/runtime/src/snapshot_package.rs
+++ b/runtime/src/snapshot_package.rs
@@ -78,12 +78,6 @@ impl AccountsPackage {
                 bank.slot() > incremental_snapshot_base_slot,
                 "Incremental snapshot base slot must be less than the bank being snapshotted!"
             );
-            assert!(
-            snapshot_storages.iter().all(|storage| storage
-                .iter()
-                .all(|entry| entry.slot() > incremental_snapshot_base_slot)),
-            "Incremental snapshot package must only contain storage entries where slot > incremental snapshot base slot (i.e. full snapshot slot)!"
-            );
         }
 
         // Hard link the snapshot into a tmpdir, to ensure its not removed prior to packaging.
@@ -136,6 +130,7 @@ impl From<AccountsPackage> for SnapshotPackage {
             "Cannot make a SnapshotPackage from an AccountsPackage when SnapshotType is None!"
         );
 
+        let mut snapshot_storages = accounts_package.snapshot_storages;
         let snapshot_archive_path = match accounts_package.snapshot_type.unwrap() {
             SnapshotType::FullSnapshot => snapshot_utils::build_full_snapshot_archive_path(
                 accounts_package.snapshot_archives_dir,
@@ -144,6 +139,18 @@ impl From<AccountsPackage> for SnapshotPackage {
                 accounts_package.archive_format,
             ),
             SnapshotType::IncrementalSnapshot(incremental_snapshot_base_slot) => {
+                snapshot_storages.retain(|storages| {
+                    storages
+                        .first() // storages are grouped by slot in the outer Vec, so all storages will have the same slot as the first
+                        .map(|storage| storage.slot() > incremental_snapshot_base_slot)
+                        .unwrap_or_default()
+                });
+                assert!(
+                    snapshot_storages.iter().all(|storage| storage
+                        .iter()
+                        .all(|entry| entry.slot() > incremental_snapshot_base_slot)),
+                    "Incremental snapshot package must only contain storage entries where slot > incremental snapshot base slot (i.e. full snapshot slot)!"
+                    );
                 snapshot_utils::build_incremental_snapshot_archive_path(
                     accounts_package.snapshot_archives_dir,
                     incremental_snapshot_base_slot,
@@ -164,7 +171,7 @@ impl From<AccountsPackage> for SnapshotPackage {
             block_height: accounts_package.block_height,
             slot_deltas: accounts_package.slot_deltas,
             snapshot_links: accounts_package.snapshot_links,
-            snapshot_storages: accounts_package.snapshot_storages,
+            snapshot_storages,
             snapshot_version: accounts_package.snapshot_version,
             snapshot_type: accounts_package.snapshot_type.unwrap(),
         }
diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index 608b2c1682..7bd2c805b8 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -1636,7 +1636,7 @@ pub fn snapshot_bank(
     hash_for_testing: Option<Hash>,
     snapshot_type: Option<SnapshotType>,
 ) -> Result<()> {
-    let snapshot_storages = get_snapshot_storages(root_bank, snapshot_type);
+    let snapshot_storages = get_snapshot_storages(root_bank);
 
     let mut add_snapshot_time = Measure::start("add-snapshot-ms");
     let bank_snapshot_info = add_bank_snapshot(
@@ -1667,44 +1667,21 @@ pub fn snapshot_bank(
     Ok(())
 }
 
-/// Get the snapshot storages for this bank and snapshot_type
-fn get_snapshot_storages(bank: &Bank, snapshot_type: Option<SnapshotType>) -> SnapshotStorages {
+/// Get the snapshot storages for this bank
+fn get_snapshot_storages(bank: &Bank) -> SnapshotStorages {
     let mut measure_snapshot_storages = Measure::start("snapshot-storages");
-    let snapshot_storages = snapshot_type.map_or_else(SnapshotStorages::default, |snapshot_type| {
-        let incremental_snapshot_base_slot = match snapshot_type {
-            SnapshotType::IncrementalSnapshot(incremental_snapshot_base_slot) => {
-                Some(incremental_snapshot_base_slot)
-            }
-            _ => None,
-        };
-        bank.get_snapshot_storages(incremental_snapshot_base_slot)
-    });
+    let snapshot_storages = bank.get_snapshot_storages(None);
     measure_snapshot_storages.stop();
-
-    if let Some(snapshot_type) = snapshot_type {
-        let snapshot_storages_count_name;
-        let snapshot_storages_time_name;
-        match snapshot_type {
-            SnapshotType::FullSnapshot => {
-                snapshot_storages_count_name = "full-snapshot-storages-count";
-                snapshot_storages_time_name = "full-snapshot-storages-time-ms";
-            }
-            SnapshotType::IncrementalSnapshot(_) => {
-                snapshot_storages_count_name = "incremental-snapshot-storages-count";
-                snapshot_storages_time_name = "incremental-snapshot-storages-time-ms";
-            }
-        }
-        let snapshot_storages_count = snapshot_storages.iter().map(Vec::len).sum::<usize>();
-        datapoint_info!(
-            "get_snapshot_storages",
-            (snapshot_storages_count_name, snapshot_storages_count, i64),
-            (
-                snapshot_storages_time_name,
-                measure_snapshot_storages.as_ms(),
-                i64
-            ),
-        );
-    }
+    let snapshot_storages_count = snapshot_storages.iter().map(Vec::len).sum::<usize>();
+    datapoint_info!(
+        "get_snapshot_storages",
+        ("snapshot-storages-count", snapshot_storages_count, i64),
+        (
+            "snapshot-storages-time-ms",
+            measure_snapshot_storages.as_ms(),
+            i64
+        ),
+    );
 
     snapshot_storages
 }