From 7602911f9a39905ecf54fef6a0416bb7d0b0962c Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Fri, 23 Apr 2021 12:00:17 -0500
Subject: [PATCH] read_only_cache remove updates cache size (#16742)

---
 runtime/src/read_only_accounts_cache.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/runtime/src/read_only_accounts_cache.rs b/runtime/src/read_only_accounts_cache.rs
index 71a0947158..bd708c5b9e 100644
--- a/runtime/src/read_only_accounts_cache.rs
+++ b/runtime/src/read_only_accounts_cache.rs
@@ -143,9 +143,10 @@ impl ReadOnlyAccountsCache {
     }
 
     pub fn remove(&self, pubkey: &Pubkey, slot: Slot) {
-        // does not keep track of data size reduction here.
-        // data size will be recomputed the next time we store and we think we may now be too large.
-        self.cache.remove(&(*pubkey, slot));
+        if let Some((_, value)) = self.cache.remove(&(*pubkey, slot)) {
+            self.data_size
+                .fetch_sub(self.account_size(&value.account), Ordering::Relaxed);
+        }
     }
 
     fn purge_lru_list(&self, lru: &[LruEntry], lru_index: &mut usize) -> bool {
@@ -294,7 +295,7 @@ pub mod tests {
         assert!(accounts_equal(&cache.load(&key2, slot).unwrap(), &account1));
         assert_eq!(1, cache.cache_len());
         cache.remove(&key2, slot);
-        assert_eq!(100 + per_account_size, cache.data_size());
+        assert_eq!(0, cache.data_size());
         assert_eq!(0, cache.cache_len());
 
         // can store 2 items, 3rd item kicks oldest item out