Modify bank_forks to support squashing/filtering new root and also don't remove parents from bank_forks when inserting, otherwise we lose potential fork points when querying blocktree for child slots

core/src/bank_forks.rs

@@ -31,15 +31,9 @@ impl BankForks {
     /// Create a map of bank slot id to the set of ancestors for the bank slot.
     pub fn ancestors(&self) -> HashMap<u64, HashSet<u64>> {
         let mut ancestors = HashMap::new();
-        let mut pending: Vec<Arc<Bank>> = self.banks.values().cloned().collect();
-        while !pending.is_empty() {
-            let bank = pending.pop().unwrap();
-            if ancestors.get(&bank.slot()).is_some() {
-                continue;
-            }
+        for bank in self.banks.values() {
             let set = bank.parents().into_iter().map(|b| b.slot()).collect();
             ancestors.insert(bank.slot(), set);
-            pending.extend(bank.parents().into_iter());
         }
         ancestors
     }
@@ -47,15 +41,7 @@ impl BankForks {
     /// Create a map of bank slot id to the set of all of its descendants
     pub fn descendants(&self) -> HashMap<u64, HashSet<u64>> {
         let mut descendants = HashMap::new();
-        let mut pending: Vec<Arc<Bank>> = self.banks.values().cloned().collect();
-        let mut done = HashSet::new();
-        assert!(!pending.is_empty());
-        while !pending.is_empty() {
-            let bank = pending.pop().unwrap();
-            if done.contains(&bank.slot()) {
-                continue;
-            }
-            done.insert(bank.slot());
+        for bank in self.banks.values() {
             let _ = descendants.entry(bank.slot()).or_insert(HashSet::new());
             for parent in bank.parents() {
                 descendants
@@ -63,22 +49,18 @@ impl BankForks {
                     .or_insert(HashSet::new())
                     .insert(bank.slot());
             }
-            pending.extend(bank.parents().into_iter());
         }
         descendants
     }
 
     pub fn frozen_banks(&self) -> HashMap<u64, Arc<Bank>> {
-        let mut frozen_banks: Vec<Arc<Bank>> = vec![];
-        frozen_banks.extend(self.banks.values().filter(|v| v.is_frozen()).cloned());
-        frozen_banks.extend(
-            self.banks
-                .iter()
-                .flat_map(|(_, v)| v.parents())
-                .filter(|v| v.is_frozen()),
-        );
-        frozen_banks.into_iter().map(|b| (b.slot(), b)).collect()
+        self.banks
+            .iter()
+            .filter(|(_, b)| b.is_frozen())
+            .map(|(k, b)| (*k, b.clone()))
+            .collect()
     }
+
     pub fn active_banks(&self) -> Vec<u64> {
         self.banks
             .iter()
@@ -104,28 +86,37 @@ impl BankForks {
 
     // TODO: use the bank's own ID instead of receiving a parameter?
     pub fn insert(&mut self, bank_slot: u64, bank: Bank) {
-        let mut bank = Arc::new(bank);
+        let bank = Arc::new(bank);
         assert_eq!(bank_slot, bank.slot());
         let prev = self.banks.insert(bank_slot, bank.clone());
         assert!(prev.is_none());
 
         self.working_bank = bank.clone();
-
-        // TODO: this really only needs to look at the first
-        //  parent if we're always calling insert()
-        //  when we construct a child bank
-        while let Some(parent) = bank.parent() {
-            if let Some(prev) = self.banks.remove(&parent.slot()) {
-                assert!(Arc::ptr_eq(&prev, &parent));
-            }
-            bank = parent;
-        }
     }
 
     // TODO: really want to kill this...
     pub fn working_bank(&self) -> Arc<Bank> {
         self.working_bank.clone()
     }
+
+    pub fn set_root(&mut self, root: u64) {
+        let root_bank = self
+            .banks
+            .get(&root)
+            .expect("root bank didn't exist in bank_forks");
+        root_bank.squash();
+        self.prune_non_root(root);
+    }
+
+    fn prune_non_root(&mut self, root: u64) {
+        self.banks.retain(|slot, bank| {
+            if *slot < root {
+                false
+            } else {
+                bank.is_descendant_of(root)
+            }
+        })
+    }
 }
 
 #[cfg(test)]

core/src/locktower.rs

@@ -152,8 +152,14 @@ impl Locktower {
         }
     }
 
-    pub fn record_vote(&mut self, slot: u64) {
+    pub fn record_vote(&mut self, slot: u64) -> Option<u64> {
+        let root_slot = self.lockouts.root_slot;
         self.lockouts.process_vote(Vote { slot });
+        if root_slot != self.lockouts.root_slot {
+            Some(self.lockouts.root_slot.unwrap())
+        } else {
+            None
+        }
     }
 
     pub fn calculate_weight(&self, stake_lockouts: &HashMap<u64, StakeLockout>) -> u128 {

core/src/replay_stage.rs

@@ -170,7 +170,9 @@ impl ReplayStage {
                                 bank.last_blockhash(),
                                 0,
                             );
-                            locktower.record_vote(bank.slot());
+                            if let Some(new_root) = locktower.record_vote(bank.slot()) {
+                                bank_forks.write().unwrap().set_root(new_root);
+                            }
                             locktower.update_epoch(&bank);
                             cluster_info.write().unwrap().push_vote(vote);
                         }
@@ -399,6 +401,7 @@ impl ReplayStage {
         let next_slots = blocktree
             .get_slots_since(&frozen_bank_slots)
             .expect("Db error");
+        // Filter out what we've already seen
         trace!("generate new forks {:?}", next_slots);
         for (parent_id, children) in next_slots {
             let parent_bank = frozen_banks
@@ -406,12 +409,8 @@ impl ReplayStage {
                 .expect("missing parent in bank forks")
                 .clone();
             for child_id in children {
-                if frozen_banks.get(&child_id).is_some() {
-                    trace!("child already frozen {}", child_id);
-                    continue;
-                }
                 if forks.get(child_id).is_some() {
-                    trace!("child already active {}", child_id);
+                    trace!("child already active or frozen {}", child_id);
                     continue;
                 }
                 let leader = leader_schedule_utils::slot_leader_at(child_id, &parent_bank).unwrap();
@@ -437,11 +436,12 @@ impl Service for ReplayStage {
 mod test {
     use super::*;
     use crate::banking_stage::create_test_recorder;
-    use crate::blocktree::create_new_tmp_ledger;
+    use crate::blocktree::{create_new_tmp_ledger, get_tmp_ledger_path};
     use crate::cluster_info::{ClusterInfo, Node};
     use crate::entry::create_ticks;
     use crate::entry::{next_entry_mut, Entry};
     use crate::fullnode::new_banks_from_blocktree;
+    use crate::packet::Blob;
     use crate::replay_stage::ReplayStage;
     use crate::result::Error;
     use solana_sdk::genesis_block::GenesisBlock;
@@ -576,4 +576,40 @@ mod test {
         }
         assert!(forward_entry_receiver.try_recv().is_err());
     }
+
+    #[test]
+    fn test_child_slots_of_same_parent() {
+        let ledger_path = get_tmp_ledger_path!();
+        {
+            let blocktree = Arc::new(
+                Blocktree::open(&ledger_path).expect("Expected to be able to open database ledger"),
+            );
+
+            let genesis_block = GenesisBlock::new(10_000).0;
+            let bank0 = Bank::new(&genesis_block);
+            let mut bank_forks = BankForks::new(0, bank0);
+            bank_forks.working_bank().freeze();
+
+            // Insert blob for slot 1, generate new forks, check result
+            let mut blob_slot_1 = Blob::default();
+            blob_slot_1.set_slot(1);
+            blob_slot_1.set_parent(0);
+            blocktree.insert_data_blobs(&vec![blob_slot_1]).unwrap();
+            assert!(bank_forks.get(1).is_none());
+            ReplayStage::generate_new_bank_forks(&blocktree, &mut bank_forks);
+            assert!(bank_forks.get(1).is_some());
+
+            // Inset blob for slot 3, generate new forks, check result
+            let mut blob_slot_2 = Blob::default();
+            blob_slot_2.set_slot(2);
+            blob_slot_2.set_parent(0);
+            blocktree.insert_data_blobs(&vec![blob_slot_2]).unwrap();
+            assert!(bank_forks.get(2).is_none());
+            ReplayStage::generate_new_bank_forks(&blocktree, &mut bank_forks);
+            assert!(bank_forks.get(1).is_some());
+            assert!(bank_forks.get(2).is_some());
+        }
+
+        let _ignored = remove_dir_all(&ledger_path);
+    }
 }

runtime/src/bank.rs

@@ -859,6 +859,24 @@ impl Bank {
         self.accounts
             .store_slow(self.accounts_id, &program_id, &bogus_account);
     }
+
+    pub fn is_descendant_of(&self, parent: u64) -> bool {
+        if self.slot() == parent {
+            return true;
+        }
+        let mut next_parent = self.parent();
+
+        while let Some(p) = next_parent {
+            if p.slot() == parent {
+                return true;
+            } else if p.slot() < parent {
+                return false;
+            }
+            next_parent = p.parent();
+        }
+
+        false
+    }
 }
 
 #[cfg(test)]
@@ -1624,4 +1642,28 @@ mod tests {
 
         assert_eq!(bank.is_votable(), true);
     }
+
+    #[test]
+    fn test_is_descendant_of() {
+        let (genesis_block, _) = GenesisBlock::new(1);
+        let parent = Arc::new(Bank::new(&genesis_block));
+        // Bank 1
+        let bank = Arc::new(new_from_parent(&parent));
+        // Bank 2
+        let bank2 = new_from_parent(&bank);
+        // Bank 5
+        let bank5 = Bank::new_from_parent(&bank, &Pubkey::default(), 5);
+
+        // Parents of bank 2: 0 -> 1 -> 2
+        assert!(bank2.is_descendant_of(0));
+        assert!(bank2.is_descendant_of(1));
+        assert!(bank2.is_descendant_of(2));
+        assert!(!bank2.is_descendant_of(3));
+
+        // Parents of bank 3: 0 -> 1 -> 5
+        assert!(bank5.is_descendant_of(0));
+        assert!(bank5.is_descendant_of(1));
+        assert!(!bank5.is_descendant_of(2));
+        assert!(!bank5.is_descendant_of(4));
+    }
 }