Repair alternate versions of dead slots (bp #9805) (#9886)

automerge

core/src/cluster_slots.rs

@@ -47,29 +47,7 @@ impl ClusterSlots {
                     self.keys.write().unwrap().insert(pubkey.clone());
                 }
                 let from = self.keys.read().unwrap().get(&pubkey).unwrap().clone();
-                let balance = self
-                    .validator_stakes
-                    .read()
-                    .unwrap()
-                    .get(&from)
-                    .map(|v| v.total_stake)
-                    .unwrap_or(0);
-
-                let mut slot_pubkeys = self.cluster_slots.read().unwrap().get(slot).cloned();
-                if slot_pubkeys.is_none() {
-                    let new_slot_pubkeys = Arc::new(RwLock::new(HashMap::default()));
-                    self.cluster_slots
-                        .write()
-                        .unwrap()
-                        .insert(*slot, new_slot_pubkeys.clone());
-                    slot_pubkeys = Some(new_slot_pubkeys);
-                }
-
-                slot_pubkeys
-                    .unwrap()
-                    .write()
-                    .unwrap()
-                    .insert(from.clone(), balance);
+                self.insert_node_id(*slot, from);
             }
         }
         self.cluster_slots.write().unwrap().retain(|x, _| *x > root);
@@ -79,6 +57,7 @@ impl ClusterSlots {
             .retain(|x| Arc::strong_count(x) > 1);
         *self.since.write().unwrap() = since;
     }
+
     pub fn collect(&self, id: &Pubkey) -> HashSet<Slot> {
         self.cluster_slots
             .read()
@@ -90,6 +69,30 @@ impl ClusterSlots {
             .collect()
     }
 
+    pub fn insert_node_id(&self, slot: Slot, node_id: Arc<Pubkey>) {
+        let balance = self
+            .validator_stakes
+            .read()
+            .unwrap()
+            .get(&node_id)
+            .map(|v| v.total_stake)
+            .unwrap_or(0);
+        let mut slot_pubkeys = self.cluster_slots.read().unwrap().get(&slot).cloned();
+        if slot_pubkeys.is_none() {
+            let new_slot_pubkeys = Arc::new(RwLock::new(HashMap::default()));
+            self.cluster_slots
+                .write()
+                .unwrap()
+                .insert(slot, new_slot_pubkeys.clone());
+            slot_pubkeys = Some(new_slot_pubkeys);
+        }
+        slot_pubkeys
+            .unwrap()
+            .write()
+            .unwrap()
+            .insert(node_id, balance);
+    }
+
     fn update_peers(&self, cluster_info: &ClusterInfo, bank_forks: &RwLock<BankForks>) {
         let root_bank = bank_forks.read().unwrap().root_bank().clone();
         let root_epoch = root_bank.epoch();
@@ -137,6 +140,23 @@ impl ClusterSlots {
             .collect()
     }
 
+    pub fn compute_weights_exclude_noncomplete(
+        &self,
+        slot: Slot,
+        repair_peers: &[ContactInfo],
+    ) -> Vec<(u64, usize)> {
+        let slot_peers = self.lookup(slot);
+        repair_peers
+            .iter()
+            .enumerate()
+            .filter_map(|(i, x)| {
+                slot_peers
+                    .as_ref()
+                    .and_then(|v| v.read().unwrap().get(&x.id).map(|stake| (*stake + 1, i)))
+            })
+            .collect()
+    }
+
     pub fn generate_repairs_for_missing_slots(
         &self,
         self_id: &Pubkey,
@@ -274,6 +294,43 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_best_completed_slot_peer() {
+        let cs = ClusterSlots::default();
+        let mut contact_infos = vec![ContactInfo::default(); 2];
+        for ci in contact_infos.iter_mut() {
+            ci.id = Pubkey::new_rand();
+        }
+        let slot = 9;
+
+        // None of these validators have completed slot 9, so should
+        // return nothing
+        assert!(cs
+            .compute_weights_exclude_noncomplete(slot, &contact_infos)
+            .is_empty());
+
+        // Give second validator max stake
+        let validator_stakes: HashMap<_, _> = vec![(
+            *Arc::new(contact_infos[1].id),
+            NodeVoteAccounts {
+                total_stake: std::u64::MAX / 2,
+                vote_accounts: vec![Pubkey::default()],
+            },
+        )]
+        .into_iter()
+        .collect();
+        *cs.validator_stakes.write().unwrap() = Arc::new(validator_stakes);
+
+        // Mark the first validator as completed slot 9, should pick that validator,
+        // even though it only has default stake, while the other validator has
+        // max stake
+        cs.insert_node_id(slot, Arc::new(contact_infos[0].id));
+        assert_eq!(
+            cs.compute_weights_exclude_noncomplete(slot, &contact_infos),
+            vec![(1, 0)]
+        );
+    }
+
     #[test]
     fn test_update_new_staked_slot() {
         let cs = ClusterSlots::default();