use stake (#4721)

core/src/staking_utils.rs

@@ -121,9 +121,8 @@ pub(crate) mod tests {
     use solana_sdk::signature::{Keypair, KeypairUtil};
     use solana_sdk::transaction::Transaction;
     use solana_stake_api::stake_instruction;
+    use solana_stake_api::stake_state::Stake;
     use solana_vote_api::vote_instruction;
-    use std::collections::HashSet;
-    use std::iter::FromIterator;
     use std::sync::Arc;
 
     fn new_from_parent(parent: &Arc<Bank>, slot: u64) -> Bank {
@@ -144,8 +143,14 @@ pub(crate) mod tests {
         let mut expected = HashMap::new();
         assert_eq!(vote_account_stakes_at_epoch(&bank, 10), None);
 
+        let leader_stake = Stake {
+            stake: BOOTSTRAP_LEADER_LAMPORTS,
+            epoch: 0,
+            ..Stake::default()
+        };
+
         // First epoch has the bootstrap leader
-        expected.insert(voting_keypair.pubkey(), BOOTSTRAP_LEADER_LAMPORTS);
+        expected.insert(voting_keypair.pubkey(), leader_stake.stake(0));
         let expected = Some(expected);
         assert_eq!(vote_account_stakes_at_epoch(&bank, 0), expected);
 
@@ -205,7 +210,12 @@ pub(crate) mod tests {
 
     #[test]
     fn test_epoch_stakes_and_lockouts() {
-        let stake = 42;
+        let stake = BOOTSTRAP_LEADER_LAMPORTS * 100;
+        let leader_stake = Stake {
+            stake: BOOTSTRAP_LEADER_LAMPORTS,
+            ..Stake::default()
+        };
+
         let validator = Keypair::new();
 
         let GenesisBlockInfo {
@@ -233,6 +243,11 @@ pub(crate) mod tests {
             stake,
         );
 
+        let other_stake = Stake {
+            stake,
+            ..Stake::default()
+        };
+
         // soonest slot that could be a new epoch is 1
         let mut slot = 1;
         let mut epoch = bank.get_stakers_epoch(0);
@@ -240,17 +255,20 @@ pub(crate) mod tests {
         while bank.get_stakers_epoch(slot) == epoch {
             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![(BOOTSTRAP_LEADER_LAMPORTS, None)]);
+        assert_eq!(result, vec![(leader_stake.stake(0), None)]);
 
-        let result: HashSet<_> = HashSet::from_iter(epoch_stakes_and_lockouts(&bank, epoch));
+        let mut result: Vec<_> = epoch_stakes_and_lockouts(&bank, epoch);
-        let expected: HashSet<_> =
+        result.sort();
-            HashSet::from_iter(vec![(BOOTSTRAP_LEADER_LAMPORTS, None), (stake, None)]);
+        let mut expected = vec![
+            (leader_stake.stake(bank.epoch()), None),
+            (other_stake.stake(bank.epoch()), None),
+        ];
+        expected.sort();
         assert_eq!(result, expected);
     }
 

programs/stake_api/src/stake_state.rs

@@ -49,26 +49,10 @@ impl StakeState {
         account.state().ok()
     }
 
-    // utility function, used by Stakes, tests
+    pub fn stake_from(account: &Account) -> Option<Stake> {
-    pub fn voter_pubkey_from(account: &Account) -> Option<Pubkey> {
-        Self::from(account).and_then(|state: Self| state.voter_pubkey())
-    }
-
-    // utility function, used by Stakes, tests
-    pub fn voter_pubkey_and_stake_from(account: &Account) -> Option<(Pubkey, u64)> {
-        Self::from(account).and_then(|state: Self| state.voter_pubkey_and_stake())
-    }
-
-    pub fn stake_from(account: &Account) -> Option<(Stake)> {
         Self::from(account).and_then(|state: Self| state.stake())
     }
 
-    pub fn voter_pubkey(&self) -> Option<Pubkey> {
-        match self {
-            StakeState::Stake(stake) => Some(stake.voter_pubkey),
-            _ => None,
-        }
-    }
     pub fn stake(&self) -> Option<Stake> {
         match self {
             StakeState::Stake(stake) => Some(stake.clone()),
@@ -76,13 +60,6 @@ impl StakeState {
         }
     }
 
-    pub fn voter_pubkey_and_stake(&self) -> Option<(Pubkey, u64)> {
-        match self {
-            StakeState::Stake(stake) => Some((stake.voter_pubkey, stake.stake)),
-            _ => None,
-        }
-    }
-
     pub fn calculate_rewards(
         credits_observed: u64,
         stake: u64,
@@ -121,27 +98,26 @@ pub struct Stake {
     pub epoch: u64,      // epoch the stake was activated
     pub prev_stake: u64, // for warmup, cooldown
 }
-const STAKE_WARMUP_COOLDOWN_EPOCHS: u64 = 3;
+pub const STAKE_WARMUP_EPOCHS: u64 = 3;
 
 impl Stake {
-    pub fn stake(&mut self, epoch: u64) -> u64 {
+    pub fn stake(&self, epoch: u64) -> u64 {
         // prev_stake for stuff in the past
         if epoch < self.epoch {
             return self.prev_stake;
         }
-        if epoch - self.epoch >= STAKE_WARMUP_COOLDOWN_EPOCHS {
+        if epoch - self.epoch >= STAKE_WARMUP_EPOCHS {
             return self.stake;
         }
 
         if self.stake != 0 {
             // warmup
             // 1/3rd, then 2/3rds...
-            (self.stake / STAKE_WARMUP_COOLDOWN_EPOCHS) * (epoch - self.epoch + 1)
+            (self.stake / STAKE_WARMUP_EPOCHS) * (epoch - self.epoch + 1)
         } else if self.prev_stake != 0 {
             // cool down
             // 3/3rds, then 2/3rds...
-            self.prev_stake
+            self.prev_stake - ((self.prev_stake / STAKE_WARMUP_EPOCHS) * (epoch - self.epoch))
-                - ((self.prev_stake / STAKE_WARMUP_COOLDOWN_EPOCHS) * (epoch - self.epoch))
         } else {
             0
         }
@@ -307,7 +283,7 @@ pub fn create_stake_account(
             credits_observed: vote_state.credits(),
             stake: lamports,
             epoch: 0,
-            prev_stake: lamports,
+            prev_stake: 0,
         }))
         .expect("set_state");
 
@@ -414,25 +390,24 @@ mod tests {
     fn test_stake_stake() {
         let mut stake = Stake::default();
         assert_eq!(stake.stake(0), 0);
-        let staked = STAKE_WARMUP_COOLDOWN_EPOCHS;
+        let staked = STAKE_WARMUP_EPOCHS;
         stake.delegate(staked, &Pubkey::default(), &VoteState::default(), 1);
         // test warmup
-        for i in 0..STAKE_WARMUP_COOLDOWN_EPOCHS {
+        for i in 0..STAKE_WARMUP_EPOCHS {
             assert_eq!(stake.stake(i), i);
         }
-        assert_eq!(stake.stake(STAKE_WARMUP_COOLDOWN_EPOCHS * 42), staked);
+        assert_eq!(stake.stake(STAKE_WARMUP_EPOCHS * 42), staked);
 
-        stake.deactivate(STAKE_WARMUP_COOLDOWN_EPOCHS);
+        stake.deactivate(STAKE_WARMUP_EPOCHS);
 
         // test cooldown
-        for i in STAKE_WARMUP_COOLDOWN_EPOCHS..STAKE_WARMUP_COOLDOWN_EPOCHS * 2 {
+        for i in STAKE_WARMUP_EPOCHS..STAKE_WARMUP_EPOCHS * 2 {
             assert_eq!(
                 stake.stake(i),
-                staked
+                staked - (staked / STAKE_WARMUP_EPOCHS) * (i - STAKE_WARMUP_EPOCHS)
-                    - (staked / STAKE_WARMUP_COOLDOWN_EPOCHS) * (i - STAKE_WARMUP_COOLDOWN_EPOCHS)
             );
         }
-        assert_eq!(stake.stake(STAKE_WARMUP_COOLDOWN_EPOCHS * 42), 0);
+        assert_eq!(stake.stake(STAKE_WARMUP_EPOCHS * 42), 0);
     }
 
     #[test]

runtime/src/bank.rs

@@ -321,7 +321,7 @@ impl Bank {
 
         self.transaction_count
             .store(parent.transaction_count() as usize, Ordering::Relaxed);
-        self.stakes = RwLock::new(parent.stakes.read().unwrap().clone());
+        self.stakes = RwLock::new(parent.stakes.read().unwrap().clone_with_epoch(self.epoch()));
         self.storage_accounts = RwLock::new(parent.storage_accounts.read().unwrap().clone());
 
         self.tick_height.store(

runtime/src/stakes.rs

@@ -17,17 +17,42 @@ pub struct Stakes {
     /// unclaimed points.
     //  a point is a credit multiplied by the stake
     points: u64,
+
+    /// current epoch, used to calculate current stake
+    epoch: u64,
 }
 
 impl Stakes {
+    pub fn clone_with_epoch(&self, epoch: u64) -> Self {
+        if self.epoch == epoch {
+            self.clone()
+        } else {
+            Stakes {
+                stake_accounts: self.stake_accounts.clone(),
+                points: self.points,
+                epoch,
+                vote_accounts: self
+                    .vote_accounts
+                    .iter()
+                    .map(|(pubkey, (_stake, account))| {
+                        (
+                            *pubkey,
+                            (self.calculate_stake(pubkey, epoch), account.clone()),
+                        )
+                    })
+                    .collect(),
+            }
+        }
+    }
+
     // sum the stakes that point to the given voter_pubkey
-    fn calculate_stake(&self, voter_pubkey: &Pubkey) -> u64 {
+    fn calculate_stake(&self, voter_pubkey: &Pubkey, epoch: u64) -> u64 {
         self.stake_accounts
             .iter()
             .map(|(_, stake_account)| {
                 StakeState::stake_from(stake_account).map_or(0, |stake| {
                     if stake.voter_pubkey == *voter_pubkey {
-                        stake.stake
+                        stake.stake(epoch)
                     } else {
                         0
                     }
@@ -47,7 +72,7 @@ impl Stakes {
             } else {
                 let old = self.vote_accounts.get(pubkey);
 
-                let stake = old.map_or_else(|| self.calculate_stake(pubkey), |v| v.0);
+                let stake = old.map_or_else(|| self.calculate_stake(pubkey, self.epoch), |v| v.0);
 
                 // count any increase in points, can only go forward
                 let old_credits = old
@@ -62,23 +87,28 @@ impl Stakes {
             }
         } else if solana_stake_api::check_id(&account.owner) {
             //  old_stake is stake lamports and voter_pubkey from the pre-store() version
-            let old_stake = self
+            let old_stake = self.stake_accounts.get(pubkey).and_then(|old_account| {
-                .stake_accounts
+                StakeState::stake_from(old_account)
-                .get(pubkey)
+                    .map(|stake| (stake.voter_pubkey, stake.stake(self.epoch)))
-                .and_then(|old_account| StakeState::voter_pubkey_and_stake_from(old_account));
+            });
 
-            let stake = if account.lamports != 0 {
+            let stake = StakeState::stake_from(account).map(|stake| {
-                StakeState::voter_pubkey_and_stake_from(account)
+                (
-            } else {
+                    stake.voter_pubkey,
-                StakeState::voter_pubkey_from(account).map(|voter_pubkey| (voter_pubkey, 0))
+                    if account.lamports != 0 {
-            };
+                        stake.stake(self.epoch)
+                    } else {
+                        0
+                    },
+                )
+            });
 
             // if adjustments need to be made...
             if stake != old_stake {
-                if let Some((old_voter_pubkey, old_stake)) = old_stake {
+                if let Some((voter_pubkey, stake)) = old_stake {
                     self.vote_accounts
-                        .entry(old_voter_pubkey)
+                        .entry(voter_pubkey)
-                        .and_modify(|e| e.0 -= old_stake);
+                        .and_modify(|e| e.0 -= stake);
                 }
                 if let Some((voter_pubkey, stake)) = stake {
                     self.vote_accounts
@@ -140,7 +170,7 @@ impl Stakes {
 pub mod tests {
     use super::*;
     use solana_sdk::pubkey::Pubkey;
-    use solana_stake_api::stake_state;
+    use solana_stake_api::stake_state::{self, STAKE_WARMUP_EPOCHS};
     use solana_vote_api::vote_state::{self, VoteState, MAX_LOCKOUT_HISTORY};
 
     //  set up some dummies for a staked node     ((     vote      )  (     stake     ))
@@ -163,43 +193,47 @@ pub mod tests {
 
     #[test]
     fn test_stakes_basic() {
-        let mut stakes = Stakes::default();
+        for i in 0..STAKE_WARMUP_EPOCHS + 1 {
+            let mut stakes = Stakes::default();
+            stakes.epoch = i;
 
-        let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
+            let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
-            create_staked_node_accounts(10);
+                create_staked_node_accounts(10);
 
-        stakes.store(&vote_pubkey, &vote_account);
+            stakes.store(&vote_pubkey, &vote_account);
-        stakes.store(&stake_pubkey, &stake_account);
+            stakes.store(&stake_pubkey, &stake_account);
+            let stake = StakeState::stake_from(&stake_account).unwrap();
+            {
+                let vote_accounts = stakes.vote_accounts();
+                assert!(vote_accounts.get(&vote_pubkey).is_some());
+                assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, stake.stake(i));
+            }
 
-        {
+            stake_account.lamports = 42;
-            let vote_accounts = stakes.vote_accounts();
+            stakes.store(&stake_pubkey, &stake_account);
-            assert!(vote_accounts.get(&vote_pubkey).is_some());
+            {
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+                let vote_accounts = stakes.vote_accounts();
-        }
+                assert!(vote_accounts.get(&vote_pubkey).is_some());
+                assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, stake.stake(i)); // stays old stake, because only 10 is activated
+            }
 
-        stake_account.lamports = 42;
+            // activate more
-        stakes.store(&stake_pubkey, &stake_account);
+            let (_stake_pubkey, mut stake_account) = create_stake_account(42, &vote_pubkey);
-        {
+            stakes.store(&stake_pubkey, &stake_account);
-            let vote_accounts = stakes.vote_accounts();
+            let stake = StakeState::stake_from(&stake_account).unwrap();
-            assert!(vote_accounts.get(&vote_pubkey).is_some());
+            {
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10); // stays old stake, because only 10 is activated
+                let vote_accounts = stakes.vote_accounts();
-        }
+                assert!(vote_accounts.get(&vote_pubkey).is_some());
+                assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, stake.stake(i)); // now stake of 42 is activated
+            }
 
-        // activate more
+            stake_account.lamports = 0;
-        let (_stake_pubkey, mut stake_account) = create_stake_account(42, &vote_pubkey);
+            stakes.store(&stake_pubkey, &stake_account);
-        stakes.store(&stake_pubkey, &stake_account);
+            {
-        {
+                let vote_accounts = stakes.vote_accounts();
-            let vote_accounts = stakes.vote_accounts();
+                assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert!(vote_accounts.get(&vote_pubkey).is_some());
+                assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 42); // now stake of 42 is activated
+            }
-        }
-
-        stake_account.lamports = 0;
-        stakes.store(&stake_pubkey, &stake_account);
-        {
-            let vote_accounts = stakes.vote_accounts();
-            assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
         }
     }
 
@@ -216,7 +250,7 @@ pub mod tests {
         stakes.store(&stake_pubkey, &stake_account);
 
         let ((vote11_pubkey, vote11_account), (stake11_pubkey, stake11_account)) =
-            create_staked_node_accounts(11);
+            create_staked_node_accounts(20);
 
         stakes.store(&vote11_pubkey, &vote11_account);
         stakes.store(&stake11_pubkey, &stake11_account);
@@ -229,6 +263,8 @@ pub mod tests {
     #[test]
     fn test_stakes_points() {
         let mut stakes = Stakes::default();
+        stakes.epoch = STAKE_WARMUP_EPOCHS + 1;
+
         let stake = 42;
         assert_eq!(stakes.points(), 0);
         assert_eq!(stakes.claim_points(), 0);
@@ -273,6 +309,7 @@ pub mod tests {
     #[test]
     fn test_stakes_vote_account_disappear_reappear() {
         let mut stakes = Stakes::default();
+        stakes.epoch = STAKE_WARMUP_EPOCHS + 1;
 
         let ((vote_pubkey, mut vote_account), (stake_pubkey, stake_account)) =
             create_staked_node_accounts(10);
@@ -306,6 +343,7 @@ pub mod tests {
     #[test]
     fn test_stakes_change_delegate() {
         let mut stakes = Stakes::default();
+        stakes.epoch = STAKE_WARMUP_EPOCHS + 1;
 
         let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
             create_staked_node_accounts(10);
@@ -319,10 +357,15 @@ pub mod tests {
         // delegates to vote_pubkey
         stakes.store(&stake_pubkey, &stake_account);
 
+        let stake = StakeState::stake_from(&stake_account).unwrap();
+
         {
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+            assert_eq!(
+                vote_accounts.get(&vote_pubkey).unwrap().0,
+                stake.stake(stakes.epoch)
+            );
             assert!(vote_accounts.get(&vote_pubkey2).is_some());
             assert_eq!(vote_accounts.get(&vote_pubkey2).unwrap().0, 0);
         }
@@ -335,12 +378,16 @@ pub mod tests {
             assert!(vote_accounts.get(&vote_pubkey).is_some());
             assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
             assert!(vote_accounts.get(&vote_pubkey2).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey2).unwrap().0, 10);
+            assert_eq!(
+                vote_accounts.get(&vote_pubkey2).unwrap().0,
+                stake.stake(stakes.epoch)
+            );
         }
     }
     #[test]
     fn test_stakes_multiple_stakers() {
         let mut stakes = Stakes::default();
+        stakes.epoch = STAKE_WARMUP_EPOCHS + 1;
 
         let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
             create_staked_node_accounts(10);
@@ -359,10 +406,38 @@ pub mod tests {
             assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 20);
         }
     }
+    #[test]
+    fn test_clone_with_epoch() {
+        let mut stakes = Stakes::default();
+
+        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
+            create_staked_node_accounts(10);
+
+        stakes.store(&vote_pubkey, &vote_account);
+        stakes.store(&stake_pubkey, &stake_account);
+        let stake = StakeState::stake_from(&stake_account).unwrap();
+
+        {
+            let vote_accounts = stakes.vote_accounts();
+            assert_eq!(
+                vote_accounts.get(&vote_pubkey).unwrap().0,
+                stake.stake(stakes.epoch)
+            );
+        }
+        let stakes = stakes.clone_with_epoch(3);
+        {
+            let vote_accounts = stakes.vote_accounts();
+            assert_eq!(
+                vote_accounts.get(&vote_pubkey).unwrap().0,
+                stake.stake(stakes.epoch)
+            );
+        }
+    }
 
     #[test]
     fn test_stakes_not_delegate() {
         let mut stakes = Stakes::default();
+        stakes.epoch = STAKE_WARMUP_EPOCHS + 1;
 
         let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
             create_staked_node_accounts(10);