From a92956d50abbe16e96f3110e403f1db785738722 Mon Sep 17 00:00:00 2001
From: Atri Sharma <atris@amazon.com>
Date: Wed, 2 Oct 2019 01:39:52 +0530
Subject: [PATCH] Reproduce across segment caching of same query

---
 .../apache/lucene/search/LRUQueryCache.java   | 10 ++++
 .../lucene/search/TestLRUQueryCache.java      | 57 +++++++++++++++++++
 2 files changed, 67 insertions(+)

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 24c5e98e723..23455129f34 100644
--- a/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
+++ b/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
@@ -22,6 +22,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.ConcurrentModificationException;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.Iterator;
@@ -104,6 +105,8 @@ public class LRUQueryCache implements QueryCache, Accountable {
   // of the same query in the same cache. We use a set because it is an invariant that
   // the entries of this data structure be unique.
   private final Set<Query> inFlightAsyncLoadQueries = new HashSet<>();
+  // For testing
+  private final Map<Query, Object> loadInitiatingSegmentMap = new HashMap<>();
   // 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
@@ -918,15 +921,22 @@ public class LRUQueryCache implements QueryCache, Accountable {
        * do not trigger another cache operation
        */
       if (inFlightAsyncLoadQueries.add(in.getQuery()) == false) {
+        if (loadInitiatingSegmentMap.get(in.getQuery()) != context.id()) {
+          throw new IllegalStateException("Loading segments are different, yet returning false");
+        }
+
         return false;
       }
 
+      loadInitiatingSegmentMap.put(in.getQuery(), context.id());
+
       FutureTask<Void> task = new FutureTask<>(() -> {
         DocIdSet localDocIdSet = cache(context);
         putIfAbsent(in.getQuery(), localDocIdSet, cacheHelper);
 
         // Remove the key from inflight -- the key is loaded now
         Object retValue = inFlightAsyncLoadQueries.remove(in.getQuery());
+        loadInitiatingSegmentMap.remove(in.getQuery());
 
         // The query should have been present in the inflight queries set before
         // we actually loaded it -- hence the removal of the key should be successful
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 fddc27089c8..48d545110bb 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestLRUQueryCache.java
@@ -299,6 +299,63 @@ public class TestLRUQueryCache extends LuceneTestCase {
     dir.close();
   }
 
+  public void testLRUConcurrentCachingAcrossSegments() throws Exception {
+    Directory dir = newDirectory();
+    final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+
+    int numDocs = atLeast(150);
+    int numIterations = atLeast(3);
+
+    for (int i = 0; i < numIterations; i++) {
+      for (int j = 0; j < numDocs; j++) {
+        Document doc = new Document();
+        StringField f = new StringField("color", "blue", Store.NO);
+        doc.add(f);
+        w.addDocument(doc);
+        w.addDocument(doc);
+        w.addDocument(doc);
+      }
+      w.commit();
+    }
+
+    final DirectoryReader reader = w.getReader();
+
+    ExecutorService service = new ThreadPoolExecutor(4, 4, 0L, TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue<Runnable>(),
+        new NamedThreadFactory("TestLRUQueryCache"));
+
+    IndexSearcher searcher = new IndexSearcher(reader, service) {
+      @Override
+      protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+        ArrayList<LeafSlice> slices = new ArrayList<>();
+        for (LeafReaderContext ctx : leaves) {
+          slices.add(new LeafSlice(Arrays.asList(ctx)));
+        }
+        return slices.toArray(new LeafSlice[0]);
+      }
+    };
+
+    final LRUQueryCache queryCache = new LRUQueryCache(2, 100000, context -> true);
+
+    final Query blue = new TermQuery(new Term("color", "blue"));
+
+    assertEquals(Collections.emptyList(), queryCache.cachedQueries());
+
+    searcher.setQueryCache(queryCache);
+    searcher.setQueryCachingPolicy(ALWAYS_CACHE);
+    assert searcher.getSlices().length > 1;
+
+    IllegalStateException exc = expectThrows(IllegalStateException.class,
+        () -> searcher.search(new ConstantScoreQuery(blue), 1));
+
+    assertTrue(exc.getMessage().contains("Loading segments are different, yet returning false"));
+
+    reader.close();
+    w.close();
+    dir.close();
+    service.shutdown();
+  }
+
   public void testLRUConcurrentLoadAndEviction() throws Exception {
     Directory dir = newDirectory();
     final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
-- 
2.17.2 (Apple Git-113)

