diff --git a/core/src/bank_forks.rs b/core/src/bank_forks.rs index 0174b0e2d0..4590d39f57 100644 --- a/core/src/bank_forks.rs +++ b/core/src/bank_forks.rs @@ -68,6 +68,7 @@ impl BankForks { .map(|(k, _v)| *k) .collect() } + pub fn get(&self, bank_slot: u64) -> Option<&Arc> { self.banks.get(&bank_slot) } diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs index e88f34aeb7..933fc3fdeb 100644 --- a/core/src/replay_stage.rs +++ b/core/src/replay_stage.rs @@ -199,6 +199,7 @@ impl ReplayStage { ); if let Some(new_root) = locktower.record_vote(bank.slot()) { bank_forks.write().unwrap().set_root(new_root); + Self::handle_new_root(&bank_forks, &mut progress); } locktower.update_epoch(&bank); cluster_info.write().unwrap().push_vote(vote); @@ -406,6 +407,14 @@ impl ReplayStage { Ok(()) } + fn handle_new_root( + bank_forks: &Arc>, + progress: &mut HashMap, + ) { + let r_bank_forks = bank_forks.read().unwrap(); + progress.retain(|k, _| r_bank_forks.get(*k).is_some()); + } + fn process_completed_bank( my_id: &Pubkey, bank: Arc, @@ -636,4 +645,14 @@ mod test { let _ignored = remove_dir_all(&ledger_path); } + + #[test] + fn test_handle_new_root() { + let bank0 = Bank::default(); + let bank_forks = Arc::new(RwLock::new(BankForks::new(0, bank0))); + let mut progress = HashMap::new(); + progress.insert(5, (Hash::default(), 0)); + ReplayStage::handle_new_root(&bank_forks, &mut progress); + assert!(progress.is_empty()); + } }