Speed up TestHnswMergeAbort and make its assertion deterministic

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();
   }
 

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,52 +36,71 @@
 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.
+ *
+ * <p>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}
    * 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));
+    doTestRollbackDuringMerge(
+        true,
+        new Lucene99HnswVectorsFormat(16, BEAM_WIDTH, TINY_SEGMENTS_THRESHOLD),
+        "build graph from " + LIVE_DOCS_AFTER_DELETES + " vectors");
   }
 
   /**
    * 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));
+    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<String> 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());