Retain executor cache counts (#22322)

runtime/src/bank.rs

@@ -367,37 +367,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() {
+    fn put(&mut self, pubkey: &Pubkey, executor: Arc<dyn Executor>) {
-                let count = entry.prev_epoch_count + entry.epoch_count.load(Relaxed);
+        let entry = if let Some(mut entry) = self.executors.remove(pubkey) {
-                if count < least {
+            entry.executor = executor;
-                    least = count;
+            entry
-                    least_key = key;
+        } 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);
     }
@@ -12791,6 +12797,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);
 
@@ -12871,6 +12882,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);