Check if merge is aborted during HNSW graph construction

lucene/CHANGES.txt

@@ -178,6 +178,9 @@ Optimizations
 
 Bug Fixes
 ---------------------
+* GITHUB#16367: HNSW graph construction now periodically checks whether the surrounding merge has
+  been aborted, so IndexWriter#rollback and abortMerges no longer block until the entire graph is
+  built. (Jeho Jeong)
 
 * GITHUB#14049: Randomize KNN codec params in RandomCodec. Fixes scalar quantization div-by-zero
   when all values are identical. (Mike Sokolov)

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java

@@ -473,7 +473,8 @@ private void buildAndWriteGraph(
             mergeState.intraMergeTaskExecutor == null
                 ? null
                 : new TaskExecutor(mergeState.intraMergeTaskExecutor),
-            numMergeWorkers);
+            numMergeWorkers,
+            mergeState::checkAborted);
     for (int i = 0; i < mergeState.liveDocs.length; i++) {
       if (hasVectorValues(mergeState.fieldInfos[i], fieldInfo.name)) {
         merger.addReader(
@@ -626,10 +627,11 @@ private HnswGraphMerger createGraphMerger(
       FieldInfo fieldInfo,
       RandomVectorScorerSupplier scorerSupplier,
       TaskExecutor parallelMergeTaskExecutor,
-      int numParallelMergeWorkers) {
+      int numParallelMergeWorkers,
+      IORunnable abortCheck) {
     if (mergeExec != null) {
       return new ConcurrentHnswMerger(
-          fieldInfo, scorerSupplier, M, beamWidth, mergeExec, numMergeWorkers);
+          fieldInfo, scorerSupplier, M, beamWidth, mergeExec, numMergeWorkers, abortCheck);
     }
     if (parallelMergeTaskExecutor != null && numParallelMergeWorkers > 1) {
       return new ConcurrentHnswMerger(
@@ -638,9 +640,10 @@ private HnswGraphMerger createGraphMerger(
           M,
           beamWidth,
           parallelMergeTaskExecutor,
-          numParallelMergeWorkers);
+          numParallelMergeWorkers,
+          abortCheck);
     }
-    return new IncrementalHnswGraphMerger(fieldInfo, scorerSupplier, M, beamWidth);
+    return new IncrementalHnswGraphMerger(fieldInfo, scorerSupplier, M, beamWidth, abortCheck);
   }
 
   @Override

lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java

@@ -29,6 +29,7 @@
 import org.apache.lucene.search.TaskExecutor;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.IORunnable;
 
 /** This merger merges graph in a concurrent manner, by using {@link HnswConcurrentMergeBuilder} */
 public class ConcurrentHnswMerger extends IncrementalHnswGraphMerger {
@@ -46,7 +47,24 @@ public ConcurrentHnswMerger(
       int beamWidth,
       TaskExecutor taskExecutor,
       int numWorker) {
-    super(fieldInfo, scorerSupplier, M, beamWidth);
+    this(fieldInfo, scorerSupplier, M, beamWidth, taskExecutor, numWorker, null);
+  }
+
+  /**
+   * @param fieldInfo FieldInfo for the field being merged
+   * @param abortCheck optional check invoked before every node insertion during graph construction;
+   *     may throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} to abort the
+   *     build when the surrounding merge has been aborted, or null
+   */
+  public ConcurrentHnswMerger(
+      FieldInfo fieldInfo,
+      RandomVectorScorerSupplier scorerSupplier,
+      int M,
+      int beamWidth,
+      TaskExecutor taskExecutor,
+      int numWorker,
+      IORunnable abortCheck) {
+    super(fieldInfo, scorerSupplier, M, beamWidth, abortCheck);
     this.taskExecutor = taskExecutor;
     this.numWorker = numWorker;
   }

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java

@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import org.apache.lucene.internal.hppc.IntHashSet;
+import org.apache.lucene.util.IORunnable;
 import org.apache.lucene.util.InfoStream;
 
 /**
@@ -46,6 +47,14 @@ public interface HnswBuilder {
   /** Set info-stream to output debugging information */
   void setInfoStream(InfoStream infoStream);
 
+  /**
+   * Sets a check that is invoked before every node insertion during graph construction. The check
+   * may throw an exception to abort the build promptly, e.g. {@link
+   * org.apache.lucene.index.MergePolicy.MergeAbortedException} when the merge that triggered the
+   * build has been aborted.
+   */
+  void setAbortCheck(IORunnable abortCheck);
+
   OnHeapHnswGraph getGraph();
 
   /**

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java

@@ -31,6 +31,7 @@
 import org.apache.lucene.search.TaskExecutor;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.IORunnable;
 import org.apache.lucene.util.InfoStream;
 
 /**
@@ -134,6 +135,13 @@ public void setInfoStream(InfoStream infoStream) {
     }
   }
 
+  @Override
+  public void setAbortCheck(IORunnable abortCheck) {
+    for (HnswBuilder worker : workers) {
+      worker.setAbortCheck(abortCheck);
+    }
+  }
+
   @Override
   public OnHeapHnswGraph getCompletedGraph() throws IOException {
     if (frozen == false) {

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java

@@ -33,6 +33,7 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.knn.KnnSearchStrategy;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.IORunnable;
 import org.apache.lucene.util.InfoStream;
 import org.apache.lucene.util.hnsw.HnswUtil.Component;
 
@@ -83,6 +84,7 @@ public class HnswGraphBuilder implements HnswBuilder {
   protected final HnswLock hnswLock;
 
   protected InfoStream infoStream = InfoStream.getDefault();
+  private IORunnable abortCheck;
   protected boolean frozen;
 
   /**
@@ -209,6 +211,11 @@ public void setInfoStream(InfoStream infoStream) {
     this.infoStream = infoStream;
   }
 
+  @Override
+  public void setAbortCheck(IORunnable abortCheck) {
+    this.abortCheck = abortCheck;
+  }
+
   @Override
   public OnHeapHnswGraph getCompletedGraph() throws IOException {
     if (!frozen) {
@@ -280,6 +287,9 @@ private void addGraphNodeInternal(int node, UpdateableRandomVectorScorer scorer,
     if (frozen) {
       throw new IllegalStateException("Graph builder is already frozen");
     }
+    if (abortCheck != null) {
+      abortCheck.run();
+    }
     final int nodeLevel = getRandomGraphLevel(ml, random);
     // first add nodes to all levels
     for (int level = nodeLevel; level >= 0; level--) {

lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java

@@ -33,6 +33,7 @@
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.FixedBitSet;
+import org.apache.lucene.util.IORunnable;
 import org.apache.lucene.util.InfoStream;
 
 /**
@@ -46,6 +47,7 @@ public class IncrementalHnswGraphMerger implements HnswGraphMerger {
   protected final RandomVectorScorerSupplier scorerSupplier;
   protected final int M;
   protected final int beamWidth;
+  protected final IORunnable abortCheck;
 
   protected List<GraphReader> graphReaders = new ArrayList<>();
   protected GraphReader largestGraphReader;
@@ -71,10 +73,26 @@ protected record GraphReader(
    */
   public IncrementalHnswGraphMerger(
       FieldInfo fieldInfo, RandomVectorScorerSupplier scorerSupplier, int M, int beamWidth) {
+    this(fieldInfo, scorerSupplier, M, beamWidth, null);
+  }
+
+  /**
+   * @param fieldInfo FieldInfo for the field being merged
+   * @param abortCheck optional check invoked before every node insertion during graph construction;
+   *     may throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} to abort the
+   *     build when the surrounding merge has been aborted, or null
+   */
+  public IncrementalHnswGraphMerger(
+      FieldInfo fieldInfo,
+      RandomVectorScorerSupplier scorerSupplier,
+      int M,
+      int beamWidth,
+      IORunnable abortCheck) {
     this.fieldInfo = fieldInfo;
     this.scorerSupplier = scorerSupplier;
     this.M = M;
     this.beamWidth = beamWidth;
+    this.abortCheck = abortCheck;
   }
 
   /**
@@ -211,6 +229,9 @@ public OnHeapHnswGraph merge(
       KnnVectorValues mergedVectorValues, InfoStream infoStream, int maxOrd) throws IOException {
     HnswBuilder builder = createBuilder(mergedVectorValues, maxOrd);
     builder.setInfoStream(infoStream);
+    if (abortCheck != null) {
+      builder.setAbortCheck(abortCheck);
+    }
     return builder.build(maxOrd);
   }
 

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

@@ -0,0 +1,163 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.index;
+
+import java.util.Random;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.InfoStream;
+import org.apache.lucene.util.NamedThreadFactory;
+
+/**
+ * Tests that aborting a merge (e.g. via {@link IndexWriter#rollback()}) promptly interrupts HNSW
+ * graph construction instead of blocking until the entire graph is built.
+ */
+public class TestHnswMergeAbort extends LuceneTestCase {
+
+  private static final int DIM = 96;
+  private static final int SEGMENTS = 4;
+  private static final int DOCS_PER_SEGMENT = 12_000;
+  private static final int BEAM_WIDTH = 250;
+
+  /**
+   * Every segment carries more than {@code IncrementalHnswGraphMerger#DELETE_PCT_THRESHOLD}
+   * deletions, so no source graph is eligible as a base and the merged graph is rebuilt from
+   * scratch via {@code HnswGraphBuilder#addVectors}.
+   */
+  public void testRollbackDuringFullRebuildMerge() throws Exception {
+    doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH));
+  }
+
+  /**
+   * All segments are deletion-free, so the merged graph is produced by joining the source graphs
+   * via {@code MergingHnswGraphBuilder}.
+   */
+  public void testRollbackDuringGraphJoinMerge() throws Exception {
+    doTestRollbackDuringMerge(false, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH));
+  }
+
+  /**
+   * The merged graph is built by {@code HnswConcurrentMergeBuilder} workers ({@code numMergeWorkers
+   * > 1}), which must forward the abort check to every worker.
+   */
+  public void testRollbackDuringConcurrentMerge() throws Exception {
+    ExecutorService mergeExec =
+        Executors.newFixedThreadPool(2, new NamedThreadFactory("hnsw-merge-worker"));
+    try {
+      doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, 2, mergeExec));
+    } finally {
+      mergeExec.shutdown();
+      assertTrue(mergeExec.awaitTermination(30, TimeUnit.SECONDS));
+    }
+  }
+
+  private void doTestRollbackDuringMerge(boolean withDeletes, KnnVectorsFormat format)
+      throws Exception {
+    try (Directory dir = newDirectory()) {
+      IndexWriterConfig cfg = new IndexWriterConfig();
+      cfg.setCodec(TestUtil.alwaysKnnVectorsFormat(format));
+      cfg.setMergePolicy(NoMergePolicy.INSTANCE);
+      try (IndexWriter w = new IndexWriter(dir, cfg)) {
+        Random r = random();
+        int id = 0;
+        for (int s = 0; s < SEGMENTS; s++) {
+          for (int i = 0; i < DOCS_PER_SEGMENT; i++, id++) {
+            Document doc = new Document();
+            doc.add(new StringField("id", Integer.toString(id), Field.Store.NO));
+            float[] v = new float[DIM];
+            for (int j = 0; j < DIM; j++) {
+              v[j] = r.nextFloat();
+            }
+            doc.add(new KnnFloatVectorField("v", v));
+            w.addDocument(doc);
+          }
+          w.flush();
+        }
+        if (withDeletes) {
+          // delete half of every segment, above the 40% base-graph eligibility threshold
+          for (int d = 0; d < id; d += 2) {
+            w.deleteDocuments(new Term("id", Integer.toString(d)));
+          }
+        }
+        w.commit();
+      }
+
+      CountDownLatch buildStarted = new CountDownLatch(1);
+      InfoStream latching =
+          new InfoStream() {
+            @Override
+            public void message(String component, String message) {
+              if ("HNSW".equals(component) && message.startsWith("build graph")) {
+                buildStarted.countDown();
+              }
+            }
+
+            @Override
+            public boolean isEnabled(String component) {
+              return "HNSW".equals(component);
+            }
+
+            @Override
+            public void close() {}
+          };
+
+      IndexWriterConfig cfg2 = new IndexWriterConfig();
+      cfg2.setCodec(TestUtil.alwaysKnnVectorsFormat(format));
+      cfg2.setMergeScheduler(new ConcurrentMergeScheduler());
+      cfg2.setInfoStream(latching);
+      IndexWriter w2 = new IndexWriter(dir, cfg2);
+      AtomicReference<Throwable> mergeFailure = new AtomicReference<>();
+      Thread merger =
+          new Thread(
+              () -> {
+                try {
+                  w2.forceMerge(1);
+                } catch (Throwable t) {
+                  mergeFailure.set(t);
+                }
+              });
+      merger.start();
+      try {
+        assertTrue(
+            "HNSW graph construction never started", buildStarted.await(120, TimeUnit.SECONDS));
+        long t0 = System.nanoTime();
+        w2.rollback();
+        long rollbackMillis = (System.nanoTime() - t0) / 1_000_000;
+        assertTrue(
+            "rollback() blocked for "
+                + rollbackMillis
+                + " ms waiting for HNSW graph construction to finish",
+            rollbackMillis < 10_000);
+      } finally {
+        merger.join(TimeUnit.MINUTES.toMillis(5));
+        assertFalse("merge thread did not terminate", merger.isAlive());
+      }
+    }
+  }
+}