From 406ae782e0652b522e1d30a7b42136ba38b14917 Mon Sep 17 00:00:00 2001
From: Sagar <99425694+sgup432@users.noreply.github.com>
Date: Thu, 6 Jun 2024 15:02:00 -0700
Subject: [PATCH] [Caching] Move cache removal notifications outside lru lock
 (#14017)

---------

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>
---
 .../org/opensearch/common/cache/Cache.java    | 53 ++++++++++++++++---
 .../common/cache/RemovalListener.java         |  5 ++
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/server/src/main/java/org/opensearch/common/cache/Cache.java b/server/src/main/java/org/opensearch/common/cache/Cache.java
index 49809a15d738d..7b5fe5cf64bce 100644
--- a/server/src/main/java/org/opensearch/common/cache/Cache.java
+++ b/server/src/main/java/org/opensearch/common/cache/Cache.java
@@ -35,9 +35,11 @@
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.util.concurrent.ReleasableLock;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
@@ -394,7 +396,12 @@ private V get(K key, long now, Consumer<Entry<K, V>> onExpiration) {
         if (entry == null) {
             return null;
         } else {
-            promote(entry, now);
+            List<RemovalNotification<K, V>> removalNotifications = promote(entry, now).v2();
+            if (!removalNotifications.isEmpty()) {
+                for (RemovalNotification<K, V> removalNotification : removalNotifications) {
+                    removalListener.onRemoval(removalNotification);
+                }
+            }
             return entry.value;
         }
     }
@@ -444,8 +451,14 @@ private V compute(K key, CacheLoader<K, V> loader) throws ExecutionException {
 
         BiFunction<? super Entry<K, V>, Throwable, ? extends V> handler = (ok, ex) -> {
             if (ok != null) {
+                List<RemovalNotification<K, V>> removalNotifications = new ArrayList<>();
                 try (ReleasableLock ignored = lruLock.acquire()) {
-                    promote(ok, now);
+                    removalNotifications = promote(ok, now).v2();
+                }
+                if (!removalNotifications.isEmpty()) {
+                    for (RemovalNotification<K, V> removalNotification : removalNotifications) {
+                        removalListener.onRemoval(removalNotification);
+                    }
                 }
                 return ok.value;
             } else {
@@ -510,16 +523,22 @@ private void put(K key, V value, long now) {
         CacheSegment<K, V> segment = getCacheSegment(key);
         Tuple<Entry<K, V>, Entry<K, V>> tuple = segment.put(key, value, now);
         boolean replaced = false;
+        List<RemovalNotification<K, V>> removalNotifications = new ArrayList<>();
         try (ReleasableLock ignored = lruLock.acquire()) {
             if (tuple.v2() != null && tuple.v2().state == State.EXISTING) {
                 if (unlink(tuple.v2())) {
                     replaced = true;
                 }
             }
-            promote(tuple.v1(), now);
+            removalNotifications = promote(tuple.v1(), now).v2();
         }
         if (replaced) {
-            removalListener.onRemoval(new RemovalNotification<>(tuple.v2().key, tuple.v2().value, RemovalReason.REPLACED));
+            removalNotifications.add(new RemovalNotification<>(tuple.v2().key, tuple.v2().value, RemovalReason.REPLACED));
+        }
+        if (!removalNotifications.isEmpty()) {
+            for (RemovalNotification<K, V> removalNotification : removalNotifications) {
+                removalListener.onRemoval(removalNotification);
+            }
         }
     }
 
@@ -764,8 +783,17 @@ public long getEvictions() {
         }
     }
 
-    private boolean promote(Entry<K, V> entry, long now) {
+    /**
+     * Promotes the desired entry to the head of the lru list and tries to see if it needs to evict any entries in
+     * case the cache size is exceeding or the entry got expired.
+     * @param entry Entry to be promoted
+     * @param now the current time
+     * @return Returns a tuple. v1 signifies whether an entry got promoted, v2 signifies the list of removal
+     * notifications that the callers needs to handle.
+     */
+    private Tuple<Boolean, List<RemovalNotification<K, V>>> promote(Entry<K, V> entry, long now) {
         boolean promoted = true;
+        List<RemovalNotification<K, V>> removalNotifications = new ArrayList<>();
         try (ReleasableLock ignored = lruLock.acquire()) {
             switch (entry.state) {
                 case DELETED:
@@ -779,10 +807,21 @@ private boolean promote(Entry<K, V> entry, long now) {
                     break;
             }
             if (promoted) {
-                evict(now);
+                while (tail != null && shouldPrune(tail, now)) {
+                    Entry<K, V> entryToBeRemoved = tail;
+                    CacheSegment<K, V> segment = getCacheSegment(entryToBeRemoved.key);
+                    if (segment != null) {
+                        segment.remove(entryToBeRemoved.key, entryToBeRemoved.value, f -> {});
+                    }
+                    if (unlink(entryToBeRemoved)) {
+                        removalNotifications.add(
+                            new RemovalNotification<>(entryToBeRemoved.key, entryToBeRemoved.value, RemovalReason.EVICTED)
+                        );
+                    }
+                }
             }
         }
-        return promoted;
+        return new Tuple<>(promoted, removalNotifications);
     }
 
     private void evict(long now) {
diff --git a/server/src/main/java/org/opensearch/common/cache/RemovalListener.java b/server/src/main/java/org/opensearch/common/cache/RemovalListener.java
index 369313f9f93f4..3c09d3de49fc4 100644
--- a/server/src/main/java/org/opensearch/common/cache/RemovalListener.java
+++ b/server/src/main/java/org/opensearch/common/cache/RemovalListener.java
@@ -39,5 +39,10 @@
  */
 @FunctionalInterface
 public interface RemovalListener<K, V> {
+
+    /**
+     * This may be called from multiple threads at once. So implementation needs to be thread safe.
+     * @param notification removal notification for desired entry.
+     */
     void onRemoval(RemovalNotification<K, V> notification);
 }