fix epoch_stakes (#5355) (#5369)

automerge

core/src/staking_utils.rs

@@ -150,13 +150,23 @@ pub(crate) mod tests {
 
         // First epoch has the bootstrap leader
         expected.insert(voting_keypair.pubkey(), leader_stake.stake(0));
-        let expected = Some(expected);
+        assert_eq!(
-        assert_eq!(vote_account_stakes_at_epoch(&bank, 0), expected);
+            vote_account_stakes_at_epoch(&bank, 0),
+            Some(expected.clone())
+        );
 
         // Second epoch carries same information
         let bank = new_from_parent(&Arc::new(bank), 1);
-        assert_eq!(vote_account_stakes_at_epoch(&bank, 0), expected);
+        assert_eq!(
-        assert_eq!(vote_account_stakes_at_epoch(&bank, 1), expected);
+            vote_account_stakes_at_epoch(&bank, 0),
+            Some(expected.clone())
+        );
+
+        expected.insert(voting_keypair.pubkey(), leader_stake.stake(1));
+        assert_eq!(
+            vote_account_stakes_at_epoch(&bank, 1),
+            Some(expected.clone())
+        );
     }
 
     pub(crate) fn setup_vote_and_stake_accounts(
@@ -223,7 +233,7 @@ pub(crate) mod tests {
             ..
         } = create_genesis_block(10_000);
 
-        let bank = Bank::new(&genesis_block);
+        let mut bank = Bank::new(&genesis_block);
         let vote_pubkey = Pubkey::new_rand();
 
         // Give the validator some stake but don't setup a staking account
@@ -244,24 +254,21 @@ pub(crate) mod tests {
 
         let other_stake = Stake {
             stake,
+            activated: bank.get_stakers_epoch(bank.slot()),
             ..Stake::default()
         };
 
-        // soonest slot that could be a new epoch is 1
+        let epoch = bank.get_stakers_epoch(bank.slot());
-        let mut slot = 1;
+        // find the first slot in the next staker's epoch
-        let mut epoch = bank.get_stakers_epoch(0);
+        while bank.epoch() <= epoch {
-        // find the first slot in the next stakers_epoch
+            let slot = bank.slot() + 1;
-        while bank.get_stakers_epoch(slot) == epoch {
+            bank = new_from_parent(&Arc::new(bank), slot);
-            slot += 1;
         }
-        epoch = bank.get_stakers_epoch(slot);
-
-        let bank = new_from_parent(&Arc::new(bank), slot);
 
         let result: Vec<_> = epoch_stakes_and_lockouts(&bank, 0);
         assert_eq!(result, vec![(leader_stake.stake(0), None)]);
 
-        let mut result: Vec<_> = epoch_stakes_and_lockouts(&bank, epoch);
+        let mut result: Vec<_> = epoch_stakes_and_lockouts(&bank, bank.epoch());
         result.sort();
         let mut expected = vec![
             (leader_stake.stake(bank.epoch()), None),

programs/stake_api/src/stake_state.rs

@@ -53,7 +53,7 @@ pub struct Stake {
     pub activated: Epoch,   // epoch the stake was activated
     pub deactivated: Epoch, // epoch the stake was deactivated, std::Epoch::MAX if not deactivated
 }
-pub const STAKE_WARMUP_EPOCHS: u64 = 3;
+pub const STAKE_WARMUP_EPOCHS: Epoch = 3;
 
 impl Default for Stake {
     fn default() -> Self {
@@ -68,7 +68,7 @@ impl Default for Stake {
 }
 
 impl Stake {
-    pub fn stake(&self, epoch: u64) -> u64 {
+    pub fn stake(&self, epoch: Epoch) -> u64 {
         // before "activated" or after deactivated?
         if epoch < self.activated || epoch >= self.deactivated {
             return 0;
@@ -217,7 +217,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
                 new_stake,
                 vote_account.unsigned_key(),
                 &vote_account.state()?,
-                clock.epoch,
+                clock.stakers_epoch,
             );
 
             self.set_state(&StakeState::Stake(stake))
@@ -231,7 +231,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
         }
 
         if let StakeState::Stake(mut stake) = self.state()? {
-            stake.deactivate(clock.epoch);
+            stake.deactivate(clock.stakers_epoch);
 
             self.set_state(&StakeState::Stake(stake))
         } else {
@@ -287,7 +287,11 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
 
         match self.state()? {
             StakeState::Stake(mut stake) => {
-                let staked = if stake.stake(clock.epoch) == 0 {
+                // still activated, no can do
+                if stake.deactivated == std::u64::MAX {
+                    return Err(InstructionError::InsufficientFunds);
+                }
+                let staked = if stake.stake(clock.stakers_epoch) == 0 {
                     0
                 } else {
                     // Assume full stake if the stake is under warmup/cooldown
@@ -353,7 +357,10 @@ mod tests {
 
     #[test]
     fn test_stake_delegate_stake() {
-        let clock = sysvar::clock::Clock::default();
+        let clock = sysvar::clock::Clock {
+            stakers_epoch: 1,
+            ..sysvar::clock::Clock::default()
+        };
 
         let vote_keypair = Keypair::new();
         let mut vote_state = VoteState::default();
@@ -399,7 +406,7 @@ mod tests {
                 voter_pubkey: vote_keypair.pubkey(),
                 credits_observed: vote_state.credits(),
                 stake: stake_lamports,
-                activated: 0,
+                activated: clock.stakers_epoch,
                 deactivated: std::u64::MAX,
             })
         );
@@ -456,7 +463,10 @@ mod tests {
         let mut stake_account =
             Account::new(stake_lamports, std::mem::size_of::<StakeState>(), &id());
 
-        let clock = sysvar::clock::Clock::default();
+        let clock = sysvar::clock::Clock {
+            stakers_epoch: 1,
+            ..sysvar::clock::Clock::default()
+        };
 
         // unsigned keyed account
         let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, false, &mut stake_account);
@@ -495,7 +505,7 @@ mod tests {
         let mut stake_account =
             Account::new(total_lamports, std::mem::size_of::<StakeState>(), &id());
 
-        let clock = sysvar::clock::Clock::default();
+        let mut clock = sysvar::clock::Clock::default();
 
         let to = Pubkey::new_rand();
         let mut to_account = Account::new(1, 0, &system_program::id());
@@ -535,6 +545,11 @@ mod tests {
             Ok(())
         );
 
+        // deactivate the stake before withdrawal
+        assert_eq!(stake_keyed_account.deactivate_stake(&clock), Ok(()));
+        // simulate time passing
+        clock.stakers_epoch += STAKE_WARMUP_EPOCHS;
+
         // Try to withdraw more than what's available
         assert_eq!(
             stake_keyed_account.withdraw(
@@ -566,7 +581,7 @@ mod tests {
 
         let clock = sysvar::clock::Clock::default();
         let mut future = sysvar::clock::Clock::default();
-        future.epoch += 16;
+        future.stakers_epoch += 16;
 
         let to = Pubkey::new_rand();
         let mut to_account = Account::new(1, 0, &system_program::id());
@@ -585,14 +600,14 @@ mod tests {
             Ok(())
         );
 
-        // Try to withdraw including staked
+        // Try to withdraw stake
         assert_eq!(
             stake_keyed_account.withdraw(
                 total_lamports - stake_lamports + 1,
                 &mut to_keyed_account,
                 &clock
             ),
-            Ok(())
+            Err(InstructionError::InsufficientFunds)
         );
     }
 

runtime/src/bank.rs

@@ -286,8 +286,9 @@ impl Bank {
         //  slot = 0 and genesis configuration
         {
             let stakes = bank.stakes.read().unwrap();
-            for i in 0..=bank.get_stakers_epoch(bank.slot) {
+            for epoch in 0..=bank.get_stakers_epoch(bank.slot) {
-                bank.epoch_stakes.insert(i, stakes.clone());
+                bank.epoch_stakes
+                    .insert(epoch, stakes.clone_with_epoch(epoch));
             }
         }
         bank.update_clock();
@@ -357,7 +358,7 @@ impl Bank {
             //  if my parent didn't populate for this epoch, we've
             //  crossed a boundary
             if epoch_stakes.get(&epoch).is_none() {
-                epoch_stakes.insert(epoch, self.stakes.read().unwrap().clone());
+                epoch_stakes.insert(epoch, self.stakes.read().unwrap().clone_with_epoch(epoch));
             }
             epoch_stakes
         };
@@ -1496,6 +1497,7 @@ mod tests {
     use solana_sdk::system_transaction;
     use solana_sdk::sysvar::{fees::Fees, rewards::Rewards};
     use solana_sdk::timing::DEFAULT_TICKS_PER_SLOT;
+    use solana_stake_api::stake_state::Stake;
     use solana_vote_api::vote_instruction;
     use solana_vote_api::vote_state::{VoteState, MAX_LOCKOUT_HISTORY};
     use std::io::Cursor;
@@ -2435,33 +2437,48 @@ mod tests {
         genesis_block.epoch_warmup = false; // allows me to do the normal division stuff below
 
         let parent = Arc::new(Bank::new(&genesis_block));
-
+        let mut leader_vote_stake: Vec<_> = parent
-        let vote_accounts0: Option<HashMap<_, _>> = parent.epoch_vote_accounts(0).map(|accounts| {
+            .epoch_vote_accounts(0)
-            accounts
+            .map(|accounts| {
-                .iter()
+                accounts
-                .filter_map(|(pubkey, (_, account))| {
+                    .iter()
-                    if let Ok(vote_state) = VoteState::deserialize(&account.data) {
+                    .filter_map(|(pubkey, (stake, account))| {
-                        if vote_state.node_pubkey == leader_pubkey {
+                        if let Ok(vote_state) = VoteState::deserialize(&account.data) {
-                            Some((*pubkey, true))
+                            if vote_state.node_pubkey == leader_pubkey {
+                                Some((*pubkey, *stake))
+                            } else {
+                                None
+                            }
                         } else {
                             None
                         }
-                    } else {
+                    })
-                        None
+                    .collect()
-                    }
+            })
-                })
+            .unwrap();
-                .collect()
+        assert_eq!(leader_vote_stake.len(), 1);
-        });
+        let (leader_vote_account, leader_stake) = leader_vote_stake.pop().unwrap();
-        assert!(vote_accounts0.is_some());
+        assert!(leader_stake > 0);
-        assert!(vote_accounts0.iter().len() != 0);
 
-        let mut i = 1;
+        let leader_stake = Stake {
+            stake: leader_lamports,
+            ..Stake::default()
+        };
+
+        let mut epoch = 1;
         loop {
-            if i > STAKERS_SLOT_OFFSET / SLOTS_PER_EPOCH {
+            if epoch > STAKERS_SLOT_OFFSET / SLOTS_PER_EPOCH {
                 break;
             }
-            assert!(parent.epoch_vote_accounts(i).is_some());
+            let vote_accounts = parent.epoch_vote_accounts(epoch);
-            i += 1;
+            assert!(vote_accounts.is_some());
+
+            assert_eq!(
+                leader_stake.stake(epoch),
+                vote_accounts.unwrap().get(&leader_vote_account).unwrap().0
+            );
+
+            epoch += 1;
         }
 
         // child crosses epoch boundary and is the first slot in the epoch
@@ -2471,15 +2488,25 @@ mod tests {
             SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH),
         );
 
-        assert!(child.epoch_vote_accounts(i).is_some());
+        assert!(child.epoch_vote_accounts(epoch).is_some());
+        assert_eq!(
+            leader_stake.stake(epoch),
+            child
+                .epoch_vote_accounts(epoch)
+                .unwrap()
+                .get(&leader_vote_account)
+                .unwrap()
+                .0
+        );
 
-        // child crosses epoch boundary but isn't the first slot in the epoch
+        // child crosses epoch boundary but isn't the first slot in the epoch, still
+        //  makes an epoch stakes
         let child = Bank::new_from_parent(
             &parent,
             &leader_pubkey,
             SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH) + 1,
         );
-        assert!(child.epoch_vote_accounts(i).is_some());
+        assert!(child.epoch_vote_accounts(epoch).is_some());
     }
 
     #[test]

runtime/src/stakes.rs

@@ -2,6 +2,7 @@
 //! node stakes
 use solana_sdk::account::Account;
 use solana_sdk::pubkey::Pubkey;
+use solana_sdk::timing::Epoch;
 use solana_stake_api::stake_state::StakeState;
 use solana_vote_api::vote_state::VoteState;
 use std::collections::HashMap;
@@ -19,11 +20,11 @@ pub struct Stakes {
     points: u64,
 
     /// current epoch, used to calculate current stake
-    epoch: u64,
+    epoch: Epoch,
 }
 
 impl Stakes {
-    pub fn clone_with_epoch(&self, epoch: u64) -> Self {
+    pub fn clone_with_epoch(&self, epoch: Epoch) -> Self {
         if self.epoch == epoch {
             self.clone()
         } else {
@@ -46,7 +47,7 @@ impl Stakes {
     }
 
     // sum the stakes that point to the given voter_pubkey
-    fn calculate_stake(&self, voter_pubkey: &Pubkey, epoch: u64) -> u64 {
+    fn calculate_stake(&self, voter_pubkey: &Pubkey, epoch: Epoch) -> u64 {
         self.stake_accounts
             .iter()
             .map(|(_, stake_account)| {
@@ -62,7 +63,9 @@ impl Stakes {
     }
 
     pub fn is_stake(account: &Account) -> bool {
-        solana_vote_api::check_id(&account.owner) || solana_stake_api::check_id(&account.owner)
+        solana_vote_api::check_id(&account.owner)
+            || solana_stake_api::check_id(&account.owner)
+                && account.data.len() >= std::mem::size_of::<StakeState>()
     }
 
     pub fn store(&mut self, pubkey: &Pubkey, account: &Account) {