From 1baae3629a585f2b53a2ff54a4ecd765f9e1db0c Mon Sep 17 00:00:00 2001
From: Grigoriy Troitskiy <gtroitskiy@gmail.com>
Date: Tue, 3 Oct 2023 16:52:12 +0300
Subject: [PATCH] Make LRUQueryCache respect Accountable queries on eviction
 and consistency check (#12614)

Given a query that implements Accountable, the LRUQueryCache would increment
its internal accounting by the amount reported by Accountable.ramBytesUsed(), but
only decrement on eviction by the default used for all other queries.  This meant that
the cache could eventually think it had run out of space, even if there were no queries
in it at all.  This commit ensures that queries that implement Accountable are always
accounted for correctly.
---
 lucene/CHANGES.txt                            |  3 ++
 .../apache/lucene/search/LRUQueryCache.java   | 25 +++++++------
 .../lucene/search/TestLRUQueryCache.java      | 37 +++++++++++++++++++
 3 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index caff0fdb9d4..93d2563b39d 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -176,6 +176,9 @@ Changes in runtime behavior
 Bug Fixes
 ---------------------
 
+* GITHUB#12614: Make LRUQueryCache respect Accountable queries on eviction and consistency check
+  (Grigoriy Troitskiy)
+
 * GITHUB#11556: HTMLStripCharFilter fails on '>' or '<' characters in attribute values. (Elliot Lin)
 
 Build
diff --git a/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java b/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
index 3479a55e6a9..763ebc12e8a 100644
--- a/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
+++ b/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
@@ -297,12 +297,7 @@ public class LRUQueryCache implements QueryCache, Accountable {
     try {
       Query singleton = uniqueQueries.putIfAbsent(query, query);
       if (singleton == null) {
-        if (query instanceof Accountable) {
-          onQueryCache(
-              query, LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY + ((Accountable) query).ramBytesUsed());
-        } else {
-          onQueryCache(query, LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY + QUERY_DEFAULT_RAM_BYTES_USED);
-        }
+        onQueryCache(query, getRamBytesUsed(query));
       } else {
         query = singleton;
       }
@@ -385,7 +380,7 @@ public class LRUQueryCache implements QueryCache, Accountable {
 
   private void onEviction(Query singleton) {
     assert lock.isHeldByCurrentThread();
-    onQueryEviction(singleton, LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY + QUERY_DEFAULT_RAM_BYTES_USED);
+    onQueryEviction(singleton, getRamBytesUsed(singleton));
     for (LeafCache leafCache : cache.values()) {
       leafCache.remove(singleton);
     }
@@ -405,6 +400,13 @@ public class LRUQueryCache implements QueryCache, Accountable {
     }
   }
 
+  private static long getRamBytesUsed(Query query) {
+    return LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY
+        + (query instanceof Accountable accountableQuery
+            ? accountableQuery.ramBytesUsed()
+            : QUERY_DEFAULT_RAM_BYTES_USED);
+  }
+
   // pkg-private for testing
   void assertConsistent() {
     lock.lock();
@@ -429,11 +431,10 @@ public class LRUQueryCache implements QueryCache, Accountable {
               "One leaf cache contains more keys than the top-level cache: " + keys);
         }
       }
-      long recomputedRamBytesUsed =
-          HASHTABLE_RAM_BYTES_PER_ENTRY * cache.size()
-              + LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY * uniqueQueries.size();
-      recomputedRamBytesUsed +=
-          mostRecentlyUsedQueries.size() * (long) QUERY_DEFAULT_RAM_BYTES_USED;
+      long recomputedRamBytesUsed = HASHTABLE_RAM_BYTES_PER_ENTRY * cache.size();
+      for (Query query : mostRecentlyUsedQueries) {
+        recomputedRamBytesUsed += getRamBytesUsed(query);
+      }
       for (LeafCache leafCache : cache.values()) {
         recomputedRamBytesUsed += HASHTABLE_RAM_BYTES_PER_ENTRY * leafCache.cache.size();
         for (CacheAndCount cached : leafCache.cache.values()) {
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java b/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java
index da68d13c669..fddeb96f10f 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java
@@ -593,6 +593,43 @@ public class TestLRUQueryCache extends LuceneTestCase {
     dir.close();
   }
 
+  public void testConsistencyWithAccountableQueries() throws IOException {
+    var queryCache = new LRUQueryCache(1, 10000000, context -> true, Float.POSITIVE_INFINITY);
+
+    var dir = newDirectory();
+    var writer = new RandomIndexWriter(random(), dir);
+    writer.addDocument(new Document());
+    var reader = writer.getReader();
+    var searcher = new IndexSearcher(reader);
+    searcher.setQueryCache(queryCache);
+    searcher.setQueryCachingPolicy(ALWAYS_CACHE);
+
+    queryCache.assertConsistent();
+
+    var accountableQuery = new AccountableDummyQuery();
+    searcher.count(accountableQuery);
+    var expectedRamBytesUsed =
+        HASHTABLE_RAM_BYTES_PER_ENTRY
+            + LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY
+            + accountableQuery.ramBytesUsed()
+            + queryCache.getChildResources().iterator().next().ramBytesUsed();
+    assertEquals(expectedRamBytesUsed, queryCache.ramBytesUsed());
+    queryCache.assertConsistent();
+
+    queryCache.clearQuery(accountableQuery);
+    assertEquals(HASHTABLE_RAM_BYTES_PER_ENTRY, queryCache.ramBytesUsed());
+    queryCache.assertConsistent();
+
+    queryCache.clearCoreCacheKey(
+        reader.getContext().leaves().get(0).reader().getCoreCacheHelper().getKey());
+    assertEquals(0, queryCache.ramBytesUsed());
+    queryCache.assertConsistent();
+
+    reader.close();
+    writer.close();
+    dir.close();
+  }
+
   public void testOnUse() throws IOException {
     final LRUQueryCache queryCache =
         new LRUQueryCache(