Handle accounts data size changes due to rent-collected accounts (#22412) (#22919)

Co-authored-by: Brooks Prumo <brooks@solana.com>

runtime/src/accounts.rs

@@ -308,12 +308,14 @@ impl Accounts {
                             .load_with_fixed_root(ancestors, key)
                             .map(|(mut account, _)| {
                                 if message.is_writable(i) {
-                                    let rent_due = rent_collector.collect_from_existing_account(
+                                    let rent_due = rent_collector
-                                        key,
+                                        .collect_from_existing_account(
-                                        &mut account,
+                                            key,
-                                        rent_for_sysvars,
+                                            &mut account,
-                                        self.accounts_db.filler_account_suffix.as_ref(),
+                                            rent_for_sysvars,
-                                    );
+                                            self.accounts_db.filler_account_suffix.as_ref(),
+                                        )
+                                        .rent_amount;
                                     (account, rent_due)
                                 } else {
                                     (account, 0)
@@ -1173,11 +1175,9 @@ impl Accounts {
 
                     if execution_status.is_ok() || is_nonce_account || is_fee_payer {
                         if account.rent_epoch() == INITIAL_RENT_EPOCH {
-                            let rent = rent_collector.collect_from_created_account(
+                            let rent = rent_collector
-                                address,
+                                .collect_from_created_account(address, account, rent_for_sysvars)
-                                account,
+                                .rent_amount;
-                                rent_for_sysvars,
-                            );
                             loaded_transaction.rent += rent;
                             loaded_transaction.rent_debits.insert(
                                 address,

runtime/src/bank.rs

@@ -54,7 +54,7 @@ use {
         epoch_stakes::{EpochStakes, NodeVoteAccounts},
         inline_spl_associated_token_account, inline_spl_token,
         message_processor::MessageProcessor,
-        rent_collector::RentCollector,
+        rent_collector::{CollectedInfo, RentCollector},
         stake_weighted_timestamp::{
             calculate_stake_weighted_timestamp, MaxAllowableDrift, MAX_ALLOWABLE_DRIFT_PERCENTAGE,
             MAX_ALLOWABLE_DRIFT_PERCENTAGE_FAST, MAX_ALLOWABLE_DRIFT_PERCENTAGE_SLOW,
@@ -153,7 +153,7 @@ use {
         sync::{
             atomic::{
                 AtomicBool, AtomicU64,
-                Ordering::{Acquire, Relaxed, Release},
+                Ordering::{AcqRel, Acquire, Relaxed, Release},
             },
             Arc, LockResult, RwLock, RwLockReadGuard, RwLockWriteGuard,
         },
@@ -4210,6 +4210,27 @@ impl Bank {
         self.accounts_data_len.store(accounts_data_len, Release)
     }
 
+    /// Update the accounts data len by adding `delta`.  Since `delta` is signed, negative values
+    /// are allowed as the means to subtract from `accounts_data_len`.
+    fn update_accounts_data_len(&self, delta: i64) {
+        /// Mixed integer ops currently not stable, so copying the impl.
+        /// Copied from: https://github.com/a1phyr/rust/blob/47edde1086412b36e9efd6098b191ec15a2a760a/library/core/src/num/uint_macros.rs#L1039-L1048
+        fn saturating_add_signed(lhs: u64, rhs: i64) -> u64 {
+            let (res, overflow) = lhs.overflowing_add(rhs as u64);
+            if overflow == (rhs < 0) {
+                res
+            } else if overflow {
+                u64::MAX
+            } else {
+                u64::MIN
+            }
+        }
+        self.accounts_data_len
+            .fetch_update(AcqRel, Acquire, |x| Some(saturating_add_signed(x, delta)))
+            // SAFETY: unwrap() is safe here since our update fn always returns `Some`
+            .unwrap();
+    }
+
     /// Calculate fee for `SanitizedMessage`
     pub fn calculate_fee(message: &SanitizedMessage, lamports_per_signature: u64) -> u64 {
         let mut num_signatures = u64::from(message.header().num_required_signatures);
@@ -4601,29 +4622,33 @@ impl Bank {
 
         // parallelize?
         let rent_for_sysvars = self.rent_for_sysvars();
-        let mut total_rent = 0;
         let mut rent_debits = RentDebits::default();
+        let mut total_collected = CollectedInfo::default();
         for (pubkey, mut account) in accounts {
-            let rent = self.rent_collector.collect_from_existing_account(
+            let collected = self.rent_collector.collect_from_existing_account(
                 &pubkey,
                 &mut account,
                 rent_for_sysvars,
                 self.rc.accounts.accounts_db.filler_account_suffix.as_ref(),
             );
-            total_rent += rent;
+            total_collected += collected;
             // Store all of them unconditionally to purge old AppendVec,
             // even if collected rent is 0 (= not updated).
             // Also, there's another subtle side-effect from this: this
             // ensures we verify the whole on-chain state (= all accounts)
             // via the account delta hash slowly once per an epoch.
             self.store_account(&pubkey, &account);
-            rent_debits.insert(&pubkey, rent, account.lamports());
+            rent_debits.insert(&pubkey, collected.rent_amount, account.lamports());
         }
-        self.collected_rent.fetch_add(total_rent, Relaxed);
+        self.collected_rent
+            .fetch_add(total_collected.rent_amount, Relaxed);
         self.rewards
             .write()
             .unwrap()
             .extend(rent_debits.into_unordered_rewards_iter());
+        if total_collected.account_data_len_reclaimed > 0 {
+            self.update_accounts_data_len(-(total_collected.account_data_len_reclaimed as i64));
+        }
 
         self.rc.accounts.hold_range_in_memory(&subrange, false);
         account_count
@@ -7657,10 +7682,12 @@ pub(crate) mod tests {
 
         let account_pubkey = solana_sdk::pubkey::new_rand();
         let account_balance = 1;
+        let data_len = 12345; // use non-zero data len to also test accounts_data_len
         let mut account =
-            AccountSharedData::new(account_balance, 0, &solana_sdk::pubkey::new_rand());
+            AccountSharedData::new(account_balance, data_len, &solana_sdk::pubkey::new_rand());
         account.set_executable(true);
         bank.store_account(&account_pubkey, &account);
+        bank.store_accounts_data_len(data_len as u64);
 
         let transfer_lamports = 1;
         let tx = system_transaction::transfer(
@@ -7675,6 +7702,7 @@ pub(crate) mod tests {
             Err(TransactionError::InvalidWritableAccount)
         );
         assert_eq!(bank.get_balance(&account_pubkey), account_balance);
+        assert_eq!(bank.load_accounts_data_len(), data_len as u64);
     }
 
     #[test]
@@ -8622,9 +8650,10 @@ pub(crate) mod tests {
         let genesis_bank2 = Arc::new(Bank::new_for_tests(&genesis_config));
         let bank1_with_zero = Arc::new(new_from_parent(&genesis_bank1));
         let bank1_without_zero = Arc::new(new_from_parent(&genesis_bank2));
-        let zero_lamports = 0;
 
-        let account = AccountSharedData::new(zero_lamports, 0, &Pubkey::default());
+        let zero_lamports = 0;
+        let data_len = 12345; // use non-zero data len to also test accounts_data_len
+        let account = AccountSharedData::new(zero_lamports, data_len, &Pubkey::default());
         bank1_with_zero.store_account(&zero_lamport_pubkey, &account);
         bank1_without_zero.store_account(&zero_lamport_pubkey, &account);
 
@@ -16065,4 +16094,40 @@ pub(crate) mod tests {
             num_accounts,
         );
     }
+
+    #[test]
+    fn test_update_accounts_data_len() {
+        let (genesis_config, _mint_keypair) = create_genesis_config(100);
+        let bank = Bank::new_for_tests(&genesis_config);
+
+        // Test: Subtraction saturates at 0
+        {
+            let data_len = 567_i64;
+            bank.store_accounts_data_len(data_len as u64);
+            bank.update_accounts_data_len(-(data_len + 1));
+            assert_eq!(bank.load_accounts_data_len(), 0);
+        }
+
+        // Test: Addition saturates at u64::MAX
+        {
+            let data_len_remaining = 567;
+            bank.store_accounts_data_len(u64::MAX - data_len_remaining);
+            bank.update_accounts_data_len((data_len_remaining + 1) as i64);
+            assert_eq!(bank.load_accounts_data_len(), u64::MAX);
+        }
+
+        // Test: Updates work as expected
+        {
+            // Set the accounts data len to be in the middle, then perform a bunch of small
+            // updates, checking the results after each one.
+            bank.store_accounts_data_len(u32::MAX as u64);
+            let mut rng = rand::thread_rng();
+            for _ in 0..100 {
+                let initial = bank.load_accounts_data_len() as i64;
+                let delta = rng.gen_range(-500, 500);
+                bank.update_accounts_data_len(delta);
+                assert_eq!(bank.load_accounts_data_len() as i64, initial + delta);
+            }
+        }
+    }
 }

runtime/src/rent_collector.rs

@@ -92,16 +92,16 @@ impl RentCollector {
         account: &mut AccountSharedData,
         rent_for_sysvars: bool,
         filler_account_suffix: Option<&Pubkey>,
-    ) -> u64 {
+    ) -> CollectedInfo {
         if self.can_skip_rent_collection(address, account, rent_for_sysvars, filler_account_suffix)
         {
-            return 0;
+            return CollectedInfo::default();
         }
 
         let rent_due = self.get_rent_due(account);
         if let RentDue::Paying(0) = rent_due {
             // maybe collect rent later, leave account alone
-            return 0;
+            return CollectedInfo::default();
         }
 
         let epoch_increment = match rent_due {
@@ -117,11 +117,16 @@ impl RentCollector {
         account.saturating_sub_lamports(rent_due.lamports());
         let end_lamports = account.lamports();
 
+        let mut account_data_len_reclaimed = 0;
         if end_lamports == 0 {
+            account_data_len_reclaimed = account.data().len() as u64;
             *account = AccountSharedData::default();
         }
 
-        begin_lamports - end_lamports
+        CollectedInfo {
+            rent_amount: begin_lamports - end_lamports,
+            account_data_len_reclaimed,
+        }
     }
 
     #[must_use = "add to Bank::collected_rent"]
@@ -130,7 +135,7 @@ impl RentCollector {
         address: &Pubkey,
         account: &mut AccountSharedData,
         rent_for_sysvars: bool,
-    ) -> u64 {
+    ) -> CollectedInfo {
         // initialize rent_epoch as created at this epoch
         account.set_rent_epoch(self.epoch);
         self.collect_from_existing_account(address, account, rent_for_sysvars, None)
@@ -153,6 +158,32 @@ impl RentCollector {
     }
 }
 
+/// Information computed during rent collection
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
+pub struct CollectedInfo {
+    /// Amount of rent collected from account
+    pub rent_amount: u64,
+    /// Size of data reclaimed from account (happens when account's lamports go to zero)
+    pub account_data_len_reclaimed: u64,
+}
+
+impl std::ops::Add for CollectedInfo {
+    type Output = Self;
+    fn add(self, other: Self) -> Self {
+        Self {
+            rent_amount: self.rent_amount + other.rent_amount,
+            account_data_len_reclaimed: self.account_data_len_reclaimed
+                + other.account_data_len_reclaimed,
+        }
+    }
+}
+
+impl std::ops::AddAssign for CollectedInfo {
+    fn add_assign(&mut self, other: Self) {
+        *self = *self + other;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use {super::*, solana_sdk::account::Account};
@@ -182,8 +213,12 @@ mod tests {
             true,
         );
         assert!(created_account.lamports() < old_lamports);
-        assert_eq!(created_account.lamports() + collected, old_lamports);
+        assert_eq!(
+            created_account.lamports() + collected.rent_amount,
+            old_lamports
+        );
         assert_ne!(created_account.rent_epoch(), old_epoch);
+        assert_eq!(collected.account_data_len_reclaimed, 0);
 
         // collect rent on a already-existing account
         let collected = rent_collector.collect_from_existing_account(
@@ -193,8 +228,12 @@ mod tests {
             None,
         );
         assert!(existing_account.lamports() < old_lamports);
-        assert_eq!(existing_account.lamports() + collected, old_lamports);
+        assert_eq!(
+            existing_account.lamports() + collected.rent_amount,
+            old_lamports
+        );
         assert_ne!(existing_account.rent_epoch(), old_epoch);
+        assert_eq!(collected.account_data_len_reclaimed, 0);
 
         // newly created account should be collected for less rent; thus more remaining balance
         assert!(created_account.lamports() > existing_account.lamports());
@@ -207,7 +246,6 @@ mod tests {
         let epoch = 3;
         let huge_lamports = 123_456_789_012;
         let tiny_lamports = 789_012;
-        let mut collected;
         let pubkey = solana_sdk::pubkey::new_rand();
 
         account.set_lamports(huge_lamports);
@@ -217,17 +255,19 @@ mod tests {
         let rent_collector = RentCollector::default().clone_with_epoch(epoch);
 
         // first mark account as being collected while being rent-exempt
-        collected = rent_collector.collect_from_existing_account(&pubkey, &mut account, true, None);
+        let collected =
+            rent_collector.collect_from_existing_account(&pubkey, &mut account, true, None);
         assert_eq!(account.lamports(), huge_lamports);
-        assert_eq!(collected, 0);
+        assert_eq!(collected, CollectedInfo::default());
 
         // decrease the balance not to be rent-exempt
         account.set_lamports(tiny_lamports);
 
         // ... and trigger another rent collection on the same epoch and check that rent is working
-        collected = rent_collector.collect_from_existing_account(&pubkey, &mut account, true, None);
+        let collected =
-        assert_eq!(account.lamports(), tiny_lamports - collected);
+            rent_collector.collect_from_existing_account(&pubkey, &mut account, true, None);
-        assert_ne!(collected, 0);
+        assert_eq!(account.lamports(), tiny_lamports - collected.rent_amount);
+        assert_ne!(collected, CollectedInfo::default());
     }
 
     #[test]
@@ -248,12 +288,42 @@ mod tests {
         let collected =
             rent_collector.collect_from_existing_account(&pubkey, &mut account, false, None);
         assert_eq!(account.lamports(), tiny_lamports);
-        assert_eq!(collected, 0);
+        assert_eq!(collected, CollectedInfo::default());
 
         // new behavior: sysvars are NOT special-cased
         let collected =
             rent_collector.collect_from_existing_account(&pubkey, &mut account, true, None);
         assert_eq!(account.lamports(), 0);
-        assert_eq!(collected, 1);
+        assert_eq!(collected.rent_amount, 1);
+    }
+
+    /// Ensure that when an account is "rent collected" away, its data len is returned.
+    #[test]
+    fn test_collect_cleans_up_account() {
+        solana_logger::setup();
+        let account_lamports = 1; // must be *below* rent amount
+        let account_data_len = 567;
+        let account_rent_epoch = 11;
+        let mut account = AccountSharedData::from(Account {
+            lamports: account_lamports, // <-- must be below rent-exempt amount
+            data: vec![u8::default(); account_data_len],
+            rent_epoch: account_rent_epoch,
+            ..Account::default()
+        });
+        let rent_collector = RentCollector::default().clone_with_epoch(account_rent_epoch + 2);
+
+        let collected = rent_collector.collect_from_existing_account(
+            &Pubkey::new_unique(),
+            &mut account,
+            true,
+            None,
+        );
+
+        assert_eq!(collected.rent_amount, account_lamports);
+        assert_eq!(
+            collected.account_data_len_reclaimed,
+            account_data_len as u64
+        );
+        assert_eq!(account, AccountSharedData::default());
     }
 }