From f7608f14a5e5059bbfdb6d340e4b89558e425c44 Mon Sep 17 00:00:00 2001 From: Jeho Jeong Date: Tue, 7 Jul 2026 13:01:30 +0900 Subject: [PATCH 1/2] Check if merge is aborted during HNSW graph construction (#16367) Graph construction is a pure-CPU loop that performs no writes, so it never observed OneMerge's abort flag: IndexWriter#rollback/abortMerges blocked until the entire graph was built (tens of minutes for large vector segments). Wire MergeState#checkAborted from Lucene99HnswVectorsWriter through the HNSW graph mergers into the builders, and invoke it before every node insertion in HnswGraphBuilder#addGraphNodeInternal - the single point that all insertion paths (full rebuild, graph-join, concurrent workers) funnel through. Flush-time graph builds are unaffected (no check set). --- lucene/CHANGES.txt | 4 + .../lucene99/Lucene99HnswVectorsWriter.java | 13 +- .../util/hnsw/ConcurrentHnswMerger.java | 20 ++- .../apache/lucene/util/hnsw/HnswBuilder.java | 9 ++ .../util/hnsw/HnswConcurrentMergeBuilder.java | 8 + .../lucene/util/hnsw/HnswGraphBuilder.java | 10 ++ .../util/hnsw/IncrementalHnswGraphMerger.java | 21 +++ .../lucene/index/TestHnswMergeAbort.java | 143 ++++++++++++++++++ 8 files changed, 222 insertions(+), 6 deletions(-) create mode 100644 lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index df8de90e6235..2780ffce443a 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -178,6 +178,10 @@ 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) diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java index 1e330fee21b9..0b0bf5bda506 100644 --- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java +++ b/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 diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java index 4b6244c18522..0803c1389ed2 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/ConcurrentHnswMerger.java +++ b/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; } diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java index 38109c9c95e2..eaf06bbc5238 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java +++ b/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(); /** diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java index f1a4666fd7e6..26fb7e420998 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java +++ b/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) { diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java index e42b69ddb278..5a9fdc7a1ac5 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java +++ b/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--) { diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java index 3caa9600ed1a..d9dc4144472e 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/IncrementalHnswGraphMerger.java +++ b/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 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); } diff --git a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java new file mode 100644 index 000000000000..a79d1f49f04e --- /dev/null +++ b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java @@ -0,0 +1,143 @@ +/* + * 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.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +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; + +/** + * 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); + } + + /** + * 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); + } + + private void doTestRollbackDuringMerge(boolean withDeletes) throws Exception { + try (Directory dir = newDirectory()) { + IndexWriterConfig cfg = new IndexWriterConfig(); + cfg.setCodec(TestUtil.alwaysKnnVectorsFormat(new Lucene99HnswVectorsFormat(16, BEAM_WIDTH))); + 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(new Lucene99HnswVectorsFormat(16, BEAM_WIDTH))); + cfg2.setMergeScheduler(new ConcurrentMergeScheduler()); + cfg2.setInfoStream(latching); + IndexWriter w2 = new IndexWriter(dir, cfg2); + AtomicReference 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 < 5_000); + } finally { + merger.join(TimeUnit.MINUTES.toMillis(5)); + assertFalse("merge thread did not terminate", merger.isAlive()); + } + } + } +} From 33a14a5e9baa322c3981e4780e386fbf669229ca Mon Sep 17 00:00:00 2001 From: Jeho Jeong Date: Wed, 8 Jul 2026 13:38:04 +0900 Subject: [PATCH 2/2] Cover the concurrent merge path in TestHnswMergeAbort Add a variant with numMergeWorkers=2 so the abort check forwarding in HnswConcurrentMergeBuilder is exercised, and relax the rollback bound from 5s to 10s for slow CI machines. Without the fix the three variants block rollback for 27s (full rebuild), 31s (graph join) and 14s (concurrent) on this index. --- .../lucene/index/TestHnswMergeAbort.java | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java index a79d1f49f04e..70653dbc16d1 100644 --- a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java +++ b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java @@ -18,8 +18,11 @@ 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; @@ -29,6 +32,7 @@ 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 @@ -47,7 +51,7 @@ public class TestHnswMergeAbort extends LuceneTestCase { * scratch via {@code HnswGraphBuilder#addVectors}. */ public void testRollbackDuringFullRebuildMerge() throws Exception { - doTestRollbackDuringMerge(true); + doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); } /** @@ -55,13 +59,29 @@ public void testRollbackDuringFullRebuildMerge() throws Exception { * via {@code MergingHnswGraphBuilder}. */ public void testRollbackDuringGraphJoinMerge() throws Exception { - doTestRollbackDuringMerge(false); + doTestRollbackDuringMerge(false, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); } - private void doTestRollbackDuringMerge(boolean withDeletes) throws Exception { + /** + * 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(new Lucene99HnswVectorsFormat(16, BEAM_WIDTH))); + cfg.setCodec(TestUtil.alwaysKnnVectorsFormat(format)); cfg.setMergePolicy(NoMergePolicy.INSTANCE); try (IndexWriter w = new IndexWriter(dir, cfg)) { Random r = random(); @@ -108,7 +128,7 @@ public void close() {} }; IndexWriterConfig cfg2 = new IndexWriterConfig(); - cfg2.setCodec(TestUtil.alwaysKnnVectorsFormat(new Lucene99HnswVectorsFormat(16, BEAM_WIDTH))); + cfg2.setCodec(TestUtil.alwaysKnnVectorsFormat(format)); cfg2.setMergeScheduler(new ConcurrentMergeScheduler()); cfg2.setInfoStream(latching); IndexWriter w2 = new IndexWriter(dir, cfg2); @@ -133,7 +153,7 @@ public void close() {} "rollback() blocked for " + rollbackMillis + " ms waiting for HNSW graph construction to finish", - rollbackMillis < 5_000); + rollbackMillis < 10_000); } finally { merger.join(TimeUnit.MINUTES.toMillis(5)); assertFalse("merge thread did not terminate", merger.isAlive());