Fix BankForks::new_rw_arc memory leak

anza-xyz · Jul 2, 2024 · 2f0a0fd · 2f0a0fd
1 parent f8630a3
commit 2f0a0fd
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 10 deletions.
diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
@@ -30,6 +30,7 @@ use {
     std::{
         collections::{hash_map::Entry, HashMap},
         fmt::{Debug, Formatter},
+        sync::Weak,
     },
 };
 
@@ -640,7 +641,7 @@ pub struct ProgramCache<FG: ForkGraph> {
     /// Statistics counters
     pub stats: ProgramCacheStats,
     /// Reference to the block store
-    pub fork_graph: Option<Arc<RwLock<FG>>>,
+    pub fork_graph: Option<Weak<RwLock<FG>>>,
     /// Coordinates TX batches waiting for others to complete their task during cooperative loading
     pub loading_task_waiter: Arc<LoadingTaskWaiter>,
 }
@@ -826,7 +827,7 @@ impl<FG: ForkGraph> ProgramCache<FG> {
     }
 
     pub fn set_fork_graph(&mut self, fork_graph: Arc<RwLock<FG>>) {
-        self.fork_graph = Some(fork_graph);
+        self.fork_graph = Some(Arc::downgrade(&fork_graph));
     }
 
     /// Returns the current environments depending on the given epoch
@@ -948,6 +949,7 @@ impl<FG: ForkGraph> ProgramCache<FG> {
             error!("Program cache doesn't have fork graph.");
             return;
         };
+        let fork_graph = fork_graph.upgrade().unwrap();
         let Ok(fork_graph) = fork_graph.read() else {
             error!("Failed to lock fork graph for reading.");
             return;
@@ -1059,7 +1061,8 @@ impl<FG: ForkGraph> ProgramCache<FG> {
         is_first_round: bool,
     ) -> Option<(Pubkey, u64)> {
         debug_assert!(self.fork_graph.is_some());
-        let locked_fork_graph = self.fork_graph.as_ref().unwrap().read().unwrap();
+        let fork_graph = self.fork_graph.as_ref().unwrap().upgrade().unwrap();
+        let locked_fork_graph = fork_graph.read().unwrap();
         let mut cooperative_loading_task = None;
         match &self.index {
             IndexImplementation::V1 {
@@ -1166,6 +1169,8 @@ impl<FG: ForkGraph> ProgramCache<FG> {
                         self.fork_graph
                             .as_ref()
                             .unwrap()
+                            .upgrade()
+                            .unwrap()
                             .read()
                             .unwrap()
                             .relationship(loaded_program.deployment_slot, slot),
@@ -2171,7 +2176,8 @@ mod tests {
         loading_slot: Slot,
         keys: &[Pubkey],
     ) -> Vec<(Pubkey, (ProgramCacheMatchCriteria, u64))> {
-        let locked_fork_graph = cache.fork_graph.as_ref().unwrap().read().unwrap();
+        let fork_graph = cache.fork_graph.as_ref().unwrap().upgrade().unwrap();
+        let locked_fork_graph = fork_graph.read().unwrap();
         let entries = cache.get_flattened_entries_for_tests();
         keys.iter()
             .filter_map(|key| {

diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs
@@ -751,6 +751,14 @@ mod tests {
         std::{sync::atomic::Ordering::Relaxed, time::Duration},
     };
 
+    #[test]
+    fn test_bank_forks_new_rw_arc_memory_leak() {
+        for _ in 0..1000 {
+            let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
+            BankForks::new_rw_arc(Bank::new_for_tests(&genesis_config));
+        }
+    }
+
     #[test]
     fn test_bank_forks_new() {
         let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);

diff --git a/svm/src/transaction_processor.rs b/svm/src/transaction_processor.rs
@@ -1305,8 +1305,9 @@ mod tests {
     fn test_replenish_program_cache_with_nonexistent_accounts() {
         let mock_bank = MockBankCallback::default();
         let batch_processor = TransactionBatchProcessor::<TestForkGraph>::default();
+        let fork_graph = Arc::new(RwLock::new(TestForkGraph {}));
         batch_processor.program_cache.write().unwrap().fork_graph =
-            Some(Arc::new(RwLock::new(TestForkGraph {})));
+            Some(Arc::downgrade(&fork_graph));
         let key = Pubkey::new_unique();
 
         let mut account_maps: HashMap<Pubkey, u64> = HashMap::new();
@@ -1319,8 +1320,9 @@ mod tests {
     fn test_replenish_program_cache() {
         let mock_bank = MockBankCallback::default();
         let batch_processor = TransactionBatchProcessor::<TestForkGraph>::default();
+        let fork_graph = Arc::new(RwLock::new(TestForkGraph {}));
         batch_processor.program_cache.write().unwrap().fork_graph =
-            Some(Arc::new(RwLock::new(TestForkGraph {})));
+            Some(Arc::downgrade(&fork_graph));
         let key = Pubkey::new_unique();
 
         let mut account_data = AccountSharedData::default();
@@ -1800,8 +1802,9 @@ mod tests {
     fn test_add_builtin() {
         let mock_bank = MockBankCallback::default();
         let batch_processor = TransactionBatchProcessor::<TestForkGraph>::default();
+        let fork_graph = Arc::new(RwLock::new(TestForkGraph {}));
         batch_processor.program_cache.write().unwrap().fork_graph =
-            Some(Arc::new(RwLock::new(TestForkGraph {})));
+            Some(Arc::downgrade(&fork_graph));
 
         let key = Pubkey::new_unique();
         let name = "a_builtin_name";
@@ -1843,8 +1846,9 @@ mod tests {
     fn fast_concur_test() {
         let mut mock_bank = MockBankCallback::default();
         let batch_processor = TransactionBatchProcessor::<TestForkGraph>::new(5, 5, HashSet::new());
+        let fork_graph = Arc::new(RwLock::new(TestForkGraph {}));
         batch_processor.program_cache.write().unwrap().fork_graph =
-            Some(Arc::new(RwLock::new(TestForkGraph {})));
+            Some(Arc::downgrade(&fork_graph));
 
         let programs = vec![
             deploy_program("hello-solana".to_string(), &mut mock_bank),

diff --git a/svm/tests/conformance.rs b/svm/tests/conformance.rs
@@ -254,7 +254,8 @@ fn run_fixture(fixture: InstrFixture, filename: OsString, execute_as_instr: bool
                 FunctionRegistry::default(),
             )),
         };
-        program_cache.fork_graph = Some(Arc::new(RwLock::new(MockForkGraph {})));
+        let fork_graph = Arc::new(RwLock::new(MockForkGraph {}));
+        program_cache.fork_graph = Some(Arc::downgrade(&fork_graph));
     }
 
     batch_processor.fill_missing_sysvar_cache_entries(&mock_bank);

diff --git a/svm/tests/integration_test.rs b/svm/tests/integration_test.rs
@@ -142,7 +142,8 @@ fn create_executable_environment(
         )),
     };
 
-    program_cache.fork_graph = Some(Arc::new(RwLock::new(MockForkGraph {})));
+    let fork_graph = Arc::new(RwLock::new(MockForkGraph {}));
+    program_cache.fork_graph = Some(Arc::downgrade(&fork_graph));
 
     // We must fill in the sysvar cache entries
     let time_now = SystemTime::now()