Abort stateless merge reads on shard close to unblock HNSW merges

docs/changelog/152342.yaml

@@ -0,0 +1,5 @@
+area: Search
+issues: []
+pr: 152342
+summary: Abort stateless merge reads on shard close to unblock HNSW merges
+type: feature

server/src/main/java/org/elasticsearch/index/engine/ElasticsearchMergeScheduler.java

@@ -24,4 +24,15 @@ public interface ElasticsearchMergeScheduler {
     void refreshConfig();
 
     MergeScheduler getMergeScheduler();
+
+    /**
+     * Marks all queued and running merges as aborted before abort-merge-reads is signalled.
+     * Implementations backed by {@link ThreadPoolMergeScheduler} call
+     * {@link org.apache.lucene.index.MergePolicy.OneMerge#setAborted()} on every pending merge so
+     * that {@code merge.isAborted()} is {@code true} before any
+     * {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} can be thrown during
+     * compound-file creation, preventing a corrupt-segment write in Lucene's {@code mergeMiddle()}.
+     * Other implementations may provide a no-op default.
+     */
+    default void abortAllMerges() {}
 }

server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java

@@ -786,6 +786,20 @@ Map<MergePolicy.OneMerge, MergeTask> getRunningMergeTasks() {
         return runningMergeTasks;
     }
 
+    /**
+     * Marks all queued and running merges as aborted. This should be called before signalling
+     * abort-merge-reads to ensure {@code merge.isAborted()} is {@code true} before any
+     * {@link MergePolicy.MergeAbortedException} can be thrown during compound-file creation,
+     * preventing Lucene's {@code mergeMiddle()} from committing a segment with
+     * {@code useCompoundFile=true} but without the corresponding {@code .cfs} file.
+     */
+    @Override
+    public synchronized void abortAllMerges() {
+        for (MergePolicy.OneMerge merge : runningMergeTasks.keySet()) {
+            merge.setAborted();
+        }
+    }
+
     private static double nsToSec(long ns) {
         return ns / (double) TimeUnit.SECONDS.toNanos(1);
     }

x-pack/plugin/stateless/src/internalClusterTest/java/org/elasticsearch/xpack/stateless/lucene/StatelessMergeAbortIT.java

@@ -0,0 +1,212 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.stateless.lucene;
+
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.elasticsearch.action.admin.indices.forcemerge.ForceMergeRequest;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xpack.stateless.AbstractStatelessPluginIntegTestCase;
+import org.elasticsearch.xpack.stateless.TestUtils;
+import org.elasticsearch.xpack.stateless.cache.StatelessSharedBlobCacheService;
+import org.elasticsearch.xpack.stateless.commits.BlobFileRanges;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.UnaryOperator;
+
+import static org.elasticsearch.index.engine.ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING;
+import static org.elasticsearch.search.vectors.KnnSearchBuilderTests.randomVector;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.lessThan;
+
+/**
+ * End-to-end tests for merge-read abort during index deletion.
+ * <p>
+ * The latch-based test follows {@code StatelessComponentsOrderIT}: a merge thread blocks in blob-cache
+ * {@code openInput} while the index is deleted. Deletion must complete promptly after
+ * {@link IndexDirectory#abortMergeReads()} is signalled from the merge scheduler close path.
+ * BBQ HNSW timing is additionally covered in {@code IndexEngineTests}.
+ */
+public class StatelessMergeAbortIT extends AbstractStatelessPluginIntegTestCase {
+
+    private static final String VECTOR_FIELD = "vector";
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        var plugins = new ArrayList<>(super.nodePlugins());
+        plugins.remove(TestUtils.StatelessPluginWithTrialLicense.class);
+        plugins.add(TestStatelessPlugin.class);
+        return plugins;
+    }
+
+    public void testDeleteIndexDuringMergeAbortsReadsQuickly() throws Exception {
+        startMasterOnlyNode();
+        var nodeSettings = Settings.builder()
+            .put(USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.getKey(), true)
+            .put(disableIndexingDiskAndMemoryControllersNodeSettings())
+            .build();
+        final String indexNode = startIndexNode(nodeSettings);
+        ensureStableCluster(2);
+
+        final String indexName = randomIdentifier();
+        createIndex(indexName, indexSettings(1, 0).put(IndexSettings.INDEX_REFRESH_INTERVAL_SETTING.getKey(), -1).build());
+        ensureGreen(indexName);
+
+        final TestStatelessPlugin plugin = findPlugin(indexNode, TestStatelessPlugin.class);
+        final var indicesService = internalCluster().getInstance(IndicesService.class, indexNode);
+        final IndexShard indexShard = findIndexShard(indexName);
+        final IndexDirectory indexDirectory = IndexDirectory.unwrapDirectory(indexShard.store().directory());
+
+        for (int i = 0; i < 11; i++) {
+            indexDocs(indexName, 10);
+            flush(indexName);
+        }
+
+        safeAwait(plugin.mergeReadStartedLatch);
+        assertThat(indexDirectory.shouldAbortMergeReads(), is(false));
+
+        final long deleteStartNanos = System.nanoTime();
+        safeGet(client().admin().indices().prepareDelete(indexName).execute());
+        assertThat(System.nanoTime() - deleteStartNanos, lessThan(TimeValue.timeValueSeconds(30).nanos()));
+        assertNull(indicesService.indexService(indexShard.shardId().getIndex()));
+        assertThat(indexDirectory.shouldAbortMergeReads(), is(true));
+
+        plugin.resumeMergeReadsLatch.countDown();
+    }
+
+    public void testDeleteIndexDuringBbqHnswMergeCompletesQuickly() throws Exception {
+        final String indexNode = startMasterAndIndexNode(disableIndexingDiskAndMemoryControllersNodeSettings());
+        ensureStableCluster(1);
+
+        final String indexName = randomIdentifier();
+        final int numDims = randomIntBetween(64, 128);
+        createBbqHnswIndex(indexName, numDims);
+        ensureGreen(indexName);
+
+        final TestStatelessPlugin plugin = findPlugin(indexNode, TestStatelessPlugin.class);
+        final IndexShard indexShard = findIndexShard(indexName);
+        final IndexDirectory indexDirectory = IndexDirectory.unwrapDirectory(indexShard.store().directory());
+
+        for (int i = 0; i < randomIntBetween(80, 150); i++) {
+            indexVectorDocs(indexName, 1, numDims);
+            if (i % 10 == 0) {
+                assertNoFailures(indicesAdmin().prepareFlush(indexName).execute().get());
+            }
+        }
+        assertNoFailures(indicesAdmin().prepareFlush(indexName).execute().get());
+
+        indexDirectory.setMergeReadCallback(() -> {
+            plugin.mergeReadStartedLatch.countDown();
+            safeAwait(plugin.resumeMergeReadsLatch);
+        });
+
+        var mergeThread = new Thread(() -> {
+            try {
+                indexShard.forceMerge(new ForceMergeRequest().maxNumSegments(1).flush(false));
+            } catch (Exception e) {
+                // merge may be aborted during index deletion
+            }
+        }, "bbq-hnsw-force-merge");
+        mergeThread.start();
+
+        safeAwait(plugin.mergeReadStartedLatch);
+
+        final long deleteStartNanos = System.nanoTime();
+        safeGet(client().admin().indices().prepareDelete(indexName).execute());
+        assertThat(System.nanoTime() - deleteStartNanos, lessThan(TimeValue.timeValueSeconds(30).nanos()));
+
+        assertBusy(() -> assertThat(indexDirectory.shouldAbortMergeReads(), is(true)));
+
+        plugin.resumeMergeReadsLatch.countDown();
+
+        mergeThread.join(TimeUnit.SECONDS.toMillis(30));
+        assertFalse("merge thread should finish after index deletion", mergeThread.isAlive());
+    }
+
+    private void createBbqHnswIndex(String indexName, int numDims) throws IOException {
+        XContentBuilder mapping = XContentFactory.jsonBuilder()
+            .startObject()
+            .startObject("properties")
+            .startObject(VECTOR_FIELD)
+            .field("type", "dense_vector")
+            .field("dims", numDims)
+            .field("index", true)
+            .field("similarity", "cosine")
+            .startObject("index_options")
+            .field("type", "bbq_hnsw")
+            .endObject()
+            .endObject()
+            .endObject()
+            .endObject();
+        assertAcked(
+            indicesAdmin().prepareCreate(indexName)
+                .setSettings(
+                    indexSettings(1, 0).put(IndexSettings.INDEX_REFRESH_INTERVAL_SETTING.getKey(), -1)
+                        .put("index.compound_format", false)
+                        .build()
+                )
+                .setMapping(mapping)
+                .get()
+        );
+    }
+
+    private void indexVectorDocs(String indexName, int numDocs, int numDims) {
+        indexDocs(indexName, numDocs, UnaryOperator.identity(), null, () -> Map.of(VECTOR_FIELD, randomVector(numDims)));
+    }
+
+    public static class TestStatelessPlugin extends TestUtils.StatelessPluginWithTrialLicense {
+
+        private final CountDownLatch mergeReadStartedLatch = new CountDownLatch(1);
+        private final CountDownLatch resumeMergeReadsLatch = new CountDownLatch(1);
+
+        public TestStatelessPlugin(Settings settings) {
+            super(settings);
+        }
+
+        @Override
+        protected IndexBlobStoreCacheDirectory createIndexBlobStoreCacheDirectory(
+            StatelessSharedBlobCacheService cacheService,
+            ShardId shardId
+        ) {
+            return new IndexBlobStoreCacheDirectory(cacheService, shardId) {
+                @Override
+                protected IndexInput doOpenInput(String name, IOContext context, BlobFileRanges blobFileRanges) {
+                    if (ThreadPool.Names.MERGE.equals(EsExecutors.executorName(Thread.currentThread()))) {
+                        mergeReadStartedLatch.countDown();
+                        try {
+                            if (resumeMergeReadsLatch.await(30, TimeUnit.SECONDS) == false) {
+                                throw new IllegalStateException("timed out waiting to resume merge reads in test");
+                            }
+                        } catch (InterruptedException e) {
+                            Thread.currentThread().interrupt();
+                            throw new IllegalStateException("interrupted waiting to resume merge reads in test", e);
+                        }
+                    }
+                    return super.doOpenInput(name, context, blobFileRanges);
+                }
+            };
+        }
+    }
+}

x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/cache/reader/CacheFileReader.java

@@ -38,6 +38,7 @@
 import java.nio.ByteBuffer;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BooleanSupplier;
 import java.util.function.LongSupplier;
 
 import static org.elasticsearch.blobcache.BlobCacheMetrics.CACHE_POPULATION_SOURCE_ATTRIBUTE_KEY;
@@ -95,6 +96,7 @@ public class CacheFileReader {
     private final long exclusiveStart;
     private final long exclusiveEnd;
     private final boolean hasSearchRole;
+    private final BooleanSupplier mergeReadAbortSupplier;
 
     public CacheFileReader(
         StatelessSharedBlobCacheService.CacheFile cacheFile,
@@ -113,7 +115,8 @@ public CacheFileReader(
             SharedBytes.MADV_NORMAL,
             0,
             0,
-            false
+            false,
+            () -> false
         );
     }
 
@@ -143,7 +146,8 @@ public CacheFileReader(
             contextToAdvice(context, hasSearchRole),
             0,
             Long.MAX_VALUE,
-            hasSearchRole
+            hasSearchRole,
+            () -> false
         );
     }
 
@@ -157,7 +161,8 @@ private CacheFileReader(
         int desiredAdvice,
         long exclusiveStart,
         long exclusiveEnd,
-        boolean hasSearchRole
+        boolean hasSearchRole,
+        BooleanSupplier mergeReadAbortSupplier
     ) {
         this.cacheFile = Objects.requireNonNull(cacheFile);
         this.cacheBlobReader = Objects.requireNonNull(cacheBlobReader);
@@ -169,6 +174,7 @@ private CacheFileReader(
         this.exclusiveStart = exclusiveStart;
         this.exclusiveEnd = exclusiveEnd;
         this.hasSearchRole = hasSearchRole;
+        this.mergeReadAbortSupplier = Objects.requireNonNull(mergeReadAbortSupplier);
     }
 
     /**
@@ -185,10 +191,41 @@ public CacheFileReader copy() {
             desiredMAdvice,
             exclusiveStart,
             exclusiveEnd,
-            hasSearchRole
+            hasSearchRole,
+            mergeReadAbortSupplier
+        );
+    }
+
+    /**
+     * Returns a copy of this reader with the given merge-read abort supplier. When the supplier returns {@code true},
+     * {@link BlobCacheIndexInput} will throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} on any
+     * read attempted with a MERGE {@link IOContext}, stopping I/O on behalf of a closing shard.
+     * The supplier is propagated to all subsequent {@link #copy()} and {@link #copyWithContext(IOContext, long, long)}
+     * instances so that compound-file sub-file slices inherit it automatically.
+     */
+    public CacheFileReader withMergeReadAbortSupplier(BooleanSupplier supplier) {
+        return new CacheFileReader(
+            cacheFile,
+            cacheBlobReader,
+            blobFileRanges,
+            blobCacheMetrics,
+            relativeTimeInMillisSupplier,
+            regionSize,
+            desiredMAdvice,
+            exclusiveStart,
+            exclusiveEnd,
+            hasSearchRole,
+            Objects.requireNonNull(supplier)
         );
     }
 
+    /**
+     * Returns {@code true} if the merge-read abort supplier indicates that in-progress merge reads should be aborted.
+     */
+    public boolean isMergeReadAborted() {
+        return mergeReadAbortSupplier.getAsBoolean();
+    }
+
     /**
      * Returns a copy of this reader for a sub-file within a compound ({@code .cfs}) blob.
      * Computes the range of cache regions that are exclusively occupied by the sub-file:
@@ -220,7 +257,8 @@ public CacheFileReader copyWithContext(IOContext context, long subFileOffset, lo
             advice,
             exclStart,
             exclEnd,
-            hasSearchRole
+            hasSearchRole,
+            mergeReadAbortSupplier
         );
     }