From 9c490e06b0dc78f1f45223fa6f88058293e7ea95 Mon Sep 17 00:00:00 2001
From: carllin <wumu727@gmail.com>
Date: Fri, 11 Sep 2020 02:03:11 -0700
Subject: [PATCH] Fix propagation on startup from snapshot (#12177)

---
 core/src/progress_map.rs |  2 +-
 core/src/replay_stage.rs | 76 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/core/src/progress_map.rs b/core/src/progress_map.rs
index 4a94cc0fe..47eba07e3 100644
--- a/core/src/progress_map.rs
+++ b/core/src/progress_map.rs
@@ -168,7 +168,7 @@ impl ForkProgress {
         num_dropped_blocks_on_fork: u64,
     ) -> Self {
         let validator_fork_info = {
-            if bank.collector_id() == my_pubkey && bank.slot() > 0 {
+            if bank.collector_id() == my_pubkey {
                 let stake = bank.epoch_vote_account_stake(voting_pubkey);
                 Some(ValidatorStakeInfo::new(
                     *voting_pubkey,
diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
index 67e8f1655..347ba1d4c 100644
--- a/core/src/replay_stage.rs
+++ b/core/src/replay_stage.rs
@@ -1594,13 +1594,13 @@ impl ReplayStage {
         loop {
             // These cases mean confirmation of propagation on any earlier
             // leader blocks must have been reached
-            if current_leader_slot == None || current_leader_slot.unwrap() <= root {
+            if current_leader_slot == None || current_leader_slot.unwrap() < root {
                 break;
             }
 
             let leader_propagated_stats = progress
                 .get_propagated_stats_mut(current_leader_slot.unwrap())
-                .expect("current_leader_slot > root, so must exist in the progress map");
+                .expect("current_leader_slot >= root, so must exist in the progress map");
 
             // If a descendant has reached propagation threshold, then
             // all its ancestor banks have also reached propagation
@@ -3319,6 +3319,10 @@ pub(crate) mod tests {
         let stake_per_validator = 10_000;
         let (mut bank_forks, mut progress_map, _) =
             initialize_state(&keypairs, stake_per_validator);
+        progress_map
+            .get_propagated_stats_mut(0)
+            .unwrap()
+            .is_leader_slot = true;
         bank_forks.set_root(0, &None, None);
         let total_epoch_stake = bank_forks.root_bank().total_epoch_stake();
 
@@ -3396,6 +3400,10 @@ pub(crate) mod tests {
         let stake_per_validator = 10_000;
         let (mut bank_forks, mut progress_map, _) =
             initialize_state(&keypairs, stake_per_validator);
+        progress_map
+            .get_propagated_stats_mut(0)
+            .unwrap()
+            .is_leader_slot = true;
         bank_forks.set_root(0, &None, None);
 
         let total_epoch_stake = num_validators as u64 * stake_per_validator;
@@ -3728,6 +3736,70 @@ pub(crate) mod tests {
         }
     }
 
+    #[test]
+    fn test_leader_snapshot_restart_propagation() {
+        let ReplayBlockstoreComponents {
+            validator_voting_keys,
+            mut progress,
+            bank_forks,
+            leader_schedule_cache,
+            ..
+        } = replay_blockstore_components();
+
+        let root_bank = bank_forks.read().unwrap().root_bank().clone();
+        let my_pubkey = leader_schedule_cache
+            .slot_leader_at(root_bank.slot(), Some(&root_bank))
+            .unwrap();
+
+        // Check that we are the leader of the root bank
+        assert!(
+            progress
+                .get_propagated_stats(root_bank.slot())
+                .unwrap()
+                .is_leader_slot
+        );
+        let ancestors = bank_forks.read().unwrap().ancestors();
+
+        // Freeze bank so it shows up in frozen banks
+        root_bank.freeze();
+        let mut frozen_banks: Vec<_> = bank_forks
+            .read()
+            .unwrap()
+            .frozen_banks()
+            .values()
+            .cloned()
+            .collect();
+
+        // Compute bank stats, make sure vote is propagated back to starting root bank
+        let vote_tracker = VoteTracker::default();
+
+        // Add votes
+        for vote_key in validator_voting_keys.values() {
+            vote_tracker.insert_vote(root_bank.slot(), Arc::new(*vote_key));
+        }
+
+        assert!(!progress.is_propagated(root_bank.slot()));
+
+        // Update propagation status
+        let tower = Tower::new_for_tests(0, 0.67);
+        ReplayStage::compute_bank_stats(
+            &my_pubkey,
+            &ancestors,
+            &mut frozen_banks,
+            &tower,
+            &mut progress,
+            &vote_tracker,
+            &ClusterSlots::default(),
+            &bank_forks,
+            &mut PubkeyReferences::default(),
+            &mut HeaviestSubtreeForkChoice::new_from_bank_forks(&bank_forks.read().unwrap()),
+            &mut BankWeightForkChoice::default(),
+        );
+
+        // Check status is true
+        assert!(progress.is_propagated(root_bank.slot()));
+    }
+
     fn setup_forks() -> (RwLock<BankForks>, ProgressMap) {
         /*
             Build fork structure: