reserialize_bank_fields_with_hash (#23916)

* reserialize_bank_with_new_accounts_hash

* Update runtime/src/serde_snapshot.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* Update runtime/src/serde_snapshot/tests.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* Update runtime/src/serde_snapshot/tests.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* pr feedback

Co-authored-by: Brooks Prumo <brooks@prumo.org>

core/src/accounts_hash_verifier.rs

@@ -144,10 +144,12 @@ impl AccountsHashVerifier {
             assert_eq!(accounts_package.expected_capitalization, lamports);
             assert_eq!(expected_hash, hash);
         };
+
         measure_hash.stop();
-        solana_runtime::serde_snapshot::reserialize_bank(
+        solana_runtime::serde_snapshot::reserialize_bank_with_new_accounts_hash(
             accounts_package.snapshot_links.path(),
             accounts_package.slot,
+            &accounts_package.accounts_hash,
         );
         datapoint_info!(
             "accounts_hash_verifier",

core/tests/snapshots.rs

@@ -283,9 +283,10 @@ mod tests {
             Some(SnapshotType::FullSnapshot),
         )
         .unwrap();
-        solana_runtime::serde_snapshot::reserialize_bank(
+        solana_runtime::serde_snapshot::reserialize_bank_with_new_accounts_hash(
             accounts_package.snapshot_links.path(),
             accounts_package.slot,
+            &last_bank.get_accounts_hash(),
         );
         let snapshot_package = SnapshotPackage::from(accounts_package);
         snapshot_utils::archive_snapshot_package(
@@ -509,9 +510,10 @@ mod tests {
                     .unwrap()
                     .take()
                     .unwrap();
-                solana_runtime::serde_snapshot::reserialize_bank(
+                solana_runtime::serde_snapshot::reserialize_bank_with_new_accounts_hash(
                     accounts_package.snapshot_links.path(),
                     accounts_package.slot,
+                    &Hash::default(),
                 );
                 let snapshot_package = SnapshotPackage::from(accounts_package);
                 pending_snapshot_package
@@ -551,7 +553,11 @@ mod tests {
         .unwrap();
 
         // files were saved off before we reserialized the bank in the hacked up accounts_hash_verifier stand-in.
-        solana_runtime::serde_snapshot::reserialize_bank(saved_snapshots_dir.path(), saved_slot);
+        solana_runtime::serde_snapshot::reserialize_bank_with_new_accounts_hash(
+            saved_snapshots_dir.path(),
+            saved_slot,
+            &Hash::default(),
+        );
 
         snapshot_utils::verify_snapshot_archive(
             saved_archive_path.unwrap(),

runtime/src/accounts_db.rs

@@ -5600,10 +5600,15 @@ impl AccountsDb {
                 expected_capitalization,
             )
             .unwrap(); // unwrap here will never fail since check_hash = false
+        self.set_accounts_hash(slot, hash);
+        (hash, total_lamports)
+    }
+
+    /// update hash for this slot in the 'bank_hashes' map
+    pub(crate) fn set_accounts_hash(&self, slot: Slot, hash: Hash) {
         let mut bank_hashes = self.bank_hashes.write().unwrap();
         let mut bank_hash_info = bank_hashes.get_mut(&slot).unwrap();
         bank_hash_info.snapshot_hash = hash;
-        (hash, total_lamports)
     }
 
     fn scan_snapshot_stores_with_cache(

runtime/src/serde_snapshot.rs

@@ -175,6 +175,18 @@ trait TypeContext<'a>: PartialEq {
     ) -> Result<AccountsDbFields<Self::SerializableAccountStorageEntry>, Error>
     where
         R: Read;
+
+    /// deserialize the bank from 'stream_reader'
+    /// modify the accounts_hash
+    /// reserialize the bank to 'stream_writer'
+    fn reserialize_bank_fields_with_hash<R, W>(
+        stream_reader: &mut BufReader<R>,
+        stream_writer: &mut BufWriter<W>,
+        accounts_hash: &Hash,
+    ) -> std::result::Result<(), Box<bincode::ErrorKind>>
+    where
+        R: Read,
+        W: Write;
 }
 
 fn deserialize_from<R, T>(reader: R) -> bincode::Result<T>
@@ -303,19 +315,55 @@ where
     })
 }
 
-pub fn reserialize_bank(bank_snapshots_dir: impl AsRef<Path>, slot: Slot) {
+/// deserialize the bank from 'stream_reader'
+/// modify the accounts_hash
+/// reserialize the bank to 'stream_writer'
+fn reserialize_bank_fields_with_new_hash<W, R>(
+    stream_reader: &mut BufReader<R>,
+    stream_writer: &mut BufWriter<W>,
+    accounts_hash: &Hash,
+) -> Result<(), Error>
+where
+    W: Write,
+    R: Read,
+{
+    newer::Context::reserialize_bank_fields_with_hash(stream_reader, stream_writer, accounts_hash)
+}
+
+/// effectively updates the accounts hash in the serialized bank file on disk
+/// read serialized bank from pre file
+/// update accounts_hash
+/// write serialized bank to post file
+/// return true if pre file found
+pub fn reserialize_bank_with_new_accounts_hash(
+    bank_snapshots_dir: impl AsRef<Path>,
+    slot: Slot,
+    accounts_hash: &Hash,
+) -> bool {
     let bank_post = snapshot_utils::get_bank_snapshots_dir(bank_snapshots_dir, slot);
     let bank_post = bank_post.join(snapshot_utils::get_snapshot_file_name(slot));
     let mut bank_pre = bank_post.clone();
     bank_pre.set_extension(BANK_SNAPSHOT_PRE_FILENAME_EXTENSION);
 
-    // some tests don't create the file
-    if bank_pre.is_file() {
-        // replace the original file with the newly serialized one
-        // atm, this just copies the file and deletes the old one
-        std::fs::copy(&bank_pre, bank_post).unwrap();
+    let mut found = false;
+    {
+        let file = std::fs::File::open(&bank_pre);
+        // some tests don't create the file
+        if let Ok(file) = file {
+            found = true;
+            let file_out = std::fs::File::create(bank_post).unwrap();
+            reserialize_bank_fields_with_new_hash(
+                &mut BufReader::new(file),
+                &mut BufWriter::new(file_out),
+                accounts_hash,
+            )
+            .unwrap();
+        }
+    }
+    if found {
         std::fs::remove_file(bank_pre).unwrap();
     }
+    found
 }
 
 struct SerializableBankAndStorage<'a, C> {

runtime/src/serde_snapshot/newer.rs

@@ -230,7 +230,7 @@ impl<'a> TypeContext<'a> for Context {
                 )
             }));
         let slot = serializable_db.slot;
-        let hash = serializable_db
+        let hash: BankHashInfo = serializable_db
             .accounts_db
             .bank_hashes
             .read()
@@ -287,4 +287,61 @@ impl<'a> TypeContext<'a> for Context {
     {
         deserialize_from(stream)
     }
+
+    /// deserialize the bank from 'stream_reader'
+    /// modify the accounts_hash
+    /// reserialize the bank to 'stream_writer'
+    fn reserialize_bank_fields_with_hash<R, W>(
+        stream_reader: &mut BufReader<R>,
+        stream_writer: &mut BufWriter<W>,
+        accounts_hash: &Hash,
+    ) -> std::result::Result<(), Box<bincode::ErrorKind>>
+    where
+        R: Read,
+        W: Write,
+    {
+        let (bank_fields, mut accounts_db_fields) =
+            Self::deserialize_bank_fields(stream_reader).unwrap();
+        accounts_db_fields.3.snapshot_hash = *accounts_hash;
+        let rhs = bank_fields;
+        let blockhash_queue = RwLock::new(rhs.blockhash_queue.clone());
+        let hard_forks = RwLock::new(rhs.hard_forks.clone());
+        let stakes_cache = StakesCache::new(rhs.stakes.clone());
+        let bank = SerializableVersionedBank {
+            blockhash_queue: &blockhash_queue,
+            ancestors: &rhs.ancestors,
+            hash: rhs.hash,
+            parent_hash: rhs.parent_hash,
+            parent_slot: rhs.parent_slot,
+            hard_forks: &hard_forks,
+            transaction_count: rhs.transaction_count,
+            tick_height: rhs.tick_height,
+            signature_count: rhs.signature_count,
+            capitalization: rhs.capitalization,
+            max_tick_height: rhs.max_tick_height,
+            hashes_per_tick: rhs.hashes_per_tick,
+            ticks_per_slot: rhs.ticks_per_slot,
+            ns_per_slot: rhs.ns_per_slot,
+            genesis_creation_time: rhs.genesis_creation_time,
+            slots_per_year: rhs.slots_per_year,
+            accounts_data_len: rhs.accounts_data_len,
+            slot: rhs.slot,
+            epoch: rhs.epoch,
+            block_height: rhs.block_height,
+            collector_id: rhs.collector_id,
+            collector_fees: rhs.collector_fees,
+            fee_calculator: rhs.fee_calculator,
+            fee_rate_governor: rhs.fee_rate_governor,
+            collected_rent: rhs.collected_rent,
+            rent_collector: rhs.rent_collector,
+            epoch_schedule: rhs.epoch_schedule,
+            inflation: rhs.inflation,
+            stakes: &stakes_cache,
+            unused_accounts: UnusedAccounts::default(),
+            epoch_stakes: &rhs.epoch_stakes,
+            is_delta: rhs.is_delta,
+        };
+
+        bincode::serialize_into(stream_writer, &(bank, accounts_db_fields))
+    }
 }

runtime/src/serde_snapshot/tests.rs

@@ -178,7 +178,11 @@ fn test_accounts_serialize_style(serde_style: SerdeStyle) {
     assert_eq!(accounts.bank_hash_at(0), daccounts.bank_hash_at(0));
 }
 
-fn test_bank_serialize_style(serde_style: SerdeStyle) {
+fn test_bank_serialize_style(
+    serde_style: SerdeStyle,
+    reserialize_accounts_hash: bool,
+    update_accounts_hash: bool,
+) {
     solana_logger::setup();
     let (genesis_config, _) = create_genesis_config(500);
     let bank0 = Arc::new(Bank::new_for_tests(&genesis_config));
@@ -214,6 +218,52 @@ fn test_bank_serialize_style(serde_style: SerdeStyle) {
     )
     .unwrap();
 
+    let accounts_hash = if update_accounts_hash {
+        let hash = Hash::new(&[1; 32]);
+        bank2
+            .accounts()
+            .accounts_db
+            .set_accounts_hash(bank2.slot(), hash);
+        hash
+    } else {
+        bank2.get_accounts_hash()
+    };
+    if reserialize_accounts_hash {
+        let slot = bank2.slot();
+        let temp_dir = TempDir::new().unwrap();
+        let slot_dir = temp_dir.path().join(slot.to_string());
+        let post_path = slot_dir.join(slot.to_string());
+        let mut pre_path = post_path.clone();
+        pre_path.set_extension(BANK_SNAPSHOT_PRE_FILENAME_EXTENSION);
+        std::fs::create_dir(&slot_dir).unwrap();
+        {
+            let mut f = std::fs::File::create(&pre_path).unwrap();
+            f.write_all(&buf).unwrap();
+        }
+        assert!(reserialize_bank_with_new_accounts_hash(
+            temp_dir.path(),
+            slot,
+            &accounts_hash
+        ));
+        let previous_len = buf.len();
+        // larger buffer than expected to make sure the file isn't larger than expected
+        let mut buf_reserialized = vec![0; previous_len + 1];
+        {
+            let mut f = std::fs::File::open(post_path).unwrap();
+            let size = f.read(&mut buf_reserialized).unwrap();
+            assert_eq!(size, previous_len);
+            buf_reserialized.truncate(size);
+        }
+        if update_accounts_hash {
+            // We cannot guarantee buffer contents are exactly the same if hash is the same.
+            // Things like hashsets/maps have randomness in their in-mem representations.
+            // This make serialized bytes not deterministic.
+            // But, we can guarantee that the buffer is different if we change the hash!
+            assert_ne!(buf, buf_reserialized);
+            std::mem::swap(&mut buf, &mut buf_reserialized);
+        }
+    }
+
     let rdr = Cursor::new(&buf[..]);
     let mut reader = std::io::BufReader::new(&buf[rdr.position() as usize..]);
 
@@ -250,6 +300,7 @@ fn test_bank_serialize_style(serde_style: SerdeStyle) {
     assert_eq!(dbank.get_balance(&key1.pubkey()), 0);
     assert_eq!(dbank.get_balance(&key2.pubkey()), 10);
     assert_eq!(dbank.get_balance(&key3.pubkey()), 0);
+    assert_eq!(dbank.get_accounts_hash(), accounts_hash);
     assert!(bank2 == dbank);
 }
 
@@ -296,7 +347,15 @@ fn test_accounts_serialize_newer() {
 
 #[test]
 fn test_bank_serialize_newer() {
-    test_bank_serialize_style(SerdeStyle::Newer)
+    for (reserialize_accounts_hash, update_accounts_hash) in
+        [(false, false), (true, false), (true, true)]
+    {
+        test_bank_serialize_style(
+            SerdeStyle::Newer,
+            reserialize_accounts_hash,
+            update_accounts_hash,
+        )
+    }
 }
 
 #[cfg(RUSTC_WITH_SPECIALIZATION)]

runtime/src/snapshot_utils.rs

@@ -1918,9 +1918,10 @@ pub fn package_and_archive_full_snapshot(
         Some(SnapshotType::FullSnapshot),
     )?;
 
-    crate::serde_snapshot::reserialize_bank(
+    crate::serde_snapshot::reserialize_bank_with_new_accounts_hash(
         accounts_package.snapshot_links.path(),
         accounts_package.slot,
+        &bank.get_accounts_hash(),
     );
 
     let snapshot_package = SnapshotPackage::from(accounts_package);
@@ -1964,9 +1965,10 @@ pub fn package_and_archive_incremental_snapshot(
         )),
     )?;
 
-    crate::serde_snapshot::reserialize_bank(
+    crate::serde_snapshot::reserialize_bank_with_new_accounts_hash(
         accounts_package.snapshot_links.path(),
         accounts_package.slot,
+        &bank.get_accounts_hash(),
     );
 
     let snapshot_package = SnapshotPackage::from(accounts_package);