From 1b67f52986c80dc32d3ed9495245e23700843d39 Mon Sep 17 00:00:00 2001 From: Jeho Jeong Date: Thu, 16 Jul 2026 20:32:33 +0900 Subject: [PATCH] Speed up TestHnswMergeAbort and make its assertion deterministic Instead of asserting that rollback() returns within a time limit, hold the merge thread at the start of the graph build until the merge is marked aborted, then assert that the build never ran to completion, observed through InfoStream. This does not depend on segment size or machine speed, so segments shrink from 12k docs to 1k and beam width from 250 to 100. MergingHnswGraphBuilder now emits a completion message like the other two builders, so the graph join path has a completion signal to assert on. --- .../util/hnsw/MergingHnswGraphBuilder.java | 13 +++ .../lucene/index/TestHnswMergeAbort.java | 92 +++++++++++++++---- 2 files changed, 85 insertions(+), 20 deletions(-) diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java index f5506dbaf8d3..e41979b1acd1 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.Arrays; +import java.util.Locale; import org.apache.lucene.internal.hppc.IntHashSet; import org.apache.lucene.util.BitSet; @@ -110,6 +111,7 @@ public OnHeapHnswGraph build(int maxOrd) throws IOException { if (frozen) { throw new IllegalStateException("This HnswGraphBuilder is frozen and cannot be updated"); } + long startTimeNs = System.nanoTime(); if (infoStream.isEnabled(HNSW_COMPONENT)) { String graphSizes = ""; for (HnswGraph g : graphs) { @@ -139,6 +141,17 @@ public OnHeapHnswGraph build(int maxOrd) throws IOException { } } + if (infoStream.isEnabled(HNSW_COMPONENT)) { + double elapsedMs = (System.nanoTime() - startTimeNs) / 1_000_000.0; + infoStream.message( + HNSW_COMPONENT, + String.format( + Locale.ROOT, + "merge completed: %d vectors from merging %d graphs in %.2f ms", + maxOrd, + graphs.length, + elapsedMs)); + } return getCompletedGraph(); } 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 70653dbc16d1..ba1e95f05df8 100644 --- a/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java +++ b/lucene/core/src/test/org/apache/lucene/index/TestHnswMergeAbort.java @@ -16,11 +16,14 @@ */ package org.apache.lucene.index; +import java.util.ArrayList; +import java.util.List; 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.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat; @@ -33,17 +36,26 @@ import org.apache.lucene.tests.util.TestUtil; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.NamedThreadFactory; +import org.apache.lucene.util.ThreadInterruptedException; /** - * 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. + * Tests that aborting a merge (e.g. via {@link IndexWriter#rollback()}) interrupts HNSW graph + * construction instead of building the entire graph. + * + *

The merge thread is held at the start of the graph build until the merge is marked aborted, + * and the test then asserts that the build never ran to completion, which holds at any segment + * size. The abort exception is deliberately not asserted on: it is swallowed as expected control + * flow, and a small unchecked build can finish without ever throwing. */ 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; + private static final int DOCS_PER_SEGMENT = 1_000; + private static final int BEAM_WIDTH = 100; + // always build a graph, no matter how small the segment is + private static final int TINY_SEGMENTS_THRESHOLD = 0; + private static final int LIVE_DOCS_AFTER_DELETES = SEGMENTS * DOCS_PER_SEGMENT / 2; /** * Every segment carries more than {@code IncrementalHnswGraphMerger#DELETE_PCT_THRESHOLD} @@ -51,7 +63,10 @@ public class TestHnswMergeAbort extends LuceneTestCase { * scratch via {@code HnswGraphBuilder#addVectors}. */ public void testRollbackDuringFullRebuildMerge() throws Exception { - doTestRollbackDuringMerge(true, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); + doTestRollbackDuringMerge( + true, + new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, TINY_SEGMENTS_THRESHOLD), + "build graph from " + LIVE_DOCS_AFTER_DELETES + " vectors"); } /** @@ -59,26 +74,33 @@ public void testRollbackDuringFullRebuildMerge() throws Exception { * via {@code MergingHnswGraphBuilder}. */ public void testRollbackDuringGraphJoinMerge() throws Exception { - doTestRollbackDuringMerge(false, new Lucene99HnswVectorsFormat(16, BEAM_WIDTH)); + doTestRollbackDuringMerge( + false, + new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, TINY_SEGMENTS_THRESHOLD), + "build graph from merging " + SEGMENTS + " graphs"); } /** * The merged graph is built by {@code HnswConcurrentMergeBuilder} workers ({@code numMergeWorkers - * > 1}), which must forward the abort check to every worker. + * > 1}), which must forward the abort check to every worker. The first worker to insert a node + * trips the forwarded check. */ 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)); + doTestRollbackDuringMerge( + true, + new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, 2, mergeExec, TINY_SEGMENTS_THRESHOLD), + "build graph from " + LIVE_DOCS_AFTER_DELETES + " vectors, with 2 workers"); } finally { mergeExec.shutdown(); assertTrue(mergeExec.awaitTermination(30, TimeUnit.SECONDS)); } } - private void doTestRollbackDuringMerge(boolean withDeletes, KnnVectorsFormat format) - throws Exception { + private void doTestRollbackDuringMerge( + boolean withDeletes, KnnVectorsFormat format, String expectedBuildStart) throws Exception { try (Directory dir = newDirectory()) { IndexWriterConfig cfg = new IndexWriterConfig(); cfg.setCodec(TestUtil.alwaysKnnVectorsFormat(format)); @@ -109,18 +131,36 @@ private void doTestRollbackDuringMerge(boolean withDeletes, KnnVectorsFormat for } CountDownLatch buildStarted = new CountDownLatch(1); + CountDownLatch mergeAborted = new CountDownLatch(1); + AtomicBoolean releasedAfterAbort = new AtomicBoolean(); + List hnswMessages = new ArrayList<>(); InfoStream latching = new InfoStream() { @Override public void message(String component, String message) { - if ("HNSW".equals(component) && message.startsWith("build graph")) { - buildStarted.countDown(); + if ("HNSW".equals(component)) { + synchronized (hnswMessages) { + hnswMessages.add(message); + } + if (message.startsWith("build graph") && buildStarted.getCount() > 0) { + buildStarted.countDown(); + // hold the merge thread until the merge is marked aborted, so the abort + // always lands mid-build + try { + releasedAfterAbort.set(mergeAborted.await(2, TimeUnit.MINUTES)); + } catch (InterruptedException e) { + throw new ThreadInterruptedException(e); + } + } + } else if ("IW".equals(component) && message.startsWith("now wait for")) { + // IndexWriter#abortMerges emits this after marking every running merge aborted + mergeAborted.countDown(); } } @Override public boolean isEnabled(String component) { - return "HNSW".equals(component); + return "HNSW".equals(component) || "IW".equals(component); } @Override @@ -146,14 +186,26 @@ public void close() {} 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); + assertTrue("merge thread was not released by the abort signal", releasedAfterAbort.get()); + synchronized (hnswMessages) { + String buildStart = null; + for (String message : hnswMessages) { + if (message.startsWith("build graph")) { + buildStart = message; + break; + } + } + assertNotNull("no graph build started", buildStart); + assertTrue( + "merge took an unexpected graph build path: " + buildStart, + buildStart.startsWith(expectedBuildStart)); + for (String message : hnswMessages) { + assertFalse( + "HNSW graph build ran to completion despite the aborted merge: " + message, + message.startsWith("addVectors [") || message.startsWith("merge completed:")); + } + } } finally { merger.join(TimeUnit.MINUTES.toMillis(5)); assertFalse("merge thread did not terminate", merger.isAlive());