Check if merge is aborted during file integrity checksums

lucene/codecs/src/java/org/apache/lucene/codecs/bloom/BloomFilteringPostingsFormat.java

@@ -34,6 +34,7 @@
 import org.apache.lucene.index.Fields;
 import org.apache.lucene.index.ImpactsEnum;
 import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -391,6 +392,11 @@ public void checkIntegrity() throws IOException {
       delegateFieldsProducer.checkIntegrity();
     }
 
+    @Override
+    public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+      delegateFieldsProducer.checkIntegrity(merge);
+    }
+
     @Override
     public String toString() {
       return getClass().getSimpleName()

lucene/codecs/src/java/org/apache/lucene/codecs/memory/DirectPostingsFormat.java

@@ -30,6 +30,7 @@
 import org.apache.lucene.index.Fields;
 import org.apache.lucene.index.ImpactsEnum;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.OrdTermState;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SegmentReadState;
@@ -153,6 +154,12 @@ public void checkIntegrity() throws IOException {
       // otherwise returned the raw postings reader
     }
 
+    @Override
+    public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+      // if we read entirely into ram, we already validated.
+      // otherwise returned the raw postings reader
+    }
+
     @Override
     public String toString() {
       return getClass().getSimpleName() + "(fields=" + fields.size() + ")";

lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextNormsFormat.java

@@ -22,6 +22,7 @@
 import org.apache.lucene.codecs.NormsProducer;
 import org.apache.lucene.index.EmptyDocValuesProducer;
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -77,6 +78,11 @@ public void checkIntegrity() throws IOException {
       impl.checkIntegrity();
     }
 
+    @Override
+    public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+      impl.checkIntegrity(merge);
+    }
+
     @Override
     public String toString() {
       return getClass().getSimpleName() + "(" + impl + ")";

lucene/core/src/java/org/apache/lucene/codecs/CodecUtil.java

@@ -22,6 +22,7 @@
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.index.IndexFormatTooNewException;
 import org.apache.lucene.index.IndexFormatTooOldException;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.store.BufferedChecksumIndexInput;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.DataInput;
@@ -597,26 +598,62 @@ private static void validateFooter(IndexInput in) throws IOException {
     }
   }
 
+  /**
+   * Number of bytes between consecutive merge abort checks during {@link
+   * #checksumEntireFile(IndexInput, MergePolicy.OneMerge)}. Files smaller than this are checksummed
+   * in one shot without checking for abort.
+   */
+  private static final long ABORT_CHECK_INTERVAL = 1024 * 1024;
+
   /**
    * Clones the provided input, reads all bytes from the file, and calls {@link #checkFooter}
    *
    * <p>Note that this method may be slow, as it must process the entire file. If you just need to
    * extract the checksum value, call {@link #retrieveChecksum}.
    */
   public static long checksumEntireFile(IndexInput input) throws IOException {
+    return checksumEntireFile(input, null, Long.MAX_VALUE);
+  }
+
+  /**
+   * Like {@link #checksumEntireFile(IndexInput)}, but periodically checks whether the provided
+   * merge has been aborted. This avoids spending a long time checksumming a large file when the
+   * merge has already been cancelled.
+   *
+   * @param input the index input to checksum
+   * @param merge the merge to check for abort, or {@code null} to behave like {@link
+   *     #checksumEntireFile(IndexInput)}
+   * @throws MergePolicy.MergeAbortedException if the merge is aborted during checksumming
+   */
+  public static long checksumEntireFile(IndexInput input, MergePolicy.OneMerge merge)
+      throws IOException {
+    return checksumEntireFile(input, merge, ABORT_CHECK_INTERVAL);
+  }
+
+  static long checksumEntireFile(
+      IndexInput input, MergePolicy.OneMerge merge, long abortCheckInterval) throws IOException {
     IndexInput clone = input.clone();
     clone.seek(0);
     ChecksumIndexInput in = new BufferedChecksumIndexInput(clone);
     assert in.getFilePointer() == 0;
-    if (in.length() < footerLength()) {
+    final long len = in.length();
+    if (len < footerLength()) {
       throw new CorruptIndexException(
           "misplaced codec footer (file truncated?): length="
-              + in.length()
+              + len
               + " but footerLength=="
               + footerLength(),
           input);
     }
-    in.seek(in.length() - footerLength());
+    final long target = len - footerLength();
+    if (merge == null || target <= abortCheckInterval) {
+      in.seek(target);
+    } else {
+      while (in.getFilePointer() < target) {
+        in.seek(Math.min(in.getFilePointer() + abortCheckInterval, target));
+        merge.checkAborted();
+      }
+    }
     return checkFooter(in);
   }
 

lucene/core/src/java/org/apache/lucene/codecs/DocValuesConsumer.java

@@ -133,7 +133,7 @@ public void merge(MergeState mergeState) throws IOException {
     for (DocValuesProducer docValuesProducer : mergeState.docValuesProducers) {
       if (docValuesProducer != null) {
         mergeState.checkAborted();
-        docValuesProducer.checkIntegrity();
+        docValuesProducer.checkIntegrity(mergeState.oneMerge);
       }
     }
 

lucene/core/src/java/org/apache/lucene/codecs/DocValuesProducer.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.index.DocValuesSkipper;
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
@@ -92,6 +93,20 @@ protected DocValuesProducer() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this producer, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * Returns an instance optimized for merging. This instance may only be consumed in the thread
    * that called {@link #getMergeInstance()}.

lucene/core/src/java/org/apache/lucene/codecs/FieldsConsumer.java

@@ -81,7 +81,7 @@ public void merge(MergeState mergeState, NormsProducer norms) throws IOException
       final int maxDoc = mergeState.maxDocs[readerIndex];
       if (f != null) {
         mergeState.checkAborted();
-        f.checkIntegrity();
+        f.checkIntegrity(mergeState.oneMerge);
         slices.add(new ReaderSlice(docBase, maxDoc, readerIndex));
         fields.add(f);
       }

lucene/core/src/java/org/apache/lucene/codecs/FieldsProducer.java

@@ -19,6 +19,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import org.apache.lucene.index.Fields;
+import org.apache.lucene.index.MergePolicy;
 
 /**
  * Abstract API that produces terms, doc, freq, prox, offset and payloads postings.
@@ -42,6 +43,20 @@ protected FieldsProducer() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this reader, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * Returns an instance optimized for merging. This instance may only be consumed in the thread
    * that called {@link #getMergeInstance()}.

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java

@@ -25,6 +25,7 @@
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.search.AcceptDocs;
 import org.apache.lucene.search.KnnCollector;
 import org.apache.lucene.search.ScoreDoc;
@@ -49,6 +50,20 @@ protected KnnVectorsReader() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this reader, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * If this reader wraps another for {@code field}, return the underlying reader, else return
    * {@code this}

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsWriter.java

@@ -114,7 +114,7 @@ public final void merge(MergeState mergeState) throws IOException {
       assert reader != null || mergeState.fieldInfos[i].hasVectorValues() == false;
       if (reader != null) {
         mergeState.checkAborted();
-        reader.checkIntegrity();
+        reader.checkIntegrity(mergeState.oneMerge);
       }
     }
 

lucene/core/src/java/org/apache/lucene/codecs/NormsConsumer.java

@@ -67,7 +67,7 @@ public void merge(MergeState mergeState) throws IOException {
     for (NormsProducer normsProducer : mergeState.normsProducers) {
       if (normsProducer != null) {
         mergeState.checkAborted();
-        normsProducer.checkIntegrity();
+        normsProducer.checkIntegrity(mergeState.oneMerge);
       }
     }
     for (FieldInfo mergeFieldInfo : mergeState.mergeFieldInfos) {

lucene/core/src/java/org/apache/lucene/codecs/NormsProducer.java

@@ -19,6 +19,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.NumericDocValues;
 
 /**
@@ -48,6 +49,20 @@ protected NormsProducer() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this producer, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * Returns an instance optimized for merging. This instance may only be used from the thread that
    * acquires it.

lucene/core/src/java/org/apache/lucene/codecs/PointsReader.java

@@ -19,6 +19,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.PointValues;
 
 /**
@@ -41,6 +42,20 @@ protected PointsReader() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this reader, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * Return {@link PointValues} for the given {@code field}. The behavior is undefined if the given
    * field doesn't have points enabled on its {@link FieldInfo}.

lucene/core/src/java/org/apache/lucene/codecs/PointsWriter.java

@@ -223,7 +223,7 @@ public void merge(MergeState mergeState) throws IOException {
     for (PointsReader reader : mergeState.pointsReaders) {
       if (reader != null) {
         mergeState.checkAborted();
-        reader.checkIntegrity();
+        reader.checkIntegrity(mergeState.oneMerge);
       }
     }
     // merge field at a time

lucene/core/src/java/org/apache/lucene/codecs/PostingsReaderBase.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.ImpactsEnum;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.store.DataInput;
@@ -84,6 +85,20 @@ public abstract ImpactsEnum impacts(FieldInfo fieldInfo, BlockTermState state, i
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this reader, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   @Override
   public abstract void close() throws IOException;
 }

lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsReader.java

@@ -18,6 +18,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import org.apache.lucene.index.MergePolicy;
 import org.apache.lucene.index.StoredFieldVisitor;
 import org.apache.lucene.index.StoredFields;
 
@@ -47,6 +48,20 @@ protected StoredFieldsReader() {}
    */
   public abstract void checkIntegrity() throws IOException;
 
+  /**
+   * Checks consistency of this reader, periodically checking if the provided merge has been
+   * aborted. Subclasses should override this to propagate the abort check into expensive checksum
+   * computations.
+   *
+   * <p>The default implementation delegates to {@link #checkIntegrity()}.
+   *
+   * @param merge the merge to check for abort, or {@code null} for non-interruptible behavior
+   * @lucene.internal
+   */
+  public void checkIntegrity(MergePolicy.OneMerge merge) throws IOException {
+    checkIntegrity();
+  }
+
   /**
    * Returns an instance optimized for merging. This instance may not be cloned.
    *

lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsWriter.java

@@ -131,7 +131,7 @@ public int merge(MergeState mergeState) throws IOException {
     for (int i = 0; i < mergeState.storedFieldsReaders.length; i++) {
       StoredFieldsReader storedFieldsReader = mergeState.storedFieldsReaders[i];
       mergeState.checkAborted();
-      storedFieldsReader.checkIntegrity();
+      storedFieldsReader.checkIntegrity(mergeState.oneMerge);
       subs.add(
           new StoredFieldsMergeSub(
               new MergeVisitor(mergeState, i),