From 372c496ce9378bf4c4396762dfa8215cce610b5c Mon Sep 17 00:00:00 2001
From: Andrei Silviu Dragnea <andreidragnea.dev@gmail.com>
Date: Thu, 27 Jun 2024 11:45:19 +0300
Subject: [PATCH] Fix BankForks::new_rw_arc memory leak

---
 program-runtime/src/loaded_programs.rs | 12 ++++++++----
 runtime/src/bank_forks.rs              | 10 +++++++++-
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
index b9c7d97cdbc..f398091fc93 100644
--- a/program-runtime/src/loaded_programs.rs
+++ b/program-runtime/src/loaded_programs.rs
@@ -25,7 +25,7 @@ use {
         fmt::{Debug, Formatter},
         sync::{
             atomic::{AtomicU64, Ordering},
-            Arc, Condvar, Mutex, RwLock,
+            Arc, Condvar, Mutex, RwLock, Weak,
         },
     },
 };
@@ -549,7 +549,7 @@ pub struct LoadedPrograms<FG: ForkGraph> {
     /// List of loaded programs which should be recompiled before the next epoch (but don't have to).
     pub programs_to_recompile: Vec<(Pubkey, Arc<LoadedProgram>)>,
     pub stats: Stats,
-    pub fork_graph: Option<Arc<RwLock<FG>>>,
+    pub fork_graph: Option<Weak<RwLock<FG>>>,
     pub loading_task_waiter: Arc<LoadingTaskWaiter>,
 }
 
@@ -691,7 +691,7 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
         }
     }
 
-    pub fn set_fork_graph(&mut self, fork_graph: Arc<RwLock<FG>>) {
+    pub fn set_fork_graph(&mut self, fork_graph: Weak<RwLock<FG>>) {
         self.fork_graph = Some(fork_graph);
     }
 
@@ -797,6 +797,7 @@ impl<FG: ForkGraph> LoadedPrograms<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;
@@ -929,7 +930,8 @@ impl<FG: ForkGraph> LoadedPrograms<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;
         search_for.retain(|(key, (match_criteria, usage_count))| {
             if let Some(second_level) = self.entries.get_mut(key) {
@@ -1031,6 +1033,8 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
                 self.fork_graph
                     .as_ref()
                     .unwrap()
+                    .upgrade()
+                    .unwrap()
                     .read()
                     .unwrap()
                     .relationship(loaded_program.deployment_slot, slot),
diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs
index d481bf1b43b..a0c33fd923e 100644
--- a/runtime/src/bank_forks.rs
+++ b/runtime/src/bank_forks.rs
@@ -130,7 +130,7 @@ impl BankForks {
             .loaded_programs_cache
             .write()
             .unwrap()
-            .set_fork_graph(bank_forks.clone());
+            .set_fork_graph(Arc::downgrade(&bank_forks));
 
         bank_forks
     }
@@ -734,6 +734,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);