Use program cache fork graph in extract() (#33806)

* Use program cache fork graph instead of WorkingSlot trait

* Fix deadlocked tests

* keep WorkingSlot trait for now

client-test/tests/client.rs

@@ -168,7 +168,7 @@ fn test_account_subscription() {
     // Transfer 100 lamports from alice to bob
     let tx = system_transaction::transfer(&alice, &bob.pubkey(), 100, blockhash);
     bank_forks
-        .write()
+        .read()
         .unwrap()
         .get(1)
         .unwrap()
@@ -373,7 +373,7 @@ fn test_program_subscription() {
     // Create new program account at bob's address
     let tx = system_transaction::create_account(&alice, &bob, blockhash, 100, 0, &program_id);
     bank_forks
-        .write()
+        .read()
         .unwrap()
         .get(1)
         .unwrap()

core/src/commitment_service.rs

@@ -510,12 +510,11 @@ mod tests {
 
         let bank0 = Bank::new_for_tests(&genesis_config);
         let bank_forks = BankForks::new_rw_arc(bank0);
-        let mut bank_forks = bank_forks.write().unwrap();
 
         // Fill bank_forks with banks with votes landing in the next slot
         // Create enough banks such that vote account will root slots 0 and 1
         for x in 0..33 {
-            let previous_bank = bank_forks.get(x).unwrap();
+            let previous_bank = bank_forks.read().unwrap().get(x).unwrap();
             let bank = Bank::new_from_parent(previous_bank.clone(), &Pubkey::default(), x + 1);
             let vote = vote_transaction::new_vote_transaction(
                 vec![x],
@@ -527,20 +526,23 @@ mod tests {
                 None,
             );
             bank.process_transaction(&vote).unwrap();
-            bank_forks.insert(bank);
+            bank_forks.write().unwrap().insert(bank);
         }
 
-        let working_bank = bank_forks.working_bank();
+        let working_bank = bank_forks.read().unwrap().working_bank();
         let root = get_vote_account_root_slot(
             validator_vote_keypairs.vote_keypair.pubkey(),
             &working_bank,
         );
         for x in 0..root {
-            bank_forks.set_root(x, &AbsRequestSender::default(), None);
+            bank_forks
+                .write()
+                .unwrap()
+                .set_root(x, &AbsRequestSender::default(), None);
         }
 
         // Add an additional bank/vote that will root slot 2
-        let bank33 = bank_forks.get(33).unwrap();
+        let bank33 = bank_forks.read().unwrap().get(33).unwrap();
         let bank34 = Bank::new_from_parent(bank33.clone(), &Pubkey::default(), 34);
         let vote33 = vote_transaction::new_vote_transaction(
             vec![33],
@@ -552,9 +554,9 @@ mod tests {
             None,
         );
         bank34.process_transaction(&vote33).unwrap();
-        bank_forks.insert(bank34);
+        bank_forks.write().unwrap().insert(bank34);
 
-        let working_bank = bank_forks.working_bank();
+        let working_bank = bank_forks.read().unwrap().working_bank();
         let root = get_vote_account_root_slot(
             validator_vote_keypairs.vote_keypair.pubkey(),
             &working_bank,
@@ -573,21 +575,22 @@ mod tests {
             .read()
             .unwrap()
             .highest_super_majority_root();
-        bank_forks.set_root(
+        bank_forks.write().unwrap().set_root(
             root,
             &AbsRequestSender::default(),
             Some(highest_super_majority_root),
         );
-        let highest_super_majority_root_bank = bank_forks.get(highest_super_majority_root);
+        let highest_super_majority_root_bank =
+            bank_forks.read().unwrap().get(highest_super_majority_root);
         assert!(highest_super_majority_root_bank.is_some());
 
         // Add a forked bank. Because the vote for bank 33 landed in the non-ancestor, the vote
         // account's root (and thus the highest_super_majority_root) rolls back to slot 1
-        let bank33 = bank_forks.get(33).unwrap();
+        let bank33 = bank_forks.read().unwrap().get(33).unwrap();
         let bank35 = Bank::new_from_parent(bank33, &Pubkey::default(), 35);
-        bank_forks.insert(bank35);
+        bank_forks.write().unwrap().insert(bank35);
 
-        let working_bank = bank_forks.working_bank();
+        let working_bank = bank_forks.read().unwrap().working_bank();
         let ancestors = working_bank.status_cache_ancestors();
         let _ = AggregateCommitmentService::update_commitment_cache(
             &block_commitment_cache,
@@ -602,13 +605,14 @@ mod tests {
             .read()
             .unwrap()
             .highest_super_majority_root();
-        let highest_super_majority_root_bank = bank_forks.get(highest_super_majority_root);
+        let highest_super_majority_root_bank =
+            bank_forks.read().unwrap().get(highest_super_majority_root);
         assert!(highest_super_majority_root_bank.is_some());
 
         // Add additional banks beyond lockout built on the new fork to ensure that behavior
         // continues normally
         for x in 35..=37 {
-            let previous_bank = bank_forks.get(x).unwrap();
+            let previous_bank = bank_forks.read().unwrap().get(x).unwrap();
             let bank = Bank::new_from_parent(previous_bank.clone(), &Pubkey::default(), x + 1);
             let vote = vote_transaction::new_vote_transaction(
                 vec![x],
@@ -620,10 +624,10 @@ mod tests {
                 None,
             );
             bank.process_transaction(&vote).unwrap();
-            bank_forks.insert(bank);
+            bank_forks.write().unwrap().insert(bank);
         }
 
-        let working_bank = bank_forks.working_bank();
+        let working_bank = bank_forks.read().unwrap().working_bank();
         let root = get_vote_account_root_slot(
             validator_vote_keypairs.vote_keypair.pubkey(),
             &working_bank,
@@ -642,12 +646,13 @@ mod tests {
             .read()
             .unwrap()
             .highest_super_majority_root();
-        bank_forks.set_root(
+        bank_forks.write().unwrap().set_root(
             root,
             &AbsRequestSender::default(),
             Some(highest_super_majority_root),
         );
-        let highest_super_majority_root_bank = bank_forks.get(highest_super_majority_root);
+        let highest_super_majority_root_bank =
+            bank_forks.read().unwrap().get(highest_super_majority_root);
         assert!(highest_super_majority_root_bank.is_some());
     }
 }

core/tests/snapshots.rs

@@ -221,12 +221,13 @@ fn run_bank_forks_snapshot_n<F>(
         accounts_package_sender,
     };
     for slot in 1..=last_slot {
-        let mut bank_forks_w = bank_forks.write().unwrap();
-        let mut bank =
-            Bank::new_from_parent(bank_forks_w[slot - 1].clone(), &Pubkey::default(), slot);
+        let mut bank = Bank::new_from_parent(
+            bank_forks.read().unwrap().get(slot - 1).unwrap().clone(),
+            &Pubkey::default(),
+            slot,
+        );
         f(&mut bank, mint_keypair);
-        let bank = bank_forks_w.insert(bank);
-        drop(bank_forks_w);
+        let bank = bank_forks.write().unwrap().insert(bank);
         // Set root to make sure we don't end up with too many account storage entries
         // and to allow snapshotting of bank and the purging logic on status_cache to
         // kick in
@@ -352,7 +353,7 @@ fn test_concurrent_snapshot_packaging(
         DISABLED_SNAPSHOT_ARCHIVE_INTERVAL,
     );
 
-    let mut bank_forks = snapshot_test_config.bank_forks.write().unwrap();
+    let bank_forks = snapshot_test_config.bank_forks.clone();
     let snapshot_config = &snapshot_test_config.snapshot_config;
     let bank_snapshots_dir = &snapshot_config.bank_snapshots_dir;
     let full_snapshot_archives_dir = &snapshot_config.full_snapshot_archives_dir;
@@ -361,7 +362,7 @@ fn test_concurrent_snapshot_packaging(
     let genesis_config = &snapshot_test_config.genesis_config_info.genesis_config;
 
     // Take snapshot of zeroth bank
-    let bank0 = bank_forks.get(0).unwrap();
+    let bank0 = bank_forks.read().unwrap().get(0).unwrap();
     let storages = bank0.get_snapshot_storages(None);
     let slot_deltas = bank0.status_cache.read().unwrap().root_slot_deltas();
     snapshot_bank_utils::add_bank_snapshot(
@@ -394,7 +395,7 @@ fn test_concurrent_snapshot_packaging(
     for i in 0..MAX_BANK_SNAPSHOTS_TO_RETAIN + 2 {
         let parent_slot = i as Slot;
         let bank = Bank::new_from_parent(
-            bank_forks[parent_slot].clone(),
+            bank_forks.read().unwrap().get(parent_slot).unwrap().clone(),
             &Pubkey::default(),
             parent_slot + 1,
         );
@@ -438,10 +439,14 @@ fn test_concurrent_snapshot_packaging(
         );
         accounts_package_sender.send(accounts_package).unwrap();
 
-        bank_forks.insert(bank);
+        bank_forks.write().unwrap().insert(bank);
         if slot == saved_slot {
             // Find the relevant snapshot storages
-            let snapshot_storage_files: HashSet<_> = bank_forks[slot]
+            let snapshot_storage_files: HashSet<_> = bank_forks
+                .read()
+                .unwrap()
+                .get(slot)
+                .unwrap()
                 .get_snapshot_storages(None)
                 .into_iter()
                 .map(|s| s.get_path())

program-runtime/src/loaded_programs.rs

@@ -53,6 +53,11 @@ pub enum BlockRelation {
 pub trait ForkGraph {
     /// Returns the BlockRelation of A to B
     fn relationship(&self, a: Slot, b: Slot) -> BlockRelation;
+
+    /// Returns the epoch of the given slot
+    fn slot_epoch(&self, _slot: Slot) -> Option<Epoch> {
+        Some(0)
+    }
 }
 
 /// Provides information about current working slot, and its ancestors
@@ -759,14 +764,29 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
         let environments = self.get_environments_for_epoch(working_slot.current_epoch());
         let mut missing = Vec::new();
         let mut unloaded = Vec::new();
+        let current_slot = working_slot.current_slot();
         let found = keys
             .filter_map(|(key, (match_criteria, count))| {
                 if let Some(second_level) = self.entries.get(&key) {
                     for entry in second_level.iter().rev() {
-                        let current_slot = working_slot.current_slot();
+                        let is_ancestor = if let Some(fork_graph) = &self.fork_graph {
+                            fork_graph
+                                .read()
+                                .map(|fork_graph_r| {
+                                    matches!(
+                                        fork_graph_r
+                                            .relationship(entry.deployment_slot, current_slot),
+                                        BlockRelation::Ancestor
+                                    )
+                                })
+                                .unwrap_or(false)
+                        } else {
+                            working_slot.is_ancestor(entry.deployment_slot)
+                        };
+
                         if entry.deployment_slot <= self.latest_root_slot
                             || entry.deployment_slot == current_slot
-                            || working_slot.is_ancestor(entry.deployment_slot)
+                            || is_ancestor
                         {
                             if current_slot >= entry.effective_slot {
                                 if !Self::is_entry_usable(entry, current_slot, &match_criteria) {
@@ -818,7 +838,7 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
         ExtractedPrograms {
             loaded: LoadedProgramsForTxBatch {
                 entries: found,
-                slot: working_slot.current_slot(),
+                slot: current_slot,
                 environments: environments.clone(),
             },
             missing,
@@ -1488,52 +1508,19 @@ mod tests {
         }
     }
 
-    struct TestWorkingSlot {
-        slot: Slot,
-        fork: Vec<Slot>,
-        slot_pos: usize,
-    }
-
-    impl TestWorkingSlot {
-        fn new(slot: Slot, fork: &[Slot]) -> Self {
-            let mut fork = fork.to_vec();
-            fork.sort();
-            let slot_pos = fork
-                .iter()
-                .position(|current| *current == slot)
-                .expect("The fork didn't have the slot in it");
-            TestWorkingSlot {
-                slot,
-                fork,
-                slot_pos,
-            }
-        }
-
-        fn update_slot(&mut self, slot: Slot) {
-            self.slot = slot;
-            self.slot_pos = self
-                .fork
-                .iter()
-                .position(|current| *current == slot)
-                .expect("The fork didn't have the slot in it");
-        }
-    }
+    struct TestWorkingSlot(pub Slot);
 
     impl WorkingSlot for TestWorkingSlot {
         fn current_slot(&self) -> Slot {
-            self.slot
+            self.0
         }
 
         fn current_epoch(&self) -> Epoch {
             0
         }
 
-        fn is_ancestor(&self, other: Slot) -> bool {
-            self.fork
-                .iter()
-                .position(|current| *current == other)
-                .map(|other_pos| other_pos < self.slot_pos)
-                .unwrap_or(false)
+        fn is_ancestor(&self, _other: Slot) -> bool {
+            false
         }
     }
 
@@ -1571,7 +1558,7 @@ mod tests {
 
         let mut fork_graph = TestForkGraphSpecific::default();
         fork_graph.insert_fork(&[0, 10, 20, 22]);
-        fork_graph.insert_fork(&[0, 5, 11, 15, 16, 19, 21, 23]);
+        fork_graph.insert_fork(&[0, 5, 11, 15, 16, 18, 19, 21, 23]);
         fork_graph.insert_fork(&[0, 5, 11, 25, 27]);
 
         let fork_graph = Arc::new(RwLock::new(fork_graph));
@@ -1628,13 +1615,12 @@ mod tests {
         //                     23
 
         // Testing fork 0 - 10 - 12 - 22 with current slot at 22
-        let working_slot = TestWorkingSlot::new(22, &[0, 10, 20, 22]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(22),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 2)),
@@ -1651,14 +1637,13 @@ mod tests {
         assert!(missing.contains(&(program3, 3)));
         assert!(unloaded.is_empty());
 
-        // Testing fork 0 - 5 - 11 - 15 - 16 with current slot at 16
-        let mut working_slot = TestWorkingSlot::new(15, &[0, 5, 11, 15, 16, 18, 19, 23]);
+        // Testing fork 0 - 5 - 11 - 15 - 16 with current slot at 15
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(15),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1681,13 +1666,12 @@ mod tests {
         assert!(unloaded.is_empty());
 
         // Testing the same fork above, but current slot is now 18 (equal to effective slot of program4).
-        working_slot.update_slot(18);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(18),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1707,13 +1691,12 @@ mod tests {
         assert!(unloaded.is_empty());
 
         // Testing the same fork above, but current slot is now 23 (future slot than effective slot of program4).
-        working_slot.update_slot(23);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(23),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1733,13 +1716,12 @@ mod tests {
         assert!(unloaded.is_empty());
 
         // Testing fork 0 - 5 - 11 - 15 - 16 with current slot at 11
-        let working_slot = TestWorkingSlot::new(11, &[0, 5, 11, 15, 16]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(11),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1772,13 +1754,12 @@ mod tests {
         assert!(!cache.replenish(program4, test_program).0);
 
         // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
-        let working_slot = TestWorkingSlot::new(19, &[0, 5, 11, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(19),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1798,13 +1779,12 @@ mod tests {
 
         // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 21
         // This would cause program4 deployed at slot 19 to be expired.
-        let working_slot = TestWorkingSlot::new(21, &[0, 5, 11, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(21),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1843,14 +1823,13 @@ mod tests {
         //                   |
         //                  23
 
-        // Testing fork 11 - 15 - 16- 19 - 22 with root at 5 and current slot at 22
-        let working_slot = TestWorkingSlot::new(22, &[5, 11, 15, 16, 19, 22, 23]);
+        // Testing fork 11 - 15 - 16- 19 - 22 with root at 5 and current slot at 21
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(21),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1861,21 +1840,20 @@ mod tests {
         );
 
         // Since the fork was pruned, we should not find the entry deployed at slot 20.
-        assert!(match_slot(&found, &program1, 0, 22));
-        assert!(match_slot(&found, &program2, 11, 22));
-        assert!(match_slot(&found, &program4, 15, 22));
+        assert!(match_slot(&found, &program1, 0, 21));
+        assert!(match_slot(&found, &program2, 11, 21));
+        assert!(match_slot(&found, &program4, 15, 21));
 
         assert!(missing.contains(&(program3, 1)));
         assert!(unloaded.is_empty());
 
         // Testing fork 0 - 5 - 11 - 25 - 27 with current slot at 27
-        let working_slot = TestWorkingSlot::new(27, &[11, 25, 27]);
         let ExtractedPrograms {
             loaded: found,
             missing: _,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(27),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1909,13 +1887,12 @@ mod tests {
         //                  23
 
         // Testing fork 16, 19, 23, with root at 15, current slot at 23
-        let working_slot = TestWorkingSlot::new(23, &[16, 19, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(23),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -1955,7 +1932,7 @@ mod tests {
 
         let mut fork_graph = TestForkGraphSpecific::default();
         fork_graph.insert_fork(&[0, 10, 20, 22]);
-        fork_graph.insert_fork(&[0, 5, 11, 15, 16, 19, 21, 23]);
+        fork_graph.insert_fork(&[0, 5, 11, 12, 15, 16, 18, 19, 21, 23]);
         fork_graph.insert_fork(&[0, 5, 11, 25, 27]);
 
         let fork_graph = Arc::new(RwLock::new(fork_graph));
@@ -1973,13 +1950,12 @@ mod tests {
         assert!(!cache.replenish(program3, new_test_loaded_program(25, 26)).0);
 
         // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
-        let working_slot = TestWorkingSlot::new(12, &[0, 5, 11, 12, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(12),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2000,7 +1976,7 @@ mod tests {
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(12),
             vec![
                 (
                     program1,
@@ -2078,13 +2054,12 @@ mod tests {
         );
 
         // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
-        let working_slot = TestWorkingSlot::new(19, &[0, 5, 11, 12, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(19),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2100,13 +2075,12 @@ mod tests {
         assert!(unloaded.is_empty());
 
         // Testing fork 0 - 5 - 11 - 25 - 27 with current slot at 27
-        let working_slot = TestWorkingSlot::new(27, &[0, 5, 11, 25, 27]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(27),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2122,13 +2096,12 @@ mod tests {
         assert!(missing.is_empty());
 
         // Testing fork 0 - 10 - 20 - 22 with current slot at 22
-        let working_slot = TestWorkingSlot::new(22, &[0, 10, 20, 22]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(22),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2164,7 +2137,7 @@ mod tests {
 
         let mut fork_graph = TestForkGraphSpecific::default();
         fork_graph.insert_fork(&[0, 10, 20, 22]);
-        fork_graph.insert_fork(&[0, 5, 11, 15, 16, 19, 21, 23]);
+        fork_graph.insert_fork(&[0, 5, 11, 12, 15, 16, 18, 19, 21, 23]);
         fork_graph.insert_fork(&[0, 5, 11, 25, 27]);
         let fork_graph = Arc::new(RwLock::new(fork_graph));
         cache.set_fork_graph(fork_graph);
@@ -2193,13 +2166,12 @@ mod tests {
         assert!(!cache.replenish(program1, test_program).0);
 
         // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
-        let working_slot = TestWorkingSlot::new(12, &[0, 5, 11, 12, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(12),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2217,13 +2189,12 @@ mod tests {
 
         // Testing fork 0 - 5 - 11 - 12 - 15 - 16 - 19 - 21 - 23 with current slot at 15
         // This would cause program4 deployed at slot 15 to be expired.
-        let working_slot = TestWorkingSlot::new(15, &[0, 5, 11, 15, 16, 18, 19, 21, 23]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(15),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2290,13 +2261,12 @@ mod tests {
 
         cache.prune(10, 0);
 
-        let working_slot = TestWorkingSlot::new(20, &[0, 10, 20]);
         let ExtractedPrograms {
             loaded: found,
             missing: _,
             unloaded,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(20),
             vec![(program1, (LoadedProgramMatchCriteria::NoCriteria, 1))].into_iter(),
         );
         assert!(unloaded.is_empty());
@@ -2328,7 +2298,7 @@ mod tests {
         // deployed at slot 0.
         let mut fork_graph = TestForkGraphSpecific::default();
         fork_graph.insert_fork(&[0, 10, 20]);
-        fork_graph.insert_fork(&[0, 5]);
+        fork_graph.insert_fork(&[0, 5, 6]);
         let fork_graph = Arc::new(RwLock::new(fork_graph));
         cache.set_fork_graph(fork_graph);
 
@@ -2339,13 +2309,12 @@ mod tests {
         let program2 = Pubkey::new_unique();
         assert!(!cache.replenish(program2, new_test_loaded_program(10, 11)).0);
 
-        let working_slot = TestWorkingSlot::new(20, &[0, 10, 20]);
         let ExtractedPrograms {
             loaded: found,
             missing: _,
             unloaded: _,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(20),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2356,13 +2325,12 @@ mod tests {
         assert!(match_slot(&found, &program1, 0, 20));
         assert!(match_slot(&found, &program2, 10, 20));
 
-        let working_slot = TestWorkingSlot::new(6, &[0, 5, 6]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded: _,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(6),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2377,13 +2345,12 @@ mod tests {
         // On fork chaining from slot 5, the entry deployed at slot 0 will become visible.
         cache.prune_by_deployment_slot(5);
 
-        let working_slot = TestWorkingSlot::new(20, &[0, 10, 20]);
         let ExtractedPrograms {
             loaded: found,
             missing: _,
             unloaded: _,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(20),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2394,13 +2361,12 @@ mod tests {
         assert!(match_slot(&found, &program1, 0, 20));
         assert!(match_slot(&found, &program2, 10, 20));
 
-        let working_slot = TestWorkingSlot::new(6, &[0, 5, 6]);
         let ExtractedPrograms {
             loaded: found,
             missing,
             unloaded: _,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(6),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),
@@ -2415,13 +2381,12 @@ mod tests {
         // As there is no other entry for program2, extract() will return it as missing.
         cache.prune_by_deployment_slot(10);
 
-        let working_slot = TestWorkingSlot::new(20, &[0, 10, 20]);
         let ExtractedPrograms {
             loaded: found,
             missing: _,
             unloaded: _,
         } = cache.extract(
-            &working_slot,
+            &TestWorkingSlot(20),
             vec![
                 (program1, (LoadedProgramMatchCriteria::NoCriteria, 1)),
                 (program2, (LoadedProgramMatchCriteria::NoCriteria, 1)),

rpc/src/rpc_pubsub.rs

@@ -659,7 +659,7 @@ mod tests {
         current_slot: Slot,
     ) -> transaction::Result<()> {
         bank_forks
-            .write()
+            .read()
             .unwrap()
             .get(current_slot)
             .unwrap()
@@ -1166,7 +1166,7 @@ mod tests {
 
         let tx = system_transaction::transfer(&alice, &bob.pubkey(), 100, blockhash);
         bank_forks
-            .write()
+            .read()
             .unwrap()
             .get(1)
             .unwrap()
@@ -1221,7 +1221,7 @@ mod tests {
 
         let tx = system_transaction::transfer(&alice, &bob.pubkey(), 100, blockhash);
         bank_forks
-            .write()
+            .read()
             .unwrap()
             .get(1)
             .unwrap()

rpc/src/rpc_subscriptions.rs

@@ -1838,7 +1838,7 @@ pub(crate) mod tests {
             &stake::program::id(),
         );
         bank_forks
-            .write()
+            .read()
             .unwrap()
             .get(0)
             .unwrap()
@@ -2436,8 +2436,7 @@ pub(crate) mod tests {
         } = create_genesis_config(100);
         let bank = Bank::new_for_tests(&genesis_config);
         let blockhash = bank.last_blockhash();
-        let bank_forks_arc = BankForks::new_rw_arc(bank);
-        let mut bank_forks = bank_forks_arc.write().unwrap();
+        let bank_forks = BankForks::new_rw_arc(bank);
         let alice = Keypair::new();
 
         let past_bank_tx =
@@ -2448,24 +2447,28 @@ pub(crate) mod tests {
             system_transaction::transfer(&mint_keypair, &alice.pubkey(), 3, blockhash);
 
         bank_forks
+            .read()
+            .unwrap()
             .get(0)
             .unwrap()
             .process_transaction(&past_bank_tx)
             .unwrap();
 
         let next_bank = Bank::new_from_parent(
-            bank_forks.get(0).unwrap(),
+            bank_forks.read().unwrap().get(0).unwrap(),
             &solana_sdk::pubkey::new_rand(),
             1,
         );
-        bank_forks.insert(next_bank);
+        bank_forks.write().unwrap().insert(next_bank);
 
         bank_forks
+            .read()
+            .unwrap()
             .get(1)
             .unwrap()
             .process_transaction(&processed_tx)
             .unwrap();
-        let bank1 = bank_forks[1].clone();
+        let bank1 = bank_forks.read().unwrap().get(1).unwrap().clone();
 
         let mut cache0 = BlockCommitment::default();
         cache0.increase_confirmation_stake(1, 10);
@@ -2483,19 +2486,16 @@ pub(crate) mod tests {
             },
         );
 
-        // Drop the write locked bank_forks
-        drop(bank_forks);
-
         let exit = Arc::new(AtomicBool::new(false));
         let optimistically_confirmed_bank =
-            OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks_arc);
+            OptimisticallyConfirmedBank::locked_from_bank_forks_root(&bank_forks);
         let max_complete_transaction_status_slot = Arc::new(AtomicU64::default());
         let max_complete_rewards_slot = Arc::new(AtomicU64::default());
         let subscriptions = Arc::new(RpcSubscriptions::new_for_tests(
             exit,
             max_complete_transaction_status_slot,
             max_complete_rewards_slot,
-            bank_forks_arc,
+            bank_forks,
             Arc::new(RwLock::new(block_commitment_cache)),
             optimistically_confirmed_bank,
         ));
@@ -2818,7 +2818,7 @@ pub(crate) mod tests {
 
         // Add the same transaction to the unfrozen 2nd bank
         bank_forks
-            .write()
+            .read()
             .unwrap()
             .get(2)
             .unwrap()

runtime/src/bank.rs

@@ -942,7 +942,6 @@ impl WorkingSlot for Bank {
         self.ancestors.contains_key(&other)
     }
 }
-
 #[derive(Debug, Default)]
 /// result of calculating the stake rewards at end of epoch
 struct StakeRewardCalculation {

runtime/src/bank_forks.rs

@@ -9,8 +9,13 @@ use {
     },
     log::*,
     solana_measure::measure::Measure,
-    solana_program_runtime::loaded_programs::{BlockRelation, ForkGraph, WorkingSlot},
-    solana_sdk::{clock::Slot, feature_set, hash::Hash, timing},
+    solana_program_runtime::loaded_programs::{BlockRelation, ForkGraph},
+    solana_sdk::{
+        clock::{Epoch, Slot},
+        feature_set,
+        hash::Hash,
+        timing,
+    },
     std::{
         collections::{hash_map::Entry, HashMap, HashSet},
         ops::Index,
@@ -680,9 +685,11 @@ impl ForkGraph for BankForks {
                 (a == b)
                     .then_some(BlockRelation::Equal)
                     .or_else(|| {
-                        self.banks
-                            .get(&b)
-                            .and_then(|bank| bank.is_ancestor(a).then_some(BlockRelation::Ancestor))
+                        self.banks.get(&b).and_then(|bank| {
+                            bank.ancestors
+                                .contains_key(&a)
+                                .then_some(BlockRelation::Ancestor)
+                        })
                     })
                     .or_else(|| {
                         self.descendants.get(&b).and_then(|slots| {
@@ -693,6 +700,10 @@ impl ForkGraph for BankForks {
             })
             .unwrap_or(BlockRelation::Unknown)
     }
+
+    fn slot_epoch(&self, slot: Slot) -> Option<Epoch> {
+        self.banks.get(&slot).map(|bank| bank.epoch())
+    }
 }
 
 #[cfg(test)]