From c90fa6643e542ad9179c3a4867dbc830b7c39a7e Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Thu, 6 Jan 2022 19:11:05 +0000
Subject: [PATCH] Retain executor cache counts (#22322) (#22340)

(cherry picked from commit f1e2598baa80a0ad4e8450c8b5e3c5ab164f501c)

Co-authored-by: Jack May <jack@solana.com>
---
 runtime/src/bank.rs | 46 ++++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 17 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 75c7dddf1bdbd0..82b4abb78982d1 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -429,37 +429,43 @@ impl CachedExecutors {
             executors: HashMap::new(),
         }
     }
+
     fn get(&self, pubkey: &Pubkey) -> Option<Arc<dyn Executor>> {
         self.executors.get(pubkey).map(|entry| {
             entry.epoch_count.fetch_add(1, Relaxed);
             entry.executor.clone()
         })
     }
+
     fn put(&mut self, pubkey: &Pubkey, executor: Arc<dyn Executor>) {
-        if !self.executors.contains_key(pubkey) && self.executors.len() >= self.max {
-            let mut least = u64::MAX;
-            let default_key = Pubkey::default();
-            let mut least_key = &default_key;
-
-            for (key, entry) in self.executors.iter() {
-                let count = entry.prev_epoch_count + entry.epoch_count.load(Relaxed);
-                if count < least {
-                    least = count;
-                    least_key = key;
+        let entry = if let Some(mut entry) = self.executors.remove(pubkey) {
+            entry.executor = executor;
+            entry
+        } else {
+            if self.executors.len() >= self.max {
+                let mut least = u64::MAX;
+                let default_key = Pubkey::default();
+                let mut least_key = &default_key;
+
+                for (key, entry) in self.executors.iter() {
+                    let count = entry.prev_epoch_count + entry.epoch_count.load(Relaxed);
+                    if count < least {
+                        least = count;
+                        least_key = key;
+                    }
                 }
+                let least_key = *least_key;
+                let _ = self.executors.remove(&least_key);
             }
-            let least_key = *least_key;
-            let _ = self.executors.remove(&least_key);
-        }
-        let _ = self.executors.insert(
-            *pubkey,
             CachedExecutorsEntry {
                 prev_epoch_count: 0,
                 epoch_count: AtomicU64::new(0),
                 executor,
-            },
-        );
+            }
+        };
+        let _ = self.executors.insert(*pubkey, entry);
     }
+
     fn remove(&mut self, pubkey: &Pubkey) {
         let _ = self.executors.remove(pubkey);
     }
@@ -11917,6 +11923,11 @@ pub(crate) mod tests {
         assert!(cache.get(&key4).is_some());
         assert!(cache.get(&key3).is_none());
 
+        cache.put(&key1, executor.clone());
+        cache.put(&key3, executor.clone());
+        assert!(cache.get(&key1).is_some());
+        assert!(cache.get(&key4).is_none());
+
         cache = cache.clone_with_epoch(2);
         assert!(cache.current_epoch == 2);
 
@@ -11992,6 +12003,7 @@ pub(crate) mod tests {
         assert!(executors.borrow().contains_key(&key3));
         assert!(executors.borrow().contains_key(&key4));
 
+        // Remove all
         bank.remove_executor(&key1);
         bank.remove_executor(&key2);
         bank.remove_executor(&key3);