Confidence implementation (#5993)

* Change confidence parameters

* Add status_cache_ancestors to get all relevant ancestors of a bank including roots from status cache

* Fix and add tests

* Clippy

core/src/confidence.rs

@@ -1,113 +1,59 @@
-use crate::consensus::StakeLockout;
+use crate::result::{Error, Result};
 use crate::service::Service;
-use std::collections::{HashMap, HashSet};
+use solana_runtime::bank::Bank;
+use solana_vote_api::vote_state::VoteState;
+use solana_vote_api::vote_state::MAX_LOCKOUT_HISTORY;
+use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::mpsc::{channel, Receiver, Sender};
+use std::sync::mpsc::{channel, Receiver, RecvTimeoutError, Sender};
 use std::sync::{Arc, RwLock};
 use std::thread::{self, Builder, JoinHandle};
+use std::time::Duration;
 
-#[derive(Debug, Default, PartialEq)]
+#[derive(Debug, Default, Eq, PartialEq)]
-pub struct Confidence {
+pub struct BankConfidence {
-    fork_stakes: u64,
+    confidence: [u64; MAX_LOCKOUT_HISTORY],
-    total_stake: u64,
-    lockouts: u64,
-    stake_weighted_lockouts: u128,
 }
 
-impl Confidence {
+impl BankConfidence {
-    pub fn new(fork_stakes: u64, total_stake: u64, lockouts: u64) -> Self {
+    pub fn increase_confirmation_stake(&mut self, confirmation_count: usize, stake: u64) {
-        Self {
+        assert!(confirmation_count > 0 && confirmation_count <= MAX_LOCKOUT_HISTORY);
-            fork_stakes,
+        self.confidence[confirmation_count - 1] += stake;
-            total_stake,
-            lockouts,
-            stake_weighted_lockouts: 0,
-        }
     }
-    pub fn new_with_stake_weighted(
+
-        fork_stakes: u64,
+    pub fn get_confirmation_stake(&mut self, confirmation_count: usize) -> u64 {
-        total_stake: u64,
+        assert!(confirmation_count > 0 && confirmation_count <= MAX_LOCKOUT_HISTORY);
-        lockouts: u64,
+        self.confidence[confirmation_count - 1]
-        stake_weighted_lockouts: u128,
-    ) -> Self {
-        Self {
-            fork_stakes,
-            total_stake,
-            lockouts,
-            stake_weighted_lockouts,
-        }
     }
 }
 
-#[derive(Default, PartialEq)]
+#[derive(Default)]
 pub struct ForkConfidenceCache {
-    confidence: HashMap<u64, Confidence>,
+    bank_confidence: HashMap<u64, BankConfidence>,
+    _total_stake: u64,
 }
 
 impl ForkConfidenceCache {
-    pub fn cache_fork_confidence(
+    pub fn new(bank_confidence: HashMap<u64, BankConfidence>, total_stake: u64) -> Self {
-        &mut self,
+        Self {
-        fork: u64,
+            bank_confidence,
-        fork_stakes: u64,
+            _total_stake: total_stake,
-        total_stake: u64,
+        }
-        lockouts: u64,
-    ) {
-        self.confidence
-            .entry(fork)
-            .and_modify(|entry| {
-                entry.fork_stakes = fork_stakes;
-                entry.total_stake = total_stake;
-                entry.lockouts = lockouts;
-            })
-            .or_insert_with(|| Confidence::new(fork_stakes, total_stake, lockouts));
     }
 
-    pub fn cache_stake_weighted_lockouts(&mut self, fork: u64, stake_weighted_lockouts: u128) {
+    pub fn get_fork_confidence(&self, fork: u64) -> Option<&BankConfidence> {
-        self.confidence
+        self.bank_confidence.get(&fork)
-            .entry(fork)
-            .and_modify(|entry| {
-                entry.stake_weighted_lockouts = stake_weighted_lockouts;
-            })
-            .or_insert(Confidence {
-                fork_stakes: 0,
-                total_stake: 0,
-                lockouts: 0,
-                stake_weighted_lockouts,
-            });
-    }
-
-    pub fn get_fork_confidence(&self, fork: u64) -> Option<&Confidence> {
-        self.confidence.get(&fork)
-    }
-
-    pub fn prune_confidence_cache(&mut self, ancestors: &HashMap<u64, HashSet<u64>>, root: u64) {
-        // For Every slot `s` in this cache must exist some bank `b` in BankForks with
-        // `b.slot() == s`, and because `ancestors` has an entry for every bank in BankForks,
-        // then there must be an entry in `ancestors` for every slot in `self.confidence`
-        self.confidence
-            .retain(|slot, _| slot == &root || ancestors[&slot].contains(&root));
     }
 }
 
 pub struct ConfidenceAggregationData {
-    lockouts: HashMap<u64, StakeLockout>,
+    bank: Arc<Bank>,
-    root: Option<u64>,
-    ancestors: Arc<HashMap<u64, HashSet<u64>>>,
     total_staked: u64,
 }
 
 impl ConfidenceAggregationData {
-    pub fn new(
+    pub fn new(bank: Arc<Bank>, total_staked: u64) -> Self {
-        lockouts: HashMap<u64, StakeLockout>,
+        Self { bank, total_staked }
-        root: Option<u64>,
-        ancestors: Arc<HashMap<u64, HashSet<u64>>>,
-        total_staked: u64,
-    ) -> Self {
-        Self {
-            lockouts,
-            root,
-            ancestors,
-            total_staked,
-        }
     }
 }
 
@@ -116,38 +62,17 @@ pub struct AggregateConfidenceService {
 }
 
 impl AggregateConfidenceService {
-    pub fn aggregate_confidence(
-        root: Option<u64>,
-        ancestors: &HashMap<u64, HashSet<u64>>,
-        stake_lockouts: &HashMap<u64, StakeLockout>,
-    ) -> HashMap<u64, u128> {
-        let mut stake_weighted_lockouts: HashMap<u64, u128> = HashMap::new();
-        for (fork, lockout) in stake_lockouts.iter() {
-            if root.is_none() || *fork >= root.unwrap() {
-                let mut slot_with_ancestors = vec![*fork];
-                slot_with_ancestors.extend(ancestors.get(&fork).unwrap_or(&HashSet::new()));
-                for slot in slot_with_ancestors {
-                    if root.is_none() || slot >= root.unwrap() {
-                        let entry = stake_weighted_lockouts.entry(slot).or_default();
-                        *entry += u128::from(lockout.lockout()) * u128::from(lockout.stake());
-                    }
-                }
-            }
-        }
-        stake_weighted_lockouts
-    }
-
     pub fn new(
         exit: &Arc<AtomicBool>,
         fork_confidence_cache: Arc<RwLock<ForkConfidenceCache>>,
     ) -> (Sender<ConfidenceAggregationData>, Self) {
-        let (lockouts_sender, lockouts_receiver): (
+        let (sender, receiver): (
             Sender<ConfidenceAggregationData>,
             Receiver<ConfidenceAggregationData>,
         ) = channel();
         let exit_ = exit.clone();
         (
-            lockouts_sender,
+            sender,
             Self {
                 t_confidence: Builder::new()
                     .name("solana-aggregate-stake-lockouts".to_string())
@@ -155,54 +80,121 @@ impl AggregateConfidenceService {
                         if exit_.load(Ordering::Relaxed) {
                             break;
                         }
-                        if let Ok(aggregation_data) = lockouts_receiver.try_recv() {
-                            let stake_weighted_lockouts = Self::aggregate_confidence(
-                                aggregation_data.root,
-                                &aggregation_data.ancestors,
-                                &aggregation_data.lockouts,
-                            );
 
-                            let mut w_fork_confidence_cache =
+                        if let Err(e) = Self::run(&receiver, &fork_confidence_cache, &exit_) {
-                                fork_confidence_cache.write().unwrap();
+                            match e {
-
+                                Error::RecvTimeoutError(RecvTimeoutError::Disconnected) => break,
-                            // Cache the confidence values
+                                Error::RecvTimeoutError(RecvTimeoutError::Timeout) => (),
-                            for (fork, stake_lockout) in aggregation_data.lockouts.iter() {
+                                _ => info!(
-                                if aggregation_data.root.is_none()
+                                    "Unexpected error from AggregateConfidenceService: {:?}",
-                                    || *fork >= aggregation_data.root.unwrap()
+                                    e
-                                {
+                                ),
-                                    w_fork_confidence_cache.cache_fork_confidence(
-                                        *fork,
-                                        stake_lockout.stake(),
-                                        aggregation_data.total_staked,
-                                        stake_lockout.lockout(),
-                                    );
-                                }
                             }
-
-                            // Cache the stake weighted lockouts
-                            for (fork, stake_weighted_lockout) in stake_weighted_lockouts.iter() {
-                                if aggregation_data.root.is_none()
-                                    || *fork >= aggregation_data.root.unwrap()
-                                {
-                                    w_fork_confidence_cache.cache_stake_weighted_lockouts(
-                                        *fork,
-                                        *stake_weighted_lockout,
-                                    )
-                                }
-                            }
-
-                            if let Some(root) = aggregation_data.root {
-                                w_fork_confidence_cache
-                                    .prune_confidence_cache(&aggregation_data.ancestors, root);
-                            }
-
-                            drop(w_fork_confidence_cache);
                         }
                     })
                     .unwrap(),
             },
         )
     }
+
+    fn run(
+        receiver: &Receiver<ConfidenceAggregationData>,
+        fork_confidence_cache: &RwLock<ForkConfidenceCache>,
+        exit: &Arc<AtomicBool>,
+    ) -> Result<()> {
+        loop {
+            if exit.load(Ordering::Relaxed) {
+                return Ok(());
+            }
+
+            let mut aggregation_data = receiver.recv_timeout(Duration::from_secs(1))?;
+
+            while let Ok(new_data) = receiver.try_recv() {
+                aggregation_data = new_data;
+            }
+
+            let ancestors = aggregation_data.bank.status_cache_ancestors();
+            if ancestors.is_empty() {
+                continue;
+            }
+
+            let bank_confidence = Self::aggregate_confidence(&ancestors, &aggregation_data.bank);
+
+            let mut new_fork_confidence =
+                ForkConfidenceCache::new(bank_confidence, aggregation_data.total_staked);
+
+            let mut w_fork_confidence_cache = fork_confidence_cache.write().unwrap();
+
+            std::mem::swap(&mut *w_fork_confidence_cache, &mut new_fork_confidence);
+        }
+    }
+
+    pub fn aggregate_confidence(ancestors: &[u64], bank: &Bank) -> HashMap<u64, BankConfidence> {
+        assert!(!ancestors.is_empty());
+
+        // Check ancestors is sorted
+        for a in ancestors.windows(2) {
+            assert!(a[0] < a[1]);
+        }
+
+        let mut confidence = HashMap::new();
+        for (_, (lamports, account)) in bank.vote_accounts().into_iter() {
+            if lamports == 0 {
+                continue;
+            }
+            let vote_state = VoteState::from(&account);
+            if vote_state.is_none() {
+                continue;
+            }
+
+            let vote_state = vote_state.unwrap();
+            Self::aggregate_confidence_for_vote_account(
+                &mut confidence,
+                &vote_state,
+                ancestors,
+                lamports,
+            );
+        }
+
+        confidence
+    }
+
+    fn aggregate_confidence_for_vote_account(
+        confidence: &mut HashMap<u64, BankConfidence>,
+        vote_state: &VoteState,
+        ancestors: &[u64],
+        lamports: u64,
+    ) {
+        assert!(!ancestors.is_empty());
+        let mut ancestors_index = 0;
+        if let Some(root) = vote_state.root_slot {
+            for (i, a) in ancestors.iter().enumerate() {
+                if *a <= root {
+                    confidence
+                        .entry(*a)
+                        .or_insert_with(BankConfidence::default)
+                        .increase_confirmation_stake(MAX_LOCKOUT_HISTORY, lamports);
+                } else {
+                    ancestors_index = i;
+                    break;
+                }
+            }
+        }
+
+        for vote in &vote_state.votes {
+            while ancestors[ancestors_index] <= vote.slot {
+                confidence
+                    .entry(ancestors[ancestors_index])
+                    .or_insert_with(BankConfidence::default)
+                    .increase_confirmation_stake(vote.confirmation_count as usize, lamports);
+                ancestors_index += 1;
+
+                if ancestors_index == ancestors.len() {
+                    return;
+                }
+            }
+        }
+    }
 }
 
 impl Service for AggregateConfidenceService {
@@ -216,64 +208,169 @@ impl Service for AggregateConfidenceService {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::genesis_utils::{create_genesis_block, GenesisBlockInfo};
+    use solana_sdk::pubkey::Pubkey;
+    use solana_stake_api::stake_state;
+    use solana_vote_api::vote_state;
 
     #[test]
-    fn test_fork_confidence_cache() {
+    fn test_bank_confidence() {
-        let mut cache = ForkConfidenceCache::default();
+        let mut cache = BankConfidence::default();
-        let fork = 0;
+        assert_eq!(cache.get_confirmation_stake(1), 0);
-        assert!(cache.confidence.get(&fork).is_none());
+        cache.increase_confirmation_stake(1, 10);
-        cache.cache_fork_confidence(fork, 11, 12, 13);
+        assert_eq!(cache.get_confirmation_stake(1), 10);
-        assert_eq!(
+        cache.increase_confirmation_stake(1, 20);
-            cache.confidence.get(&fork).unwrap(),
+        assert_eq!(cache.get_confirmation_stake(1), 30);
-            &Confidence {
-                fork_stakes: 11,
-                total_stake: 12,
-                lockouts: 13,
-                stake_weighted_lockouts: 0,
-            }
-        );
-        // Ensure that {fork_stakes, total_stake, lockouts} and stake_weighted_lockouts
-        // can be updated separately
-        cache.cache_stake_weighted_lockouts(fork, 20);
-        assert_eq!(
-            cache.confidence.get(&fork).unwrap(),
-            &Confidence {
-                fork_stakes: 11,
-                total_stake: 12,
-                lockouts: 13,
-                stake_weighted_lockouts: 20,
-            }
-        );
-        cache.cache_fork_confidence(fork, 21, 22, 23);
-        assert_eq!(
-            cache.confidence.get(&fork).unwrap().stake_weighted_lockouts,
-            20,
-        );
     }
 
     #[test]
-    fn test_aggregate_confidence() {
+    fn test_aggregate_confidence_for_vote_account_1() {
-        let stakes = vec![
+        let ancestors = vec![3, 4, 5, 7, 9, 11];
-            (0, StakeLockout::new(1, 32)),
+        let mut confidence = HashMap::new();
-            (1, StakeLockout::new(1, 24)),
+        let lamports = 5;
-            (2, StakeLockout::new(1, 16)),
+        let mut vote_state = VoteState::new(&Pubkey::default(), &Pubkey::default(), 0);
-            (3, StakeLockout::new(1, 8)),
+
-        ]
+        let root = ancestors.last().unwrap();
-        .into_iter()
+        vote_state.root_slot = Some(*root);
-        .collect();
+        AggregateConfidenceService::aggregate_confidence_for_vote_account(
-        let ancestors = vec![
+            &mut confidence,
-            (0, HashSet::new()),
+            &vote_state,
-            (1, vec![0].into_iter().collect()),
+            &ancestors,
-            (2, vec![0, 1].into_iter().collect()),
+            lamports,
-            (3, vec![0, 1, 2].into_iter().collect()),
+        );
-        ]
+
-        .into_iter()
+        for a in ancestors {
-        .collect();
+            let mut expected = BankConfidence::default();
-        let stake_weighted_lockouts =
+            expected.increase_confirmation_stake(MAX_LOCKOUT_HISTORY, lamports);
-            AggregateConfidenceService::aggregate_confidence(Some(1), &ancestors, &stakes);
+            assert_eq!(*confidence.get(&a).unwrap(), expected);
-        assert!(stake_weighted_lockouts.get(&0).is_none());
+        }
-        assert_eq!(*stake_weighted_lockouts.get(&1).unwrap(), 8 + 16 + 24);
+    }
-        assert_eq!(*stake_weighted_lockouts.get(&2).unwrap(), 8 + 16);
+
-        assert_eq!(*stake_weighted_lockouts.get(&3).unwrap(), 8);
+    #[test]
+    fn test_aggregate_confidence_for_vote_account_2() {
+        let ancestors = vec![3, 4, 5, 7, 9, 11];
+        let mut confidence = HashMap::new();
+        let lamports = 5;
+        let mut vote_state = VoteState::new(&Pubkey::default(), &Pubkey::default(), 0);
+
+        let root = ancestors[2];
+        vote_state.root_slot = Some(root);
+        vote_state.process_slot_vote_unchecked(*ancestors.last().unwrap());
+        AggregateConfidenceService::aggregate_confidence_for_vote_account(
+            &mut confidence,
+            &vote_state,
+            &ancestors,
+            lamports,
+        );
+
+        for a in ancestors {
+            if a <= root {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(MAX_LOCKOUT_HISTORY, lamports);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(1, lamports);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            }
+        }
+    }
+
+    #[test]
+    fn test_aggregate_confidence_for_vote_account_3() {
+        let ancestors = vec![3, 4, 5, 7, 9, 10, 11];
+        let mut confidence = HashMap::new();
+        let lamports = 5;
+        let mut vote_state = VoteState::new(&Pubkey::default(), &Pubkey::default(), 0);
+
+        let root = ancestors[2];
+        vote_state.root_slot = Some(root);
+        assert!(ancestors[4] + 2 >= ancestors[6]);
+        vote_state.process_slot_vote_unchecked(ancestors[4]);
+        vote_state.process_slot_vote_unchecked(ancestors[6]);
+        AggregateConfidenceService::aggregate_confidence_for_vote_account(
+            &mut confidence,
+            &vote_state,
+            &ancestors,
+            lamports,
+        );
+
+        for (i, a) in ancestors.iter().enumerate() {
+            if *a <= root {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(MAX_LOCKOUT_HISTORY, lamports);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else if i <= 4 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(2, lamports);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else if i <= 6 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(1, lamports);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            }
+        }
+    }
+
+    #[test]
+    fn test_aggregate_confidence_validity() {
+        let ancestors = vec![3, 4, 5, 7, 9, 10, 11];
+        let GenesisBlockInfo {
+            mut genesis_block, ..
+        } = create_genesis_block(10_000);
+
+        let pk1 = Pubkey::new_rand();
+        let mut vote_account1 = vote_state::create_account(&pk1, &Pubkey::new_rand(), 0, 100);
+        let stake_account1 = stake_state::create_account(&pk1, &vote_account1, 100);
+        let pk2 = Pubkey::new_rand();
+        let mut vote_account2 = vote_state::create_account(&pk2, &Pubkey::new_rand(), 0, 50);
+        let stake_account2 = stake_state::create_account(&pk2, &vote_account2, 50);
+
+        genesis_block.accounts.extend(vec![
+            (pk1, vote_account1.clone()),
+            (Pubkey::new_rand(), stake_account1),
+            (pk2, vote_account2.clone()),
+            (Pubkey::new_rand(), stake_account2),
+        ]);
+
+        // Create bank
+        let bank = Arc::new(Bank::new(&genesis_block));
+
+        let mut vote_state1 = VoteState::from(&vote_account1).unwrap();
+        vote_state1.process_slot_vote_unchecked(3);
+        vote_state1.process_slot_vote_unchecked(5);
+        vote_state1.to(&mut vote_account1).unwrap();
+        bank.store_account(&pk1, &vote_account1);
+
+        let mut vote_state2 = VoteState::from(&vote_account2).unwrap();
+        vote_state2.process_slot_vote_unchecked(9);
+        vote_state2.process_slot_vote_unchecked(10);
+        vote_state2.to(&mut vote_account2).unwrap();
+        bank.store_account(&pk2, &vote_account2);
+
+        let confidence = AggregateConfidenceService::aggregate_confidence(&ancestors, &bank);
+
+        for a in ancestors {
+            if a <= 3 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(2, 150);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else if a <= 5 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(1, 100);
+                expected.increase_confirmation_stake(2, 50);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else if a <= 9 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(2, 50);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else if a <= 10 {
+                let mut expected = BankConfidence::default();
+                expected.increase_confirmation_stake(1, 50);
+                assert_eq!(*confidence.get(&a).unwrap(), expected);
+            } else {
+                assert!(confidence.get(&a).is_none());
+            }
+        }
     }
 }

core/src/replay_stage.rs

@@ -156,7 +156,7 @@ impl ReplayStage {
                         &mut progress,
                     );
 
-                    if let Some((_, bank, lockouts, total_staked)) = votable.into_iter().last() {
+                    if let Some((_, bank, _, total_staked)) = votable.into_iter().last() {
                         subscriptions.notify_subscribers(bank.slot(), &bank_forks);
 
                         if let Some(votable_leader) =
@@ -184,7 +184,6 @@ impl ReplayStage {
                         Self::handle_votable_bank(
                             &bank,
                             &bank_forks,
-                            &ancestors,
                             &mut tower,
                             &mut progress,
                             &vote_account,
@@ -193,7 +192,6 @@ impl ReplayStage {
                             &blocktree,
                             &leader_schedule_cache,
                             &root_bank_sender,
-                            lockouts,
                             total_staked,
                             &lockouts_sender,
                             &snapshot_package_sender,
@@ -406,7 +404,6 @@ impl ReplayStage {
     fn handle_votable_bank<T>(
         bank: &Arc<Bank>,
         bank_forks: &Arc<RwLock<BankForks>>,
-        ancestors: &Arc<HashMap<u64, HashSet<u64>>>,
         tower: &mut Tower,
         progress: &mut HashMap<u64, ForkProgress>,
         vote_account: &Pubkey,
@@ -415,7 +412,6 @@ impl ReplayStage {
         blocktree: &Arc<Blocktree>,
         leader_schedule_cache: &Arc<LeaderScheduleCache>,
         root_bank_sender: &Sender<Vec<Arc<Bank>>>,
-        lockouts: HashMap<u64, StakeLockout>,
         total_staked: u64,
         lockouts_sender: &Sender<ConfidenceAggregationData>,
         snapshot_package_sender: &Option<SnapshotPackageSender>,
@@ -455,7 +451,7 @@ impl ReplayStage {
                 Err(e)?;
             }
         }
-        Self::update_confidence_cache(ancestors, tower, lockouts, total_staked, lockouts_sender);
+        Self::update_confidence_cache(bank.clone(), total_staked, lockouts_sender);
 
         if let Some(ref voting_keypair) = voting_keypair {
             let node_keypair = cluster_info.read().unwrap().keypair.clone();
@@ -476,18 +472,11 @@ impl ReplayStage {
     }
 
     fn update_confidence_cache(
-        ancestors: &Arc<HashMap<u64, HashSet<u64>>>,
+        bank: Arc<Bank>,
-        tower: &Tower,
-        lockouts: HashMap<u64, StakeLockout>,
         total_staked: u64,
         lockouts_sender: &Sender<ConfidenceAggregationData>,
     ) {
-        if let Err(e) = lockouts_sender.send(ConfidenceAggregationData::new(
+        if let Err(e) = lockouts_sender.send(ConfidenceAggregationData::new(bank, total_staked)) {
-            lockouts,
-            tower.root(),
-            ancestors.clone(),
-            total_staked,
-        )) {
             trace!("lockouts_sender failed: {:?}", e);
         }
     }
@@ -809,7 +798,7 @@ mod test {
     use super::*;
     use crate::blocktree::tests::make_slot_entries;
     use crate::blocktree::{entries_to_test_shreds, get_tmp_ledger_path};
-    use crate::confidence::Confidence;
+    use crate::confidence::BankConfidence;
     use crate::entry;
     use crate::genesis_utils::{create_genesis_block, create_genesis_block_with_leader};
     use crate::replay_stage::ReplayStage;
@@ -1028,43 +1017,18 @@ mod test {
             &[arc_bank0.clone()],
             vec![0],
         )));
-        let pubkey = Pubkey::new_rand();
-        let mut tower = Tower::new(&pubkey, &Pubkey::new_rand(), &bank_forks.read().unwrap());
-        let mut progress = HashMap::new();
 
-        leader_vote(&arc_bank0, &leader_voting_pubkey);
+        assert!(fork_confidence_cache
-        let ancestors = Arc::new(bank_forks.read().unwrap().ancestors());
+            .read()
-
+            .unwrap()
-        let votable =
+            .get_fork_confidence(0)
-            ReplayStage::generate_votable_banks(&ancestors, &bank_forks, &tower, &mut progress);
+            .is_none());
-        if let Some((_, _, lockouts, total_staked)) = votable.into_iter().last() {
-            ReplayStage::update_confidence_cache(
-                &ancestors,
-                &tower,
-                lockouts,
-                total_staked,
-                &lockouts_sender,
-            );
-        }
-
-        thread::sleep(Duration::from_millis(200));
-
-        assert_eq!(
-            fork_confidence_cache
-                .read()
-                .unwrap()
-                .get_fork_confidence(0)
-                .unwrap(),
-            &Confidence::new(0, 3, 2)
-        );
         assert!(fork_confidence_cache
             .read()
             .unwrap()
             .get_fork_confidence(1)
             .is_none());
 
-        tower.record_vote(arc_bank0.slot(), arc_bank0.hash());
-
         let bank1 = Bank::new_from_parent(&arc_bank0, &Pubkey::default(), arc_bank0.slot() + 1);
         let _res = bank1.transfer(10, &genesis_block_info.mint_keypair, &Pubkey::new_rand());
         for _ in 0..genesis_block.ticks_per_slot {
@@ -1074,20 +1038,7 @@ mod test {
         bank_forks.write().unwrap().insert(bank1);
         let arc_bank1 = bank_forks.read().unwrap().get(1).unwrap().clone();
         leader_vote(&arc_bank1, &leader_voting_pubkey);
-        let ancestors = Arc::new(bank_forks.read().unwrap().ancestors());
+        ReplayStage::update_confidence_cache(arc_bank1.clone(), leader_lamports, &lockouts_sender);
-        let votable =
-            ReplayStage::generate_votable_banks(&ancestors, &bank_forks, &tower, &mut progress);
-        if let Some((_, _, lockouts, total_staked)) = votable.into_iter().last() {
-            ReplayStage::update_confidence_cache(
-                &ancestors,
-                &tower,
-                lockouts,
-                total_staked,
-                &lockouts_sender,
-            );
-        }
-
-        tower.record_vote(arc_bank1.slot(), arc_bank1.hash());
 
         let bank2 = Bank::new_from_parent(&arc_bank1, &Pubkey::default(), arc_bank1.slot() + 1);
         let _res = bank2.transfer(10, &genesis_block_info.mint_keypair, &Pubkey::new_rand());
@@ -1098,43 +1049,38 @@ mod test {
         bank_forks.write().unwrap().insert(bank2);
         let arc_bank2 = bank_forks.read().unwrap().get(2).unwrap().clone();
         leader_vote(&arc_bank2, &leader_voting_pubkey);
-        let ancestors = Arc::new(bank_forks.read().unwrap().ancestors());
+        ReplayStage::update_confidence_cache(arc_bank2.clone(), leader_lamports, &lockouts_sender);
-        let votable =
-            ReplayStage::generate_votable_banks(&ancestors, &bank_forks, &tower, &mut progress);
-        if let Some((_, _, lockouts, total_staked)) = votable.into_iter().last() {
-            ReplayStage::update_confidence_cache(
-                &ancestors,
-                &tower,
-                lockouts,
-                total_staked,
-                &lockouts_sender,
-            );
-        }
         thread::sleep(Duration::from_millis(200));
 
+        let mut expected0 = BankConfidence::default();
+        expected0.increase_confirmation_stake(2, leader_lamports);
         assert_eq!(
             fork_confidence_cache
                 .read()
                 .unwrap()
                 .get_fork_confidence(0)
                 .unwrap(),
-            &Confidence::new_with_stake_weighted(3, 3, 14, 60)
+            &expected0,
         );
+        let mut expected1 = BankConfidence::default();
+        expected1.increase_confirmation_stake(2, leader_lamports);
         assert_eq!(
             fork_confidence_cache
                 .read()
                 .unwrap()
                 .get_fork_confidence(1)
                 .unwrap(),
-            &Confidence::new_with_stake_weighted(3, 3, 6, 18)
+            &expected1
         );
+        let mut expected2 = BankConfidence::default();
+        expected2.increase_confirmation_stake(1, leader_lamports);
         assert_eq!(
             fork_confidence_cache
                 .read()
                 .unwrap()
                 .get_fork_confidence(2)
                 .unwrap(),
-            &Confidence::new_with_stake_weighted(0, 3, 2, 0)
+            &expected2
         );
     }
 }

runtime/src/bank.rs

@@ -402,6 +402,20 @@ impl Bank {
         *self.hash.read().unwrap() != Hash::default()
     }
 
+    pub fn status_cache_ancestors(&self) -> Vec<u64> {
+        let mut roots = self.src.status_cache.read().unwrap().roots().clone();
+        let min = roots.iter().min().cloned().unwrap_or(0);
+        for ancestor in self.ancestors.keys() {
+            if *ancestor >= min {
+                roots.insert(*ancestor);
+            }
+        }
+
+        let mut ancestors: Vec<_> = roots.into_iter().collect();
+        ancestors.sort();
+        ancestors
+    }
+
     fn update_clock(&self) {
         self.store_account(
             &clock::id(),
@@ -1539,6 +1553,7 @@ mod tests {
     use crate::genesis_utils::{
         create_genesis_block_with_leader, GenesisBlockInfo, BOOTSTRAP_LEADER_LAMPORTS,
     };
+    use crate::status_cache::MAX_CACHE_ENTRIES;
     use bincode::{deserialize_from, serialize_into, serialized_size};
     use solana_sdk::clock::DEFAULT_TICKS_PER_SLOT;
     use solana_sdk::genesis_block::create_genesis_block;
@@ -3021,4 +3036,22 @@ mod tests {
         assert_eq!(bank1.get_program_accounts(&program_id).len(), 2);
         assert_eq!(bank3.get_program_accounts(&program_id).len(), 2);
     }
+
+    #[test]
+    fn test_status_cache_ancestors() {
+        let (genesis_block, _mint_keypair) = create_genesis_block(500);
+        let parent = Arc::new(Bank::new(&genesis_block));
+        let bank1 = Arc::new(new_from_parent(&parent));
+        let mut bank = bank1;
+        for _ in 0..MAX_CACHE_ENTRIES * 2 {
+            bank = Arc::new(new_from_parent(&bank));
+            bank.squash();
+        }
+
+        let bank = new_from_parent(&bank);
+        assert_eq!(
+            bank.status_cache_ancestors(),
+            (bank.slot() - MAX_CACHE_ENTRIES as u64..=bank.slot()).collect::<Vec<_>>()
+        );
+    }
 }