From 94d3ec37d2d35b5a4915759e071c9400e2f3d567 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <wash678@gmail.com>
Date: Mon, 18 Apr 2022 22:44:19 -0500
Subject: [PATCH] collect_rent_eagerly can just do rewrites (#24438)

---
 runtime/src/bank.rs           | 29 ++++++++++++++++++++++-------
 runtime/src/snapshot_utils.rs |  2 ++
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index a2b5fc2f61..9853e30e11 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -3035,7 +3035,7 @@ impl Bank {
         let mut hash = self.hash.write().unwrap();
         if *hash == Hash::default() {
             // finish up any deferred changes to account state
-            self.collect_rent_eagerly();
+            self.collect_rent_eagerly(false);
             self.collect_fees();
             self.distribute_rent();
             self.update_slot_history();
@@ -4741,7 +4741,12 @@ impl Bank {
         }
     }
 
-    fn collect_rent_eagerly(&self) {
+    /// after deserialize, populate rewrites with accounts that would normally have had their data rewritten in this slot due to rent collection (but didn't)
+    pub fn prepare_rewrites_for_hash(&self) {
+        self.collect_rent_eagerly(true);
+    }
+
+    fn collect_rent_eagerly(&self, just_rewrites: bool) {
         if self.lazy_rent_collection.load(Relaxed) {
             return;
         }
@@ -4751,7 +4756,7 @@ impl Bank {
         let count = partitions.len();
         let account_count: usize = partitions
             .into_iter()
-            .map(|partition| self.collect_rent_in_partition(partition))
+            .map(|partition| self.collect_rent_in_partition(partition, just_rewrites))
             .sum();
         measure.stop();
         datapoint_info!(
@@ -4787,7 +4792,17 @@ impl Bank {
         }
     }
 
-    fn collect_rent_in_partition(&self, partition: Partition) -> usize {
+    /// load accounts with pubkeys in 'partition'
+    /// collect rent
+    /// store accounts, whether rent was collected or not
+    /// update bank's rewrites set for all rewrites that were skipped
+    /// if 'just_rewrites', function will only update bank's rewrites set and not actually store any accounts
+    /// return # accounts loaded
+    fn collect_rent_in_partition(&self, partition: Partition, just_rewrites: bool) -> usize {
+        if just_rewrites {
+            // this is not implemented yet. In the meantime, this call can have no side effects.
+            return 0;
+        }
         let subrange = Self::pubkey_range_from_partition(partition);
 
         let thread_pool = &self.rc.accounts.accounts_db.thread_pool;
@@ -8781,7 +8796,7 @@ pub(crate) mod tests {
             vec![genesis_slot]
         );
 
-        bank.collect_rent_in_partition((0, 0, 1)); // all range
+        bank.collect_rent_in_partition((0, 0, 1), false); // all range
 
         assert_eq!(bank.collected_rent.load(Relaxed), rent_collected);
         assert_eq!(
@@ -8864,8 +8879,8 @@ pub(crate) mod tests {
         assert_eq!(hash1_with_zero, hash1_without_zero);
         assert_ne!(hash1_with_zero, Hash::default());
 
-        bank2_with_zero.collect_rent_in_partition((0, 0, 1)); // all
-        bank2_without_zero.collect_rent_in_partition((0, 0, 1)); // all
+        bank2_with_zero.collect_rent_in_partition((0, 0, 1), false); // all
+        bank2_without_zero.collect_rent_in_partition((0, 0, 1), false); // all
 
         bank2_with_zero.freeze();
         let hash2_with_zero = bank2_with_zero.hash();
diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index c1fd8afdd0..864f9d9ffa 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -1632,6 +1632,8 @@ fn rebuild_bank_from_snapshots(
 
     bank.src.append(&slot_deltas);
 
+    bank.prepare_rewrites_for_hash();
+
     info!("Loaded bank for slot: {}", bank.slot());
     Ok(bank)
 }