Add a thread safe CachingLeafSlicesSupplier to compute and cache the …

…LeafSlices used with concurrent segment

search. It uses the protected method `slices` by default to compute the slices which can be
overriden by the sub classes of IndexSearcher

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

@@ -29,6 +29,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.ThreadPoolExecutor;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
@@ -116,8 +117,11 @@ public class IndexSearcher {
   protected final IndexReaderContext readerContext;
   protected final List<LeafReaderContext> leafContexts;
 
-  /** used with executor - each slice holds a set of leafs executed within one thread */
-  private final LeafSlice[] leafSlices;
+  /**
+   * used with executor - LeafSlice supplier where each slice holds a set of leafs executed within
+   * one thread
+   */
+  private final CachingLeafSlicesSupplier leafSlicesSupplier;
 
   // These are only used for multi-threaded search
   private final Executor executor;
@@ -236,7 +240,7 @@ public IndexSearcher(IndexReaderContext context, Executor executor) {
     this.sliceExecutor = sliceExecutor;
     this.readerContext = context;
     leafContexts = context.leaves();
-    this.leafSlices = executor == null ? null : slices(leafContexts);
+    leafSlicesSupplier = new CachingLeafSlicesSupplier(executor, this::slices, leafContexts);
   }
 
   /**
@@ -431,7 +435,7 @@ public int count(Query query) throws IOException {
    * @lucene.experimental
    */
   public LeafSlice[] getSlices() {
-    return leafSlices;
+    return leafSlicesSupplier.get();
   }
 
   /**
@@ -456,16 +460,17 @@ public TopDocs searchAfter(ScoreDoc after, Query query, int numHits) throws IOEx
 
     final int cappedNumHits = Math.min(numHits, limit);
 
+    final LeafSlice[] leafSlices = leafSlicesSupplier.get();
     final CollectorManager<TopScoreDocCollector, TopDocs> manager =
         new CollectorManager<TopScoreDocCollector, TopDocs>() {
 
           private final HitsThresholdChecker hitsThresholdChecker =
-              (executor == null || leafSlices.length <= 1)
+              (leafSlices == null || leafSlices.length <= 1)
                   ? HitsThresholdChecker.create(Math.max(TOTAL_HITS_THRESHOLD, numHits))
                   : HitsThresholdChecker.createShared(Math.max(TOTAL_HITS_THRESHOLD, numHits));
 
           private final MaxScoreAccumulator minScoreAcc =
-              (executor == null || leafSlices.length <= 1) ? null : new MaxScoreAccumulator();
+              (leafSlices == null || leafSlices.length <= 1) ? null : new MaxScoreAccumulator();
 
           @Override
           public TopScoreDocCollector newCollector() throws IOException {
@@ -599,17 +604,18 @@ private TopFieldDocs searchAfter(
     }
     final int cappedNumHits = Math.min(numHits, limit);
     final Sort rewrittenSort = sort.rewrite(this);
+    final LeafSlice[] leafSlices = leafSlicesSupplier.get();
 
     final CollectorManager<TopFieldCollector, TopFieldDocs> manager =
         new CollectorManager<>() {
 
           private final HitsThresholdChecker hitsThresholdChecker =
-              (executor == null || leafSlices.length <= 1)
+              (leafSlices == null || leafSlices.length <= 1)
                   ? HitsThresholdChecker.create(Math.max(TOTAL_HITS_THRESHOLD, numHits))
                   : HitsThresholdChecker.createShared(Math.max(TOTAL_HITS_THRESHOLD, numHits));
 
           private final MaxScoreAccumulator minScoreAcc =
-              (executor == null || leafSlices.length <= 1) ? null : new MaxScoreAccumulator();
+              (leafSlices == null || leafSlices.length <= 1) ? null : new MaxScoreAccumulator();
 
           @Override
           public TopFieldCollector newCollector() throws IOException {
@@ -639,7 +645,7 @@ public TopFieldDocs reduce(Collection<TopFieldCollector> collectors) throws IOEx
   /**
    * Lower-level search API. Search all leaves using the given {@link CollectorManager}. In contrast
    * to {@link #search(Query, Collector)}, this method will use the searcher's {@link Executor} in
-   * order to parallelize execution of the collection on the configured {@link #leafSlices}.
+   * order to parallelize execution of the collection on the configured {@link #getSlices()}.
    *
    * @see CollectorManager
    * @lucene.experimental
@@ -654,7 +660,8 @@ public <C extends Collector, T> T search(Query query, CollectorManager<C, T> col
 
   private <C extends Collector, T> T search(
       Weight weight, CollectorManager<C, T> collectorManager, C firstCollector) throws IOException {
-    if (executor == null || leafSlices.length <= 1) {
+    final LeafSlice[] leafSlices = leafSlicesSupplier.get();
+    if (leafSlices == null || leafSlices.length <= 1) {
       search(leafContexts, weight, firstCollector);
       return collectorManager.reduce(Collections.singletonList(firstCollector));
     } else {
@@ -1014,4 +1021,48 @@ private static SliceExecutor getSliceExecutionControlPlane(Executor executor) {
 
     return new SliceExecutor(executor);
   }
+
+  /**
+   * Supplier for {@link LeafSlice} slices when passed in executor is non-null. This supplier
+   * computes and caches the value on first invocation and returns cached value on subsequent
+   * invocation. If the passed in provider for slice computation throws exception then same will be
+   * passed to the caller of this supplier on each invocation. If the provider returns null then
+   * {@link NullPointerException} will be thrown to the caller.
+   */
+  private static class CachingLeafSlicesSupplier implements Supplier<LeafSlice[]> {
+    private volatile LeafSlice[] leafSlices;
+
+    private final Executor executor;
+
+    private final Function<List<LeafReaderContext>, LeafSlice[]> sliceProvider;
+
+    private final List<LeafReaderContext> leaves;
+
+    private CachingLeafSlicesSupplier(
+        Executor executor,
+        Function<List<LeafReaderContext>, LeafSlice[]> provider,
+        List<LeafReaderContext> leaves) {
+      this.executor = executor;
+      this.sliceProvider = Objects.requireNonNull(provider, "leaf slice provider cannot be null");
+      this.leaves = Objects.requireNonNull(leaves, "list of LeafReaderContext cannot be null");
+    }
+
+    @Override
+    public LeafSlice[] get() {
+      if (executor == null) {
+        return null;
+      }
+
+      if (leafSlices == null) {
+        synchronized (this) {
+          if (leafSlices == null) {
+            leafSlices =
+                Objects.requireNonNull(
+                    sliceProvider.apply(leaves), "slices computed by the provider is null");
+          }
+        }
+      }
+      return leafSlices;
+    }
+  }
 }