Move finishMerge to a finally block so it always runs (#14977)

There are some cases where KnnVectorsReader.finishMerge is not always run, if an exception is thrown. This can lead to ReadAdvice not being reset properly. This puts finishMerge in a finally block, so it is always run regardless of merge outcome.

Author: thecoop

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

@@ -109,18 +109,9 @@ public final void merge(MergeState mergeState) throws IOException {
         }
       }
     }
-    finishMerge(mergeState);
     finish();
   }
 
-  private void finishMerge(MergeState mergeState) throws IOException {
-    for (KnnVectorsReader reader : mergeState.knnVectorsReaders) {
-      if (reader != null) {
-        reader.finishMerge();
-      }
-    }
-  }
-
   /** Tracks state of one sub-reader that we are merging */
   private static class FloatVectorValuesSub extends DocIDMerger.Sub {
 

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

@@ -3447,23 +3447,26 @@ public void addIndexesReaderMerge(MergePolicy.OneMerge merge) throws IOException
             globalFieldNumberMap,
             context,
             intraMergeExecutor);
-
-    if (!merger.shouldMerge()) {
-      return;
-    }
-
-    merge.checkAborted();
-    synchronized (this) {
-      runningAddIndexesMerges.add(merger);
-    }
-    merge.mergeStartNS = System.nanoTime();
     try {
-      merger.merge(); // merge 'em
-    } finally {
+      if (!merger.shouldMerge()) {
+        return;
+      }
+
+      merge.checkAborted();
       synchronized (this) {
-        runningAddIndexesMerges.remove(merger);
-        notifyAll();
+        runningAddIndexesMerges.add(merger);
+      }
+      merge.mergeStartNS = System.nanoTime();
+      try {
+        merger.merge(); // merge 'em
+      } finally {
+        synchronized (this) {
+          runningAddIndexesMerges.remove(merger);
+          notifyAll();
+        }
       }
+    } finally {
+      merger.cleanupMerge();
     }
 
     merge.setMergeInfo(
@@ -5238,32 +5241,36 @@ public int length() {
               globalFieldNumberMap,
               context,
               intraMergeExecutor);
-      merge.info.setSoftDelCount(Math.toIntExact(softDeleteCount.get()));
-      merge.checkAborted();
-
       MergeState mergeState = merger.mergeState;
       MergeState.DocMap[] docMaps;
-      if (reorderDocMaps == null) {
-        docMaps = mergeState.docMaps;
-      } else {
-        // Since the reader was reordered, we passed a merged view to MergeState and from its
-        // perspective there is a single input segment to the merge and the
-        // SlowCompositeCodecReaderWrapper is effectively doing the merge.
-        assert mergeState.docMaps.length == 1
-            : "Got " + mergeState.docMaps.length + " docMaps, but expected 1";
-        MergeState.DocMap compactionDocMap = mergeState.docMaps[0];
-        docMaps = new MergeState.DocMap[reorderDocMaps.length];
-        for (int i = 0; i < docMaps.length; ++i) {
-          MergeState.DocMap reorderDocMap = reorderDocMaps[i];
-          docMaps[i] = docID -> compactionDocMap.get(reorderDocMap.get(docID));
+      try {
+        merge.info.setSoftDelCount(Math.toIntExact(softDeleteCount.get()));
+        merge.checkAborted();
+
+        if (reorderDocMaps == null) {
+          docMaps = mergeState.docMaps;
+        } else {
+          // Since the reader was reordered, we passed a merged view to MergeState and from its
+          // perspective there is a single input segment to the merge and the
+          // SlowCompositeCodecReaderWrapper is effectively doing the merge.
+          assert mergeState.docMaps.length == 1
+              : "Got " + mergeState.docMaps.length + " docMaps, but expected 1";
+          MergeState.DocMap compactionDocMap = mergeState.docMaps[0];
+          docMaps = new MergeState.DocMap[reorderDocMaps.length];
+          for (int i = 0; i < docMaps.length; ++i) {
+            MergeState.DocMap reorderDocMap = reorderDocMaps[i];
+            docMaps[i] = docID -> compactionDocMap.get(reorderDocMap.get(docID));
+          }
         }
-      }
 
-      merge.mergeStartNS = System.nanoTime();
+        merge.mergeStartNS = System.nanoTime();
 
-      // This is where all the work happens:
-      if (merger.shouldMerge()) {
-        merger.merge();
+        // This is where all the work happens:
+        if (merger.shouldMerge()) {
+          merger.merge();
+        }
+      } finally {
+        merger.cleanupMerge();
       }
 
       assert mergeState.segmentInfo == merge.info.info;

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

@@ -23,6 +23,7 @@
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.codecs.DocValuesConsumer;
 import org.apache.lucene.codecs.FieldsConsumer;
+import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.NormsConsumer;
 import org.apache.lucene.codecs.NormsProducer;
@@ -225,7 +226,7 @@ private void mergeTerms(SegmentWriteState segmentWriteState, SegmentReadState se
     }
   }
 
-  public void mergeFieldInfos() {
+  private void mergeFieldInfos() {
     for (FieldInfos readerFieldInfos : mergeState.fieldInfos) {
       for (FieldInfo fi : readerFieldInfos) {
         fieldInfosBuilder.add(fi);
@@ -324,4 +325,12 @@ private void mergeWithLogging(
               + " docs]");
     }
   }
+
+  void cleanupMerge() throws IOException {
+    for (KnnVectorsReader reader : mergeState.knnVectorsReaders) {
+      if (reader != null) {
+        reader.finishMerge();
+      }
+    }
+  }
 }

lucene/core/src/test/org/apache/lucene/index/TestDoc.java

@@ -242,6 +242,7 @@ private SegmentCommitInfo merge(
             new SameThreadExecutorService());
 
     merger.merge();
+    merger.cleanupMerge();
     r1.close();
     r2.close();
     si.setFiles(new HashSet<>(trackingDir.getCreatedFiles()));

lucene/core/src/test/org/apache/lucene/index/TestSegmentMerger.java

@@ -109,6 +109,7 @@ public void testMerge() throws IOException {
             newIOContext(random(), IOContext.merge(new MergeInfo(-1, -1, false, -1))),
             new SameThreadExecutorService());
     MergeState mergeState = merger.merge();
+    merger.cleanupMerge();
     int docsMerged = mergeState.segmentInfo.maxDoc();
     assertTrue(docsMerged == 2);
     // Should be able to open a new SegmentReader against the new directory

lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java

@@ -238,7 +238,7 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
     public void close() throws IOException {
       delegate.close();
       delegate.close(); // impls should be able to handle multiple closes
-      assert finishMergeCount.get() <= 0 || mergeInstanceCount.get() == finishMergeCount.get();
+      assert mergeInstanceCount.get() == finishMergeCount.get();
     }
 
     @Override