From 0e7b0597db3533d304a009d4276164ddb78049df Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <wash678@gmail.com>
Date: Thu, 14 Apr 2022 13:46:18 -0500
Subject: [PATCH] check for rewrites skipped in closure (#24330)

---
 runtime/src/expected_rent_collection.rs | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/runtime/src/expected_rent_collection.rs b/runtime/src/expected_rent_collection.rs
index 25f63c0749..d5d7c6e23a 100644
--- a/runtime/src/expected_rent_collection.rs
+++ b/runtime/src/expected_rent_collection.rs
@@ -280,14 +280,18 @@ impl ExpectedRentCollection {
                     possibly_update = false;
                 }
             } // if more than 1 epoch old, then we need to collect rent because we clearly skipped it.
+
+            let rewrites_skipped_this_pubkey_this_slot = || {
+                rewrites_skipped_this_slot
+                    .read()
+                    .unwrap()
+                    .contains_key(pubkey)
+            };
             let rent_epoch = account.rent_epoch();
             if possibly_update
                 && rent_epoch == 0
                 && current_epoch > 1
-                && !rewrites_skipped_this_slot
-                    .read()
-                    .unwrap()
-                    .contains_key(pubkey)
+                && !rewrites_skipped_this_pubkey_this_slot()
             {
                 // we know we're done
                 return None;
@@ -296,10 +300,8 @@ impl ExpectedRentCollection {
             // if an account was written >= its rent collection slot within the last epoch worth of slots, then we don't want to update it here
             if possibly_update && rent_epoch < current_epoch {
                 let new_rent_epoch = if partition_from_pubkey < partition_from_current_slot
-                    || rewrites_skipped_this_slot
-                        .read()
-                        .unwrap()
-                        .contains_key(pubkey)
+                    || (partition_from_pubkey == partition_from_current_slot
+                        && rewrites_skipped_this_pubkey_this_slot())
                 {
                     // partition_from_pubkey < partition_from_current_slot:
                     //  we already would have done a rewrite on this account IN this epoch
@@ -316,7 +318,7 @@ impl ExpectedRentCollection {
                 }
             } else if !possibly_update {
                 // This is a non-trivial lookup. Would be nice to skip this.
-                assert!(!rewrites_skipped_this_slot.read().unwrap().contains_key(pubkey), "did not update rent_epoch: {}, new value for rent_epoch: {}, old: {}, current epoch: {}", pubkey, rent_epoch, next_epoch, current_epoch);
+                assert!(!rewrites_skipped_this_pubkey_this_slot(), "did not update rent_epoch: {}, new value for rent_epoch: {}, old: {}, current epoch: {}", pubkey, rent_epoch, next_epoch, current_epoch);
             }
         }
         None