From cb55820f6638fd1a3ef488baedb6e95e2aa1eedc Mon Sep 17 00:00:00 2001
From: Kiran Prakash <awskiran@amazon.com>
Date: Wed, 13 Mar 2024 14:37:26 -0700
Subject: [PATCH] add updateCleanupKeyToCountMapOnCacheEviction

Signed-off-by: Kiran Prakash <awskiran@amazon.com>
---
 .../indices/IndicesRequestCache.java          | 33 +++++++++++++++++--
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java b/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java
index ce576b225094d..3870e7a33b8d2 100644
--- a/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java
+++ b/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java
@@ -193,7 +193,11 @@ void clear(CacheEntity entity) {
     public void onRemoval(RemovalNotification<Key, BytesReference> notification) {
         // In case this event happens for an old shard, we can safely ignore this as we don't keep track for old
         // shards as part of request cache.
-        cacheEntityLookup.apply(notification.getKey().shardId).ifPresent(entity -> entity.onRemoval(notification));
+        Key key = notification.getKey();
+        cacheEntityLookup.apply(key.shardId).ifPresent(entity -> entity.onRemoval(notification));
+        cacheCleanupManager.updateCleanupKeyToCountMapOnCacheEviction(
+            new CleanupKey(cacheEntityLookup.apply(key.shardId).orElse(null), key.readerCacheKeyId)
+        );
     }
 
     BytesReference getOrCompute(
@@ -222,7 +226,7 @@ BytesReference getOrCompute(
                     OpenSearchDirectoryReader.addReaderCloseListener(reader, cleanupKey);
                 }
             }
-            cacheCleanupManager.updateCleanupKeyToCountMap(cleanupKey);
+            cacheCleanupManager.updateCleanupKeyToCountMapOnCacheInsertion(cleanupKey);
         } else {
             cacheEntity.onHit();
         }
@@ -463,7 +467,7 @@ void enqueueCleanupKey(CleanupKey cleanupKey) {
          *
          * @param cleanupKey the CleanupKey to be updated in the map
          */
-        private void updateCleanupKeyToCountMap(CleanupKey cleanupKey) {
+        private void updateCleanupKeyToCountMapOnCacheInsertion(CleanupKey cleanupKey) {
             if (stalenessThreshold == 0.0) {
                 return;
             }
@@ -480,6 +484,29 @@ private void updateCleanupKeyToCountMap(CleanupKey cleanupKey) {
                 .merge(cleanupKey.readerCacheKeyId, 1, Integer::sum);
         }
 
+        private void updateCleanupKeyToCountMapOnCacheEviction(CleanupKey cleanupKey) {
+            if (stalenessThreshold == 0.0) {
+                return;
+            }
+            IndexShard indexShard = (IndexShard) cleanupKey.entity.getCacheIdentity();
+            if (indexShard == null) {
+                logger.warn("IndexShard is null for CleanupKey: {} while cleaning Indices Request Cache", cleanupKey.readerCacheKeyId);
+                return;
+            }
+            ShardId shardId = indexShard.shardId();
+
+            writeLock.lock();
+            try {
+                // If the key doesn't exist, ignore
+                ConcurrentMap<String, Integer> keyCountMap = cleanupKeyToCountMap.get(shardId);
+                if (keyCountMap != null) {
+                    keyCountMap.computeIfPresent(cleanupKey.readerCacheKeyId, (key, value) -> value - 1);
+                }
+            } finally {
+                writeLock.unlock();
+            }
+        }
+
         /**
          * Updates the count of stale keys in the cache.
          * This method is called when a CleanupKey is added to the keysToClean set.