remove forks.working_bank() where possible (#3010)

src/bank_forks.rs

@@ -2,32 +2,41 @@
 
 use solana_runtime::bank::Bank;
 use std::collections::HashMap;
+use std::ops::Index;
 use std::sync::Arc;
 
 pub struct BankForks {
-    working_bank_id: u64,
     banks: HashMap<u64, Arc<Bank>>,
+    working_bank: Arc<Bank>,
+}
+
+impl Index<u64> for BankForks {
+    type Output = Arc<Bank>;
+    fn index(&self, bank_id: u64) -> &Arc<Bank> {
+        &self.banks[&bank_id]
+    }
 }
 
 impl BankForks {
-    pub fn new(working_bank_id: u64, bank: Bank) -> Self {
+    pub fn new(bank_id: u64, bank: Bank) -> Self {
         let mut banks = HashMap::new();
-        banks.insert(working_bank_id, Arc::new(bank));
+        let working_bank = Arc::new(bank);
+        banks.insert(bank_id, working_bank.clone());
         Self {
-            working_bank_id,
             banks,
+            working_bank,
         }
     }
 
-    pub fn working_bank(&self) -> Arc<Bank> {
+    // TODO: use the bank's own ID instead of receiving a parameter?
-        self.banks[&self.working_bank_id].clone()
-    }
-
-    // TODO: use the bank's own ID instead of receiving a parameter
     pub fn insert(&mut self, bank_id: u64, bank: Bank) {
         let mut bank = Arc::new(bank);
         self.banks.insert(bank_id, bank.clone());
 
+        if bank_id > self.working_bank.id() {
+            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
@@ -37,10 +46,9 @@ impl BankForks {
         }
     }
 
-    pub fn set_working_bank_id(&mut self, bank_id: u64) {
+    // TODO: really want to kill this...
-        if self.banks.contains_key(&bank_id) {
+    pub fn working_bank(&self) -> Arc<Bank> {
-            self.working_bank_id = bank_id;
+        self.working_bank.clone()
-        }
     }
 }
 
@@ -51,21 +59,14 @@ mod tests {
     use solana_sdk::hash::Hash;
 
     #[test]
-    fn test_bank_forks_root() {
+    fn test_bank_forks() {
-        let bank = Bank::default();
-        let bank_forks = BankForks::new(0, bank);
-        assert_eq!(bank_forks.working_bank().tick_height(), 0);
-    }
-
-    #[test]
-    fn test_bank_forks_parent() {
         let (genesis_block, _) = GenesisBlock::new(10_000);
         let bank = Bank::new(&genesis_block);
         let mut bank_forks = BankForks::new(0, bank);
-        let child_bank = Bank::new_from_parent(&bank_forks.working_bank());
+        let child_bank = Bank::new_from_parent(&bank_forks[0u64]);
         child_bank.register_tick(&Hash::default());
         bank_forks.insert(1, child_bank);
-        bank_forks.set_working_bank_id(1);
+        assert_eq!(bank_forks[1u64].tick_height(), 1);
         assert_eq!(bank_forks.working_bank().tick_height(), 1);
     }
 }

src/blocktree_processor.rs

@@ -114,13 +114,13 @@ pub fn process_blocktree(
 
     // Setup bank for slot 0
     let (mut bank_forks, mut pending_slots) = {
-        let bank0 = Bank::new_with_paths(&genesis_block, account_paths);
+        let bank = Bank::new_with_paths(&genesis_block, account_paths);
         let slot = 0;
         let entry_height = 0;
-        let last_entry_id = bank0.last_id();
+        let last_entry_id = bank.last_id();
 
         (
-            BankForks::new(slot, bank0),
+            BankForks::new(slot, bank),
             vec![(slot, entry_height, last_entry_id)],
         )
     };
@@ -129,8 +129,7 @@ pub fn process_blocktree(
     while !pending_slots.is_empty() {
         let (slot, mut entry_height, mut last_entry_id) = pending_slots.pop().unwrap();
 
-        bank_forks.set_working_bank_id(slot);
+        let bank = bank_forks[slot].clone();
-        let bank = bank_forks.working_bank();
 
         // Load the metadata for this slot
         let meta = blocktree
@@ -355,7 +354,7 @@ mod tests {
         info!("last_fork1_entry_id: {:?}", last_fork1_entry_id);
         info!("last_fork2_entry_id: {:?}", last_fork2_entry_id);
 
-        let (mut bank_forks, bank_forks_info) =
+        let (bank_forks, bank_forks_info) =
             process_blocktree(&genesis_block, &blocktree, None).unwrap();
 
         assert_eq!(bank_forks_info.len(), 2); // There are two forks
@@ -380,8 +379,7 @@ mod tests {
 
         // Ensure bank_forks holds the right banks
         for info in bank_forks_info {
-            bank_forks.set_working_bank_id(info.bank_id);
+            assert_eq!(bank_forks[info.bank_id].last_id(), info.last_entry_id)
-            assert_eq!(bank_forks.working_bank().last_id(), info.last_entry_id)
         }
     }
 
@@ -484,7 +482,7 @@ mod tests {
             }
         );
 
-        let bank = bank_forks.working_bank();
+        let bank = bank_forks[1].clone();
         assert_eq!(bank.get_balance(&mint_keypair.pubkey()), 50 - 3);
         assert_eq!(bank.tick_height(), 2 * genesis_block.ticks_per_slot - 1);
         assert_eq!(bank.last_id(), entries.last().unwrap().id);

src/fullnode.rs

@@ -103,7 +103,6 @@ pub struct Fullnode {
     node_services: NodeServices,
     rotation_receiver: TvuRotationReceiver,
     blocktree: Arc<Blocktree>,
-    bank_forks: Arc<RwLock<BankForks>>,
     poh_service: PohService,
     poh_recorder: Arc<Mutex<PohRecorder>>,
 }
@@ -125,13 +124,12 @@ impl Fullnode {
         let id = keypair.pubkey();
         assert_eq!(id, node.info.id);
 
-        let (mut bank_forks, bank_forks_info, blocktree, ledger_signal_receiver) =
+        let (bank_forks, bank_forks_info, blocktree, ledger_signal_receiver) =
             new_banks_from_blocktree(ledger_path, config.account_paths.clone());
 
         let exit = Arc::new(AtomicBool::new(false));
         let bank_info = &bank_forks_info[0];
-        bank_forks.set_working_bank_id(bank_info.bank_id);
+        let bank = bank_forks[bank_info.bank_id].clone();
-        let bank = bank_forks.working_bank();
 
         info!(
             "starting PoH... {} {}",
@@ -265,7 +263,6 @@ impl Fullnode {
             broadcast_socket: node.sockets.broadcast,
             rotation_receiver,
             blocktree,
-            bank_forks,
             poh_service,
             poh_recorder,
         }
@@ -285,7 +282,7 @@ impl Fullnode {
             // TODO: This is not the correct bank.  Instead TVU should pass along the
             // frozen Bank for each completed block for RPC to use from it's notion of the "best"
             // available fork (until we want to surface multiple forks to RPC)
-            rpc_service.set_bank(&self.bank_forks.read().unwrap().working_bank());
+            rpc_service.set_bank(&rotation_info.bank);
         }
 
         if rotation_info.leader_id == self.id {
@@ -297,7 +294,7 @@ impl Fullnode {
                 FullnodeReturnType::ValidatorToLeaderRotation
             };
             self.node_services.tpu.switch_to_leader(
-                &self.bank_forks.read().unwrap().working_bank(),
+                &rotation_info.bank,
                 &self.poh_recorder,
                 self.tpu_sockets
                     .iter()

src/replay_stage.rs

@@ -187,9 +187,7 @@ impl ReplayStage {
 
         // Gather up all the metadata about the current state of the ledger
         let (mut bank, tick_height, mut last_entry_id, mut current_blob_index) = {
-            let mut bank_forks = bank_forks.write().unwrap();
+            let bank = bank_forks.read().unwrap()[bank_forks_info[0].bank_id].clone();
-            bank_forks.set_working_bank_id(bank_forks_info[0].bank_id);
-            let bank = bank_forks.working_bank();
             let tick_height = bank.tick_height();
             (
                 bank,
@@ -222,7 +220,7 @@ impl ReplayStage {
             // be updated by the TPU
             to_leader_sender
                 .send(TvuRotationInfo {
-                    bank: old_bank,
+                    bank: bank.clone(),
                     last_entry_id,
                     slot,
                     leader_id,
@@ -383,7 +381,7 @@ impl ReplayStage {
                         // RPC can be made aware of last slot's bank
                         to_leader_sender
                             .send(TvuRotationInfo {
-                                bank: old_bank,
+                                bank: bank.clone(),
                                 last_entry_id,
                                 slot: next_slot,
                                 leader_id,
@@ -423,8 +421,7 @@ impl ReplayStage {
         new_bank.squash();
         let mut bank_forks = bank_forks.write().unwrap();
         bank_forks.insert(slot, new_bank);
-        bank_forks.set_working_bank_id(slot);
+        bank_forks[slot].clone()
-        bank_forks.working_bank()
     }
 
     fn reset_state(