From fe9e771b9bd62949afda627876a2ee556f6d49d4 Mon Sep 17 00:00:00 2001
From: carllin <wumu727@gmail.com>
Date: Tue, 19 Mar 2019 17:30:36 -0700
Subject: [PATCH] Clear progress map on squash (#3377)

---
 core/src/bank_forks.rs   |  1 +
 core/src/replay_stage.rs | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+)

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<Bank>> {
         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<RwLock<BankForks>>,
+        progress: &mut HashMap<u64, (Hash, usize)>,
+    ) {
+        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<Bank>,
@@ -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());
+    }
 }