Add root to bank forks (#4206)

core/src/bank_forks.rs

@@ -11,6 +11,7 @@ use std::time::Instant;
 pub struct BankForks {
     banks: HashMap<u64, Arc<Bank>>,
     working_bank: Arc<Bank>,
+    root: u64,
 }
 
 impl Index<u64> for BankForks {
@@ -28,6 +29,7 @@ impl BankForks {
         Self {
             banks,
             working_bank,
+            root: 0,
         }
     }
 
@@ -79,13 +81,14 @@ impl BankForks {
         self.banks.get(&bank_slot)
     }
 
-    pub fn new_from_banks(initial_banks: &[Arc<Bank>]) -> Self {
+    pub fn new_from_banks(initial_banks: &[Arc<Bank>], root: u64) -> Self {
         let mut banks = HashMap::new();
         let working_bank = initial_banks[0].clone();
         for bank in initial_banks {
             banks.insert(bank.slot(), bank.clone());
         }
         Self {
+            root,
             banks,
             working_bank,
         }
@@ -105,6 +108,7 @@ impl BankForks {
     }
 
     pub fn set_root(&mut self, root: u64) {
+        self.root = root;
         let set_root_start = Instant::now();
         let root_bank = self
             .banks
@@ -112,12 +116,17 @@ impl BankForks {
             .expect("root bank didn't exist in bank_forks");
         root_bank.squash();
         self.prune_non_root(root);
+
         inc_new_counter_info!(
             "bank-forks_set_root_ms",
             timing::duration_as_ms(&set_root_start.elapsed()) as usize
         );
     }
 
+    pub fn root(&self) -> u64 {
+        self.root
+    }
+
     fn prune_non_root(&mut self, root: u64) {
         let descendants = self.descendants();
         self.banks

core/src/blocktree_processor.rs

@@ -157,6 +157,7 @@ pub fn process_blocktree(
 
     let mut fork_info = vec![];
     let mut last_status_report = Instant::now();
+    let mut root = 0;
     while !pending_slots.is_empty() {
         let (slot, meta, bank, mut entry_height, mut last_entry_hash) =
             pending_slots.pop().unwrap();
@@ -210,6 +211,7 @@ pub fn process_blocktree(
         bank.freeze(); // all banks handled by this routine are created from complete slots
 
         if blocktree.is_root(slot) {
+            root = slot;
             leader_schedule_cache.set_root(slot);
             bank.squash();
             pending_slots.clear();
@@ -270,7 +272,7 @@ pub fn process_blocktree(
     }
 
     let (banks, bank_forks_info): (Vec<_>, Vec<_>) = fork_info.into_iter().unzip();
-    let bank_forks = BankForks::new_from_banks(&banks);
+    let bank_forks = BankForks::new_from_banks(&banks, root);
     info!(
         "processing ledger...complete in {}ms, forks={}...",
         duration_as_ms(&now.elapsed()),
@@ -437,6 +439,8 @@ mod tests {
             assert_eq!(bank_forks[info.bank_slot].slot(), info.bank_slot);
             assert!(bank_forks[info.bank_slot].is_frozen());
         }
+
+        assert_eq!(bank_forks.root(), 4);
     }
 
     #[test]
@@ -521,6 +525,8 @@ mod tests {
             &[1]
         );
 
+        assert_eq!(bank_forks.root(), 1);
+
         // Ensure bank_forks holds the right banks
         for info in bank_forks_info {
             assert_eq!(bank_forks[info.bank_slot].slot(), info.bank_slot);
@@ -643,6 +649,7 @@ mod tests {
             process_blocktree(&genesis_block, &blocktree, None).unwrap();
 
         assert_eq!(bank_forks_info.len(), 1);
+        assert_eq!(bank_forks.root(), 0);
         assert_eq!(
             bank_forks_info[0],
             BankForksInfo {

core/src/storage_stage.rs

@@ -480,7 +480,7 @@ mod tests {
         let cluster_info = test_cluster_info(&keypair.pubkey());
         let (genesis_block, _mint_keypair) = create_genesis_block(1000);
         let bank = Arc::new(Bank::new(&genesis_block));
-        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank])));
+        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank], 0)));
         let (_slot_sender, slot_receiver) = channel();
         let storage_state = StorageState::new();
         let storage_stage = StorageStage::new(
@@ -519,7 +519,7 @@ mod tests {
         let blocktree = Arc::new(Blocktree::open(&ledger_path).unwrap());
         let slot = 1;
         let bank = Arc::new(Bank::new(&genesis_block));
-        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank])));
+        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank], 0)));
         blocktree
             .write_entries(slot, 0, 0, ticks_per_slot, &entries)
             .unwrap();
@@ -594,7 +594,7 @@ mod tests {
             .write_entries(1, 0, 0, ticks_per_slot, &entries)
             .unwrap();
         let bank = Arc::new(Bank::new(&genesis_block));
-        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank])));
+        let bank_forks = Arc::new(RwLock::new(BankForks::new_from_banks(&[bank], 0)));
         let cluster_info = test_cluster_info(&keypair.pubkey());
 
         let (slot_sender, slot_receiver) = channel();