LUCENE-9662: CheckIndex should be concurrent - parallelizing index check across segments

lucene/core/src/java/org/apache/lucene/index/CheckIndex.java

@@ -31,6 +31,13 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.codecs.DocValuesProducer;
 import org.apache.lucene.codecs.NormsProducer;
@@ -60,6 +67,7 @@
 import org.apache.lucene.util.FixedBitSet;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.LongBitSet;
+import org.apache.lucene.util.NamedThreadFactory;
 import org.apache.lucene.util.StringHelper;
 import org.apache.lucene.util.SuppressForbidden;
 import org.apache.lucene.util.Version;
@@ -450,6 +458,13 @@ public void setChecksumsOnly(boolean v) {
 
   private boolean checksumsOnly;
 
+  /** Set threadCount used for parallelizing index integrity checking. */
+  public void setThreadCount(int tc) {
+    threadCount = tc;
+  }
+
+  private int threadCount = Runtime.getRuntime().availableProcessors();
+
   /**
    * Set infoStream where messages should go. If null, no messages are printed. If verbose is true
    * then more details are printed.
@@ -488,8 +503,35 @@ public Status checkIndex() throws IOException {
    *     quite a long time to run.
    */
   public Status checkIndex(List<String> onlySegments) throws IOException {
+    ExecutorService executorService =
+        Executors.newFixedThreadPool(threadCount, new NamedThreadFactory("async-check-index"));
+    try {
+      return checkIndex(onlySegments, executorService);
+    } finally {
+      executorService.shutdown();
+      try {
+        executorService.awaitTermination(5, TimeUnit.SECONDS);
+      } catch (
+          @SuppressWarnings("unused")
+          InterruptedException e) {
+      } finally {
+        executorService.shutdownNow();
+      }
+    }
+  }
+
+  /**
+   * Returns a {@link Status} instance detailing the state of the index.
+   *
+   * <p>This method allows caller to pass in customized ExecutorService to speed up the check.
+   *
+   * <p><b>WARNING</b>: make sure you only call this when the index is not opened by any writer.
+   */
+  public Status checkIndex(List<String> onlySegments, ExecutorService executorService)
+      throws IOException {
     ensureOpen();
     long startNS = System.nanoTime();
+
     NumberFormat nf = NumberFormat.getInstance(Locale.ROOT);
     SegmentInfos sis = null;
     Status result = new Status();
@@ -605,6 +647,11 @@ public Status checkIndex(List<String> onlySegments) throws IOException {
     result.newSegments.clear();
     result.maxSegmentName = -1;
 
+    // nocommit the msg statements with infoStream inside this loop (as well as inside each
+    // index part checking methods such as testLiveDocs) may be
+    // interleaved when this section of code run concurrently. Is it ok to not synchronize for
+    // these msg statements and instead have each msg content to include segment id information for
+    // identification?
     for (int i = 0; i < numSegments; i++) {
       final SegmentCommitInfo info = sis.info(i);
       long segmentName = Long.parseLong(info.info.name.substring(1), Character.MAX_RADIX);
@@ -639,6 +686,8 @@ public Status checkIndex(List<String> onlySegments) throws IOException {
       SegmentReader reader = null;
 
       try {
+        // nocommit these msg statements may require synchronization if printed without segment
+        // identifier
         msg(infoStream, "    version=" + (version == null ? "3.0" : version));
         msg(infoStream, "    id=" + StringHelper.idToString(info.info.getId()));
         final Codec codec = info.info.getCodec();
@@ -731,40 +780,170 @@ public Status checkIndex(List<String> onlySegments) throws IOException {
         }
 
         if (checksumsOnly == false) {
+          // This redundant assignment is done to make compiler happy
+          SegmentReader finalReader = reader;
+
           // Test Livedocs
-          segInfoStat.liveDocStatus = testLiveDocs(reader, infoStream, failFast);
+          CompletableFuture<Void> testliveDocs =
+              CompletableFuture.supplyAsync(
+                      callableToSupplier(() -> testLiveDocs(finalReader, infoStream, failFast)),
+                      executorService)
+                  .thenAccept(
+                      liveDocStatus -> {
+                        segInfoStat.liveDocStatus = liveDocStatus;
+                      });
 
           // Test Fieldinfos
-          segInfoStat.fieldInfoStatus = testFieldInfos(reader, infoStream, failFast);
+          CompletableFuture<Void> testFieldInfos =
+              CompletableFuture.supplyAsync(
+                      callableToSupplier(() -> testFieldInfos(finalReader, infoStream, failFast)),
+                      executorService)
+                  .thenAccept(
+                      fieldInfoStatus -> {
+                        segInfoStat.fieldInfoStatus = fieldInfoStatus;
+                      });
 
           // Test Field Norms
-          segInfoStat.fieldNormStatus = testFieldNorms(reader, infoStream, failFast);
+          CompletableFuture<Void> testFieldNorms =
+              CompletableFuture.supplyAsync(
+                      callableToSupplier(() -> testFieldNorms(finalReader, infoStream, failFast)),
+                      executorService)
+                  .thenAccept(
+                      fieldNormStatus -> {
+                        segInfoStat.fieldNormStatus = fieldNormStatus;
+                      });
 
           // Test the Term Index
-          segInfoStat.termIndexStatus =
-              testPostings(reader, infoStream, verbose, doSlowChecks, failFast);
+          CompletableFuture<Void> testTermIndex =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testPostings(
+                              finalReader, infoStream, verbose, doSlowChecks, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      termIndexStatus -> {
+                        segInfoStat.termIndexStatus = termIndexStatus;
+                      });
 
           // Test Stored Fields
-          segInfoStat.storedFieldStatus = testStoredFields(reader, infoStream, failFast);
+          CompletableFuture<Void> testStoredFields =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testStoredFields(finalReader, infoStream, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      storedFieldStatus -> {
+                        segInfoStat.storedFieldStatus = storedFieldStatus;
+                      });
 
           // Test Term Vectors
-          segInfoStat.termVectorStatus =
-              testTermVectors(reader, infoStream, verbose, doSlowChecks, failFast);
+          CompletableFuture<Void> testTermVectors =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testTermVectors(
+                              finalReader, infoStream, verbose, doSlowChecks, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      termVectorStatus -> {
+                        segInfoStat.termVectorStatus = termVectorStatus;
+                      });
 
           // Test Docvalues
-          segInfoStat.docValuesStatus = testDocValues(reader, infoStream, failFast);
+          CompletableFuture<Void> testDocValues =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testDocValues(finalReader, infoStream, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      docValuesStatus -> {
+                        segInfoStat.docValuesStatus = docValuesStatus;
+                      });
 
           // Test PointValues
-          segInfoStat.pointsStatus = testPoints(reader, infoStream, failFast);
+          CompletableFuture<Void> testPointvalues =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testPoints(finalReader, infoStream, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      pointsStatus -> {
+                        segInfoStat.pointsStatus = pointsStatus;
+                      });
 
           // Test VectorValues
-          segInfoStat.vectorValuesStatus = testVectors(reader, infoStream, failFast);
+          CompletableFuture<Void> testVectors =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testVectors(finalReader, infoStream, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      vectorValuesStatus -> {
+                        segInfoStat.vectorValuesStatus = vectorValuesStatus;
+                      });
 
           // Test index sort
-          segInfoStat.indexSortStatus = testSort(reader, indexSort, infoStream, failFast);
+          CompletableFuture<Void> testSort =
+              CompletableFuture.supplyAsync(
+                      () -> {
+                        try {
+                          return testSort(finalReader, indexSort, infoStream, failFast);
+                        } catch (IOException e) {
+                          throw new CompletionException(e);
+                        }
+                      },
+                      executorService)
+                  .thenAccept(
+                      indexSortStatus -> {
+                        segInfoStat.indexSortStatus = indexSortStatus;
+                      });
 
           // Rethrow the first exception we encountered
           //  This will cause stats for failed segments to be incremented properly
+          // nocommit The error != null check ordering below requires sequencing the above async
+          // calls.
+          // Does the order really matter here or can be done differently?
+          CompletableFuture.allOf(
+                  testliveDocs,
+                  testFieldInfos,
+                  testFieldNorms,
+                  testTermIndex,
+                  testStoredFields,
+                  testTermVectors,
+                  testDocValues,
+                  testPointvalues,
+                  testVectors,
+                  testSort)
+              .join();
           if (segInfoStat.liveDocStatus.error != null) {
             throw new RuntimeException("Live docs test failed");
           } else if (segInfoStat.fieldInfoStatus.error != null) {
@@ -783,6 +962,8 @@ public Status checkIndex(List<String> onlySegments) throws IOException {
             throw new RuntimeException("Points test failed");
           }
         }
+
+        // nocommit parallelize this as well?
         final String softDeletesField = reader.getFieldInfos().getSoftDeletesField();
         if (softDeletesField != null) {
           checkSoftDeletes(softDeletesField, info, reader, infoStream, failFast);
@@ -843,6 +1024,18 @@ public Status checkIndex(List<String> onlySegments) throws IOException {
     return result;
   }
 
+  private <T> Supplier<T> callableToSupplier(Callable<T> callable) {
+    return () -> {
+      try {
+        return callable.call();
+      } catch (RuntimeException | Error e) {
+        throw e;
+      } catch (Throwable e) {
+        throw new CompletionException(e);
+      }
+    };
+  }
+
   /**
    * Tests index sort order.
    *
@@ -931,6 +1124,8 @@ public static Status.LiveDocStatus testLiveDocs(
     final Status.LiveDocStatus status = new Status.LiveDocStatus();
 
     try {
+      // nocommit these msg statements may require synchronization if printed without segment
+      // identifier
       if (infoStream != null) infoStream.print("    test: check live docs.....");
       final int numDocs = reader.numDocs();
       if (reader.hasDeletions()) {
@@ -969,6 +1164,8 @@ public static Status.LiveDocStatus testLiveDocs(
             }
           }
         }
+        // nocommit these msg statements may require synchronization if printed without segment
+        // identifier
         msg(
             infoStream,
             String.format(
@@ -3622,6 +3819,7 @@ public static class Options {
     boolean doSlowChecks = false;
     boolean verbose = false;
     boolean doChecksumsOnly = false;
+    int threadCount;
     List<String> onlySegments = new ArrayList<>();
     String indexPath = null;
     String dirImpl = null;
@@ -3720,6 +3918,12 @@ public static Options parseOptions(String[] args) {
         }
         i++;
         opts.dirImpl = args[i];
+      } else if ("-threadCount".equals(arg)) {
+        if (i == args.length - 1) {
+          throw new IllegalArgumentException("-threadCount requires a following number");
+        }
+        i++;
+        opts.threadCount = Integer.parseInt(args[i]);
       } else {
         if (opts.indexPath != null) {
           throw new IllegalArgumentException("ERROR: unexpected extra argument '" + args[i] + "'");
@@ -3787,8 +3991,10 @@ public int doCheck(Options opts) throws IOException, InterruptedException {
     setDoSlowChecks(opts.doSlowChecks);
     setChecksumsOnly(opts.doChecksumsOnly);
     setInfoStream(opts.out, opts.verbose);
+    setThreadCount(opts.threadCount);
 
     Status result = checkIndex(opts.onlySegments);
+
     if (result.missingSegments) {
       return 1;
     }