Fix fork detection (#10839)

* Fix fork detection

Co-authored-by: Carl <carl@solana.com>

core/src/replay_stage.rs

@@ -300,7 +300,7 @@ impl ReplayStage {
                 let mut tower = Tower::new(&my_pubkey, &vote_account, root, &heaviest_bank);
                 let mut current_leader = None;
                 let mut last_reset = Hash::default();
-                let mut partition = false;
+                let mut partition_exists = false;
                 let mut skipped_slots_info = SkippedSlotsInfo::default();
                 let mut replay_timing = ReplayTiming::default();
                 loop {
@@ -533,34 +533,39 @@ impl ReplayStage {
                             last_reset = reset_bank.last_blockhash();
                             tpu_has_bank = false;
 
-                            if !partition
+                            if let Some(last_voted_slot) = tower.last_voted_slot() {
-                                && vote_bank.as_ref().map(|(b, _)| b.slot())
+                                // If the current heaviest bank is not a descendant of the last voted slot,
-                                    != Some(reset_bank.slot())
+                                // there must be a partition
+                                let partition_detected = Self::is_partition_detected(&ancestors, last_voted_slot, heaviest_bank.slot());
+
+                                if !partition_exists && partition_detected
                                 {
                                     warn!(
-                                    "PARTITION DETECTED waiting to join fork: {} last vote: {:?}",
+                                        "PARTITION DETECTED waiting to join heaviest fork: {} last vote: {:?}, reset slot: {}",
+                                        heaviest_bank.slot(),
+                                        last_voted_slot,
                                         reset_bank.slot(),
-                                    tower.last_vote()
                                     );
                                     inc_new_counter_info!("replay_stage-partition_detected", 1);
                                     datapoint_info!(
                                         "replay_stage-partition",
                                         ("slot", reset_bank.slot() as i64, i64)
                                     );
-                                partition = true;
+                                    partition_exists = true;
-                            } else if partition
+                                } else if partition_exists
-                                && vote_bank.as_ref().map(|(b, _)| b.slot())
+                                    && !partition_detected
-                                    == Some(reset_bank.slot())
                                 {
                                     warn!(
-                                    "PARTITION resolved fork: {} last vote: {:?}",
+                                        "PARTITION resolved heaviest fork: {} last vote: {:?}, reset slot: {}",
-                                    reset_bank.slot(),
+                                        heaviest_bank.slot(),
-                                    tower.last_vote()
+                                        last_voted_slot,
+                                        reset_bank.slot()
                                     );
-                                partition = false;
+                                    partition_exists = false;
                                     inc_new_counter_info!("replay_stage-partition_resolved", 1);
                                 }
                             }
+                        }
                         Self::report_memory(&allocated, "reset_bank", start);
                     }
                     reset_bank_time.stop();
@@ -629,6 +634,18 @@ impl ReplayStage {
         }
     }
 
+    fn is_partition_detected(
+        ancestors: &HashMap<Slot, HashSet<Slot>>,
+        last_voted_slot: Slot,
+        heaviest_slot: Slot,
+    ) -> bool {
+        last_voted_slot != heaviest_slot
+            && !ancestors
+                .get(&heaviest_slot)
+                .map(|ancestors| ancestors.contains(&last_voted_slot))
+                .unwrap_or(true)
+    }
+
     fn report_memory(
         allocated: &solana_measure::thread_mem_usage::Allocatedp,
         name: &'static str,
@@ -1903,6 +1920,22 @@ pub(crate) mod tests {
     };
     use trees::tr;
 
+    #[test]
+    fn test_is_partition_detected() {
+        let (bank_forks, _) = setup_forks();
+        let ancestors = bank_forks.read().unwrap().ancestors();
+        // Last vote 1 is an ancestor of the heaviest slot 3, no partition
+        assert!(!ReplayStage::is_partition_detected(&ancestors, 1, 3));
+        // Last vote 1 is an ancestor of the from heaviest slot 1, no partition
+        assert!(!ReplayStage::is_partition_detected(&ancestors, 3, 3));
+        // Last vote 2 is not an ancestor of the heaviest slot 3,
+        // partition detected!
+        assert!(ReplayStage::is_partition_detected(&ancestors, 2, 3));
+        // Last vote 4 is not an ancestor of the heaviest slot 3,
+        // partition detected!
+        assert!(ReplayStage::is_partition_detected(&ancestors, 4, 3));
+    }
+
     #[test]
     fn test_child_slots_of_same_parent() {
         let ledger_path = get_tmp_ledger_path!();