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

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();