[QueryCache] Use RamUsageEstimator instead of defaulting to 1024 bytes for non-accountable query size (#15124)

* Use RamUsageEstimator to calculate query size instead of the default 1024 bytes for non accountable queries

* Use try-with-resources for directory and indexWriter

* Adding changes to cache query size and queries per clause to reduce impact of repeated visit() calls during RamUsageEstimator.sizeOf()

* Adding changelog entry

* Making queries per clause list immutable

* Adding unit test to verify query size is cached correctly

* Renmaing QueryMetadata to Record

* Changing query metadata to record type and removing boolean query changes

Author: sgup432

lucene/CHANGES.txt

@@ -64,6 +64,8 @@ Improvements
   SIGSEGV with the old unsafe "mmap hack", which has been replaced by MemorySegments.
   (Uwe Schindler, Robert Muir)
 
+* GITHUB#15124: Use RamUsageEstimator to calculate size for non-accountable queries. (Sagar Upadhyaya)
+
 Optimizations
 ---------------------
 * GITHUB#14011: Reduce allocation rate in HNSW concurrent merge. (Viliam Durina)

lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java

@@ -18,7 +18,6 @@
 
 import static org.apache.lucene.util.RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY;
 import static org.apache.lucene.util.RamUsageEstimator.LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY;
-import static org.apache.lucene.util.RamUsageEstimator.QUERY_DEFAULT_RAM_BYTES_USED;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -91,7 +90,7 @@ public class LRUQueryCache implements QueryCache, Accountable {
   private final Predicate<LeafReaderContext> leavesToCache;
   // maps queries that are contained in the cache to a singleton so that this
   // cache does not store several copies of the same query
-  private final Map<Query, Query> uniqueQueries;
+  private final Map<Query, QueryMetadata> uniqueQueries;
   // The contract between this set and the per-leaf caches is that per-leaf caches
   // are only allowed to store sub-sets of the queries that are contained in
   // mostRecentlyUsedQueries. This is why write operations are performed under a lock
@@ -176,6 +175,11 @@ public LRUQueryCache(int maxSize, long maxRamBytesUsed) {
     this(maxSize, maxRamBytesUsed, new MinSegmentSizePredicate(10000), 10);
   }
 
+  // pkg-private for testing
+  Map<Query, QueryMetadata> getUniqueQueries() {
+    return uniqueQueries;
+  }
+
   // pkg-private for testing
   static class MinSegmentSizePredicate implements Predicate<LeafReaderContext> {
     private final int minSize;
@@ -300,32 +304,35 @@ CacheAndCount get(Query key, IndexReader.CacheHelper cacheHelper) {
       return null;
     }
     // this get call moves the query to the most-recently-used position
-    final Query singleton = uniqueQueries.get(key);
-    if (singleton == null) {
+    final QueryMetadata record = uniqueQueries.get(key);
+    if (record == null || record.query == null) {
       onMiss(readerKey, key);
       return null;
     }
-    final CacheAndCount cached = leafCache.get(singleton);
+    final CacheAndCount cached = leafCache.get(record.query);
     if (cached == null) {
-      onMiss(readerKey, singleton);
+      onMiss(readerKey, record.query);
     } else {
-      onHit(readerKey, singleton);
+      onHit(readerKey, record.query);
     }
     return cached;
   }
 
-  private void putIfAbsent(Query query, CacheAndCount cached, IndexReader.CacheHelper cacheHelper) {
+  public void putIfAbsent(Query query, CacheAndCount cached, IndexReader.CacheHelper cacheHelper) {
     assert query instanceof BoostQuery == false;
     assert query instanceof ConstantScoreQuery == false;
     // under a lock to make sure that mostRecentlyUsedQueries and cache remain sync'ed
     writeLock.lock();
     try {
-      Query singleton = uniqueQueries.putIfAbsent(query, query);
-      if (singleton == null) {
-        onQueryCache(query, getRamBytesUsed(query));
-      } else {
-        query = singleton;
-      }
+      QueryMetadata record =
+          uniqueQueries.computeIfAbsent(
+              query,
+              q -> {
+                long queryRamBytesUsed = getRamBytesUsed(q);
+                onQueryCache(q, queryRamBytesUsed);
+                return new QueryMetadata(q, queryRamBytesUsed);
+              });
+      query = record.query;
       final IndexReader.CacheKey key = cacheHelper.getKey();
       LeafCache leafCache = cache.get(key);
       if (leafCache == null) {
@@ -347,25 +354,24 @@ private void evictIfNecessary() {
     assert writeLock.isHeldByCurrentThread();
     // under a lock to make sure that mostRecentlyUsedQueries and cache keep sync'ed
     if (requiresEviction()) {
-
-      Iterator<Query> iterator = mostRecentlyUsedQueries.iterator();
+      Iterator<Map.Entry<Query, QueryMetadata>> iterator = uniqueQueries.entrySet().iterator();
       do {
-        final Query query = iterator.next();
-        final int size = mostRecentlyUsedQueries.size();
+        final Map.Entry<Query, QueryMetadata> entry = iterator.next();
+        final int size = uniqueQueries.size();
         iterator.remove();
-        if (size == mostRecentlyUsedQueries.size()) {
+        if (size == uniqueQueries.size()) {
           // size did not decrease, because the hash of the query changed since it has been
           // put into the cache
           throw new ConcurrentModificationException(
               "Removal from the cache failed! This "
                   + "is probably due to a query which has been modified after having been put into "
                   + " the cache or a badly implemented clone(). Query class: ["
-                  + query.getClass()
+                  + entry.getKey().getClass()
                   + "], query: ["
-                  + query
+                  + entry.getKey()
                   + "]");
         }
-        onEviction(query);
+        onEviction(entry.getKey(), entry.getValue().queryRamBytesUsed);
       } while (iterator.hasNext() && requiresEviction());
     }
   }
@@ -394,18 +400,18 @@ public void clearCoreCacheKey(Object coreKey) {
   public void clearQuery(Query query) {
     writeLock.lock();
     try {
-      final Query singleton = uniqueQueries.remove(query);
-      if (singleton != null) {
-        onEviction(singleton);
+      final QueryMetadata record = uniqueQueries.remove(query);
+      if (record != null && record.query != null) {
+        onEviction(record.query, record.queryRamBytesUsed);
       }
     } finally {
       writeLock.unlock();
     }
   }
 
-  private void onEviction(Query singleton) {
+  private void onEviction(Query singleton, long querySizeInBytes) {
     assert writeLock.isHeldByCurrentThread();
-    onQueryEviction(singleton, getRamBytesUsed(singleton));
+    onQueryEviction(singleton, querySizeInBytes);
     for (LeafCache leafCache : cache.values()) {
       leafCache.remove(singleton);
     }
@@ -426,10 +432,9 @@ public void clear() {
   }
 
   private static long getRamBytesUsed(Query query) {
-    return LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY
-        + (query instanceof Accountable accountableQuery
-            ? accountableQuery.ramBytesUsed()
-            : QUERY_DEFAULT_RAM_BYTES_USED);
+    // Here 32 represents a rough shallow size for a query object
+    long queryRamBytesUsed = RamUsageEstimator.sizeOf(query, 32);
+    return LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY + queryRamBytesUsed;
   }
 
   // pkg-private for testing
@@ -734,6 +739,9 @@ public long ramBytesUsed() {
     }
   }
 
+  // pkg-private for testing
+  record QueryMetadata(Query query, long queryRamBytesUsed) {}
+
   private class CachingWrapperWeight extends ConstantScoreWeight {
 
     private final Weight in;

lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java

@@ -74,6 +74,7 @@
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.SuppressForbidden;
 
 public class TestLRUQueryCache extends LuceneTestCase {
@@ -1062,6 +1063,65 @@ public void testBooleanQueryCachesSubClauses() throws IOException {
     dir.close();
   }
 
+  public void testQuerySizeBytesAreCached() throws IOException {
+    try (Directory dir = newDirectory();
+        RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
+      Document doc = new Document();
+      doc.add(new StringField("foo", "bar", Store.YES));
+      doc.add(new StringField("foo", "quux", Store.YES));
+      w.addDocument(doc);
+      w.commit();
+      final IndexReader reader = w.getReader();
+      final IndexSearcher searcher = newSearcher(reader);
+      final LRUQueryCache queryCache =
+          new LRUQueryCache(1000000, 10000000, _ -> true, Float.POSITIVE_INFINITY);
+      searcher.setQueryCache(queryCache);
+      searcher.setQueryCachingPolicy(ALWAYS_CACHE);
+      StringBuilder sb = new StringBuilder();
+      sb.append("a".repeat(100));
+      String term = sb.toString();
+      TermQuery termQuery = new TermQuery(new Term("foo", term));
+      long expectedQueryInBytes =
+          LINKED_HASHTABLE_RAM_BYTES_PER_ENTRY + RamUsageEstimator.sizeOf(termQuery, 32);
+      searcher.search(new ConstantScoreQuery(termQuery), 1);
+
+      assertEquals(1, queryCache.getUniqueQueries().size());
+      assertEquals(
+          expectedQueryInBytes, queryCache.getUniqueQueries().get(termQuery).queryRamBytesUsed());
+      reader.close();
+    }
+  }
+
+  public void testCacheRamBytesWithALargeTermQuery() throws IOException {
+    try (Directory dir = newDirectory();
+        RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
+      Document doc = new Document();
+      doc.add(new StringField("foo", "bar", Store.YES));
+      doc.add(new StringField("foo", "quux", Store.YES));
+      w.addDocument(doc);
+      w.commit();
+      final IndexReader reader = w.getReader();
+      final IndexSearcher searcher = newSearcher(reader);
+      final LRUQueryCache queryCache =
+          new LRUQueryCache(1000000, 10000000, _ -> true, Float.POSITIVE_INFINITY);
+      searcher.setQueryCache(queryCache);
+      searcher.setQueryCachingPolicy(ALWAYS_CACHE);
+      StringBuilder sb = new StringBuilder();
+      // Create a large string for the field value so it certainly exceeds the default query size we
+      // use ie 1024 bytes.
+      sb.append("a".repeat(1200));
+      String longTerm = sb.toString();
+      TermQuery must = new TermQuery(new Term("foo", longTerm));
+      long queryInBytes = RamUsageEstimator.sizeOf(must, 32);
+      assertTrue(queryInBytes > QUERY_DEFAULT_RAM_BYTES_USED);
+      searcher.search(new ConstantScoreQuery(must), 1);
+
+      assertEquals(1, queryCache.cachedQueries().size());
+      assertTrue(queryCache.ramBytesUsed() >= queryInBytes);
+      reader.close();
+    }
+  }
+
   private static Term randomTerm() {
     final String term = RandomPicks.randomFrom(random(), Arrays.asList("foo", "bar", "baz"));
     return new Term("foo", term);