Rewrite stake accounts for clear migration (#13461)

* Reduce overage stake by rewritng stake accounts

* Write tests and finish implemention

* Create and use new feature gate

* Clean up logging

* Fix typo

* Simplify enable_rewrite_stake

* Fix typo...

* Even simplify gating

* Add metrics

ledger-tool/src/main.rs

@@ -2052,7 +2052,15 @@ fn main() {
                                     feature_account_balance,
                                 ),
                             );
+                            base_bank.store_account(
+                                &feature_set::rewrite_stake::id(),
+                                &feature::create_account(
+                                    &Feature { activated_at: None },
+                                    feature_account_balance,
+                                ),
+                            );
 
+                            let mut store_failed_count = 0;
                             if base_bank
                                 .get_account(&feature_set::secp256k1_program_enabled::id())
                                 .is_some()
@@ -2064,6 +2072,21 @@ fn main() {
                                     &Account::default(),
                                 );
                             } else {
+                                store_failed_count += 1;
+                            }
+
+                            if base_bank
+                                .get_account(&feature_set::instructions_sysvar_enabled::id())
+                                .is_some()
+                            {
+                                base_bank.store_account(
+                                    &feature_set::instructions_sysvar_enabled::id(),
+                                    &Account::default(),
+                                );
+                            } else {
+                                store_failed_count += 1;
+                            }
+                            if store_failed_count >= 1 {
                                 // we have no choice; maybe locally created blank cluster with
                                 // not-Development cluster type.
                                 let old_cap = base_bank.set_capitalization();
@@ -2073,7 +2096,10 @@ fn main() {
                                      requested: increasing {} from {} to {}",
                                     feature_account_balance, old_cap, new_cap,
                                 );
-                                assert_eq!(old_cap + feature_account_balance, new_cap);
+                                assert_eq!(
+                                    old_cap + feature_account_balance * store_failed_count,
+                                    new_cap
+                                );
                             }
                         }
 

programs/stake/src/stake_state.rs

@@ -166,6 +166,24 @@ impl Meta {
         }
         Ok(())
     }
+
+    pub fn rewrite_rent_exempt_reserve(
+        &mut self,
+        rent: &Rent,
+        data_len: usize,
+    ) -> Option<(u64, u64)> {
+        let corrected_rent_exempt_reserve = rent.minimum_balance(data_len);
+        if corrected_rent_exempt_reserve != self.rent_exempt_reserve {
+            // We forcibly update rent_excempt_reserve even
+            // if rent_exempt_reserve > account_balance, hoping user might restore
+            // rent_exempt status by depositing.
+            let (old, new) = (self.rent_exempt_reserve, corrected_rent_exempt_reserve);
+            self.rent_exempt_reserve = corrected_rent_exempt_reserve;
+            Some((old, new))
+        } else {
+            None
+        }
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Clone, Copy, AbiExample)]
@@ -392,6 +410,27 @@ impl Delegation {
             (self.stake, 0)
         }
     }
+
+    fn rewrite_stake(
+        &mut self,
+        account_balance: u64,
+        rent_exempt_balance: u64,
+    ) -> Option<(u64, u64)> {
+        // note that this will intentionally overwrite innocent
+        // deactivated-then-immeditealy-withdrawn stake accounts as well
+        // this is chosen to minimize the risks from complicated logic,
+        // over some unneeded rewrites
+        let corrected_stake = account_balance.saturating_sub(rent_exempt_balance);
+        if self.stake != corrected_stake {
+            // this could result in creating a 0-staked account;
+            // rewards and staking calc can handle it.
+            let (old, new) = (self.stake, corrected_stake);
+            self.stake = corrected_stake;
+            Some((old, new))
+        } else {
+            None
+        }
+    }
 }
 
 #[derive(Debug, Default, Serialize, Deserialize, PartialEq, Clone, Copy, AbiExample)]
@@ -1186,6 +1225,44 @@ fn calculate_split_rent_exempt_reserve(
     lamports_per_byte_year * (split_data_len + ACCOUNT_STORAGE_OVERHEAD)
 }
 
+pub type RewriteStakeStatus = (&'static str, (u64, u64), (u64, u64));
+
+pub fn rewrite_stakes(
+    stake_account: &mut Account,
+    rent: &Rent,
+) -> Result<RewriteStakeStatus, InstructionError> {
+    match stake_account.state()? {
+        StakeState::Initialized(mut meta) => {
+            let meta_status = meta.rewrite_rent_exempt_reserve(rent, stake_account.data.len());
+
+            if meta_status.is_none() {
+                return Err(InstructionError::InvalidAccountData);
+            }
+
+            stake_account.set_state(&StakeState::Initialized(meta))?;
+            Ok(("initialized", meta_status.unwrap_or_default(), (0, 0)))
+        }
+        StakeState::Stake(mut meta, mut stake) => {
+            let meta_status = meta.rewrite_rent_exempt_reserve(rent, stake_account.data.len());
+            let stake_status = stake
+                .delegation
+                .rewrite_stake(stake_account.lamports, meta.rent_exempt_reserve);
+
+            if meta_status.is_none() && stake_status.is_none() {
+                return Err(InstructionError::InvalidAccountData);
+            }
+
+            stake_account.set_state(&StakeState::Stake(meta, stake))?;
+            Ok((
+                "stake",
+                meta_status.unwrap_or_default(),
+                stake_status.unwrap_or_default(),
+            ))
+        }
+        _ => Err(InstructionError::InvalidAccountData),
+    }
+}
+
 // utility function, used by runtime::Stakes, tests
 pub fn new_stake_history_entry<'a, I>(
     epoch: Epoch,
@@ -4868,6 +4945,133 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_meta_rewrite_rent_exempt_reserve() {
+        let right_data_len = std::mem::size_of::<StakeState>() as u64;
+        let rent = Rent::default();
+        let expected_rent_exempt_reserve = rent.minimum_balance(right_data_len as usize);
+
+        let test_cases = [
+            (
+                right_data_len + 100,
+                Some((
+                    rent.minimum_balance(right_data_len as usize + 100),
+                    expected_rent_exempt_reserve,
+                )),
+            ), // large data_len, too small rent exempt
+            (right_data_len, None), // correct
+            (
+                right_data_len - 100,
+                Some((
+                    rent.minimum_balance(right_data_len as usize - 100),
+                    expected_rent_exempt_reserve,
+                )),
+            ), // small data_len, too large rent exempt
+        ];
+        for (data_len, expected_rewrite) in &test_cases {
+            let rent_exempt_reserve = rent.minimum_balance(*data_len as usize);
+            let mut meta = Meta {
+                rent_exempt_reserve,
+                ..Meta::default()
+            };
+            let actual_rewrite = meta.rewrite_rent_exempt_reserve(&rent, right_data_len as usize);
+            assert_eq!(actual_rewrite, *expected_rewrite);
+            assert_eq!(meta.rent_exempt_reserve, expected_rent_exempt_reserve);
+        }
+    }
+
+    #[test]
+    fn test_stake_rewrite_stake() {
+        let right_data_len = std::mem::size_of::<StakeState>() as u64;
+        let rent = Rent::default();
+        let rent_exempt_reserve = rent.minimum_balance(right_data_len as usize);
+        let expected_stake = 1000;
+        let account_balance = rent_exempt_reserve + expected_stake;
+
+        let test_cases = [
+            (9999, Some((9999, expected_stake))), // large stake
+            (1000, None),                         // correct
+            (42, Some((42, expected_stake))),     // small stake
+        ];
+        for (staked_amount, expected_rewrite) in &test_cases {
+            let mut delegation = Delegation {
+                stake: *staked_amount,
+                ..Delegation::default()
+            };
+            let actual_rewrite = delegation.rewrite_stake(account_balance, rent_exempt_reserve);
+            assert_eq!(actual_rewrite, *expected_rewrite);
+            assert_eq!(delegation.stake, expected_stake);
+        }
+    }
+
+    enum ExpectedRewriteResult {
+        NotRewritten,
+        Rewritten,
+    }
+
+    #[test]
+    fn test_rewrite_stakes_initialized() {
+        let right_data_len = std::mem::size_of::<StakeState>();
+        let rent = Rent::default();
+        let rent_exempt_reserve = rent.minimum_balance(right_data_len as usize);
+        let expected_stake = 1000;
+        let account_balance = rent_exempt_reserve + expected_stake;
+
+        let test_cases = [
+            (1, ExpectedRewriteResult::Rewritten),
+            (0, ExpectedRewriteResult::NotRewritten),
+        ];
+        for (offset, expected_rewrite) in &test_cases {
+            let meta = Meta {
+                rent_exempt_reserve: rent_exempt_reserve + offset,
+                ..Meta::default()
+            };
+            let mut account = Account::new(account_balance, right_data_len, &id());
+            account.set_state(&StakeState::Initialized(meta)).unwrap();
+            let result = rewrite_stakes(&mut account, &rent);
+            match expected_rewrite {
+                ExpectedRewriteResult::NotRewritten => assert!(result.is_err()),
+                ExpectedRewriteResult::Rewritten => assert!(result.is_ok()),
+            }
+        }
+    }
+
+    #[test]
+    fn test_rewrite_stakes_stake() {
+        let right_data_len = std::mem::size_of::<StakeState>();
+        let rent = Rent::default();
+        let rent_exempt_reserve = rent.minimum_balance(right_data_len as usize);
+        let expected_stake = 1000;
+        let account_balance = rent_exempt_reserve + expected_stake;
+
+        let test_cases = [
+            (1, 9999, ExpectedRewriteResult::Rewritten), // bad meta, bad stake
+            (1, 1000, ExpectedRewriteResult::Rewritten), // bad meta, good stake
+            (0, 9999, ExpectedRewriteResult::Rewritten), // good meta, bad stake
+            (0, 1000, ExpectedRewriteResult::NotRewritten), // good meta, good stake
+        ];
+        for (offset, staked_amount, expected_rewrite) in &test_cases {
+            let meta = Meta {
+                rent_exempt_reserve: rent_exempt_reserve + offset,
+                ..Meta::default()
+            };
+            let stake = Stake {
+                delegation: (Delegation {
+                    stake: *staked_amount,
+                    ..Delegation::default()
+                }),
+                ..Stake::default()
+            };
+            let mut account = Account::new(account_balance, right_data_len, &id());
+            account.set_state(&StakeState::Stake(meta, stake)).unwrap();
+            let result = rewrite_stakes(&mut account, &rent);
+            match expected_rewrite {
+                ExpectedRewriteResult::NotRewritten => assert!(result.is_err()),
+                ExpectedRewriteResult::Rewritten => assert!(result.is_ok()),
+            }
+        }
+    }
+
     #[test]
     fn test_calculate_lamports_per_byte_year() {
         let rent = Rent::default();

runtime/src/bank.rs

@@ -701,6 +701,8 @@ pub struct Bank {
 
     pub lazy_rent_collection: AtomicBool,
 
+    pub no_stake_rewrite: AtomicBool,
+
     // this is temporary field only to remove rewards_pool entirely
     pub rewards_pool_pubkeys: Arc<HashSet<Pubkey>>,
 
@@ -828,6 +830,7 @@ impl Bank {
             capitalization: AtomicU64::new(parent.capitalization()),
             inflation: parent.inflation.clone(),
             transaction_count: AtomicU64::new(parent.transaction_count()),
+            // we will .clone_with_epoch() this soon after stake data update; so just .clone() for now
             stakes: RwLock::new(parent.stakes.read().unwrap().clone()),
             epoch_stakes: parent.epoch_stakes.clone(),
             parent_hash: parent.hash(),
@@ -848,6 +851,7 @@ impl Bank {
             skip_drop: AtomicBool::new(false),
             cluster_type: parent.cluster_type,
             lazy_rent_collection: AtomicBool::new(parent.lazy_rent_collection.load(Relaxed)),
+            no_stake_rewrite: AtomicBool::new(parent.no_stake_rewrite.load(Relaxed)),
             rewards_pool_pubkeys: parent.rewards_pool_pubkeys.clone(),
             cached_executors: RwLock::new((*parent.cached_executors.read().unwrap()).clone()),
             transaction_debug_keys: parent.transaction_debug_keys.clone(),
@@ -866,9 +870,11 @@ impl Bank {
         });
 
         // Following code may touch AccountsDB, requiring proper ancestors
-        if parent.epoch() < new.epoch() {
+        let parent_epoch = parent.epoch();
+        if parent_epoch < new.epoch() {
             new.apply_feature_activations(false);
         }
+
         let cloned = new
             .stakes
             .read()
@@ -879,9 +885,9 @@ impl Bank {
         let leader_schedule_epoch = epoch_schedule.get_leader_schedule_epoch(slot);
         new.update_epoch_stakes(leader_schedule_epoch);
         new.update_slot_hashes();
-        new.update_rewards(parent.epoch(), reward_calc_tracer);
-        new.update_stake_history(Some(parent.epoch()));
-        new.update_clock(Some(parent.epoch()));
+        new.update_rewards(parent_epoch, reward_calc_tracer);
+        new.update_stake_history(Some(parent_epoch));
+        new.update_clock(Some(parent_epoch));
         new.update_fees();
         if !new.fix_recent_blockhashes_sysvar_delay() {
             new.update_recent_blockhashes();
@@ -959,6 +965,7 @@ impl Bank {
             skip_drop: new(),
             cluster_type: Some(genesis_config.cluster_type),
             lazy_rent_collection: new(),
+            no_stake_rewrite: new(),
             rewards_pool_pubkeys: new(),
             cached_executors: RwLock::new(CowCachedExecutors::new(Arc::new(RwLock::new(
                 CachedExecutors::new(MAX_CACHED_EXECUTORS),
@@ -1293,6 +1300,41 @@ impl Bank {
         self.epoch_schedule.get_slots_in_epoch(prev_epoch) as f64 / self.slots_per_year
     }
 
+    fn rewrite_stakes(&self) -> (usize, usize) {
+        let mut examined_count = 0;
+        let mut rewritten_count = 0;
+        self.cloned_stake_delegations()
+            .into_iter()
+            .for_each(|(stake_pubkey, _delegation)| {
+                examined_count += 1;
+                if let Some(mut stake_account) = self.get_account(&stake_pubkey) {
+                    if let Ok(result) =
+                        stake_state::rewrite_stakes(&mut stake_account, &self.rent_collector.rent)
+                    {
+                        self.store_account(&stake_pubkey, &stake_account);
+                        let message = format!("rewrote stake: {}, {:?}", stake_pubkey, result);
+                        info!("{}", message);
+                        datapoint_info!("stake_info", ("info", message, String));
+                        rewritten_count += 1;
+                    }
+                }
+            });
+
+        info!(
+            "bank (slot: {}): rewrite_stakes: {} accounts rewritten / {} accounts examined",
+            self.slot(),
+            rewritten_count,
+            examined_count,
+        );
+        datapoint_info!(
+            "rewrite-stakes",
+            ("examined_count", examined_count, i64),
+            ("rewritten_count", rewritten_count, i64)
+        );
+
+        (examined_count, rewritten_count)
+    }
+
     // update rewards based on the previous epoch
     fn update_rewards(
         &mut self,
@@ -2963,6 +3005,12 @@ impl Bank {
         inc_new_counter_info!("collect_rent_eagerly-ms", measure.as_ms() as usize);
     }
 
+    #[cfg(test)]
+    fn restore_old_behavior_for_fragile_tests(&self) {
+        self.lazy_rent_collection.store(true, Relaxed);
+        self.no_stake_rewrite.store(true, Relaxed);
+    }
+
     fn enable_eager_rent_collection(&self) -> bool {
         if self.lazy_rent_collection.load(Relaxed) {
             return false;
@@ -3907,8 +3955,7 @@ impl Bank {
     }
 
     /// current stake delegations for this bank
-    /// Note: this method is exposed publicly for external usage
-    pub fn stake_delegations(&self) -> HashMap<Pubkey, Delegation> {
+    pub fn cloned_stake_delegations(&self) -> HashMap<Pubkey, Delegation> {
         self.stakes.read().unwrap().stake_delegations().clone()
     }
 
@@ -4104,6 +4151,16 @@ impl Bank {
         if new_feature_activations.contains(&feature_set::spl_token_v2_multisig_fix::id()) {
             self.apply_spl_token_v2_multisig_fix();
         }
+        // Remove me after a while around v1.6
+        if !self.no_stake_rewrite.load(Relaxed)
+            && new_feature_activations.contains(&feature_set::rewrite_stake::id())
+        {
+            // to avoid any potential risk of wrongly rewriting accounts in the future,
+            // only do this once, taking small risk of unknown
+            // bugs which again creates bad stake accounts..
+
+            self.rewrite_stakes();
+        }
 
         self.ensure_feature_builtins(init_finish_or_warp, &new_feature_activations);
         self.reconfigure_token2_native_mint();
@@ -5113,7 +5170,7 @@ pub(crate) mod tests {
         let root_bank = Bank::new(&genesis_config);
         // until we completely transition to the eager rent collection,
         // we must ensure lazy rent collection doens't get broken!
-        root_bank.lazy_rent_collection.store(true, Relaxed);
+        root_bank.restore_old_behavior_for_fragile_tests();
         let root_bank = Arc::new(root_bank);
         let bank = create_child_bank_for_rent_test(&root_bank, &genesis_config, mock_program_id);
 
@@ -6063,7 +6120,7 @@ pub(crate) mod tests {
 
         // enable lazy rent collection because this test depends on rent-due accounts
         // not being eagerly-collected for exact rewards calculation
-        bank.lazy_rent_collection.store(true, Relaxed);
+        bank.restore_old_behavior_for_fragile_tests();
 
         assert_eq!(bank.capitalization(), 42 * 1_000_000_000);
         assert!(bank.rewards.read().unwrap().is_empty());
@@ -6182,7 +6239,7 @@ pub(crate) mod tests {
 
         // enable lazy rent collection because this test depends on rent-due accounts
         // not being eagerly-collected for exact rewards calculation
-        bank.lazy_rent_collection.store(true, Relaxed);
+        bank.restore_old_behavior_for_fragile_tests();
 
         assert_eq!(bank.capitalization(), 42 * 1_000_000_000);
         assert!(bank.rewards.read().unwrap().is_empty());
@@ -7750,7 +7807,7 @@ pub(crate) mod tests {
     }
 
     #[test]
-    fn test_bank_stake_delegations() {
+    fn test_bank_cloned_stake_delegations() {
         let GenesisConfigInfo {
             genesis_config,
             mint_keypair,
@@ -7758,7 +7815,7 @@ pub(crate) mod tests {
         } = create_genesis_config_with_leader(500, &solana_sdk::pubkey::new_rand(), 1);
         let bank = Arc::new(Bank::new(&genesis_config));
 
-        let stake_delegations = bank.stake_delegations();
+        let stake_delegations = bank.cloned_stake_delegations();
         assert_eq!(stake_delegations.len(), 1); // bootstrap validator has
                                                 // to have a stake delegation
 
@@ -7794,7 +7851,7 @@ pub(crate) mod tests {
 
         bank.process_transaction(&transaction).unwrap();
 
-        let stake_delegations = bank.stake_delegations();
+        let stake_delegations = bank.cloned_stake_delegations();
         assert_eq!(stake_delegations.len(), 2);
         assert!(stake_delegations.get(&stake_keypair.pubkey()).is_some());
     }
@@ -7853,7 +7910,7 @@ pub(crate) mod tests {
     fn test_bank_get_program_accounts() {
         let (genesis_config, mint_keypair) = create_genesis_config(500);
         let parent = Arc::new(Bank::new(&genesis_config));
-        parent.lazy_rent_collection.store(true, Relaxed);
+        parent.restore_old_behavior_for_fragile_tests();
 
         let genesis_accounts: Vec<_> = parent.get_all_accounts_with_modified_slots();
         assert!(
@@ -9247,7 +9304,7 @@ pub(crate) mod tests {
         let pubkey2 = solana_sdk::pubkey::new_rand();
 
         let mut bank = Arc::new(Bank::new(&genesis_config));
-        bank.lazy_rent_collection.store(true, Relaxed);
+        bank.restore_old_behavior_for_fragile_tests();
         assert_eq!(bank.process_stale_slot_with_budget(0, 0), 0);
         assert_eq!(bank.process_stale_slot_with_budget(133, 0), 133);
 
@@ -10297,7 +10354,7 @@ pub(crate) mod tests {
         // Make sure rent collection doesn't overwrite `large_account_pubkey`, which
         // keeps slot 1 alive in the accounts database. Otherwise, slot 1 and it's bank
         // hash would be removed from accounts, preventing `rehash()` from succeeding
-        bank1.lazy_rent_collection.store(true, Relaxed);
+        bank1.restore_old_behavior_for_fragile_tests();
         bank1.freeze();
         let bank1_hash = bank1.hash();
 
@@ -10625,4 +10682,24 @@ pub(crate) mod tests {
             },
         );
     }
+
+    #[test]
+    fn test_stake_rewrite() {
+        let GenesisConfigInfo { genesis_config, .. } =
+            create_genesis_config_with_leader(500, &solana_sdk::pubkey::new_rand(), 1);
+        let bank = Arc::new(Bank::new(&genesis_config));
+
+        // quickest way of creting bad stake account
+        let bootstrap_stake_pubkey = bank
+            .cloned_stake_delegations()
+            .keys()
+            .next()
+            .copied()
+            .unwrap();
+        let mut bootstrap_stake_account = bank.get_account(&bootstrap_stake_pubkey).unwrap();
+        bootstrap_stake_account.lamports = 10000000;
+        bank.store_account(&bootstrap_stake_pubkey, &bootstrap_stake_account);
+
+        assert_eq!(bank.rewrite_stakes(), (1, 1));
+    }
 }

sdk/src/feature_set.rs

@@ -86,6 +86,10 @@ pub mod stake_program_v2 {
     solana_sdk::declare_id!("Gvd9gGJZDHGMNf1b3jkxrfBQSR5etrfTQSBNKCvLSFJN");
 }
 
+pub mod rewrite_stake {
+    solana_sdk::declare_id!("6ap2eGy7wx5JmsWUmQ5sHwEWrFSDUxSti2k5Hbfv5BZG");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -109,6 +113,7 @@ lazy_static! {
         (pull_request_ping_pong_check::id(), "ping-pong packet check #12794"),
         (timestamp_bounding::id(), "add timestamp-correction bounding #13120"),
         (stake_program_v2::id(), "solana_stake_program v2"),
+        (rewrite_stake::id(), "rewrite stake"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()