diff --git a/docs/changelog/152342.yaml b/docs/changelog/152342.yaml new file mode 100644 index 0000000000000..f4cb92ed2ba0a --- /dev/null +++ b/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 diff --git a/server/src/main/java/org/elasticsearch/index/engine/ElasticsearchMergeScheduler.java b/server/src/main/java/org/elasticsearch/index/engine/ElasticsearchMergeScheduler.java index ac72c7a21da75..343c46fcfaf51 100644 --- a/server/src/main/java/org/elasticsearch/index/engine/ElasticsearchMergeScheduler.java +++ b/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() {} } diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java index 236584a3f0db4..f5a3046733241 100644 --- a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java +++ b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java @@ -786,6 +786,20 @@ Map 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); } diff --git a/x-pack/plugin/stateless/src/internalClusterTest/java/org/elasticsearch/xpack/stateless/lucene/StatelessMergeAbortIT.java b/x-pack/plugin/stateless/src/internalClusterTest/java/org/elasticsearch/xpack/stateless/lucene/StatelessMergeAbortIT.java new file mode 100644 index 0000000000000..ee72e3c5c0787 --- /dev/null +++ b/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. + *

+ * 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> 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); + } + }; + } + } +} diff --git a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/cache/reader/CacheFileReader.java b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/cache/reader/CacheFileReader.java index f9ed810f39278..f7eafa4fddb8b 100644 --- a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/cache/reader/CacheFileReader.java +++ b/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 ); } diff --git a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/engine/IndexEngine.java b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/engine/IndexEngine.java index 6cc430fc59a25..c879b5bcd5340 100644 --- a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/engine/IndexEngine.java +++ b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/engine/IndexEngine.java @@ -16,6 +16,8 @@ import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.MergePolicy; +import org.apache.lucene.index.MergeScheduler; +import org.apache.lucene.index.MergeTrigger; import org.apache.lucene.index.OneMergeWrappingMergePolicy; import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.index.SegmentReadState; @@ -54,6 +56,7 @@ import org.elasticsearch.index.engine.ThreadPoolMergeExecutorService; import org.elasticsearch.index.mapper.DateFieldMapper; import org.elasticsearch.index.mapper.ParsedDocument; +import org.elasticsearch.index.merge.MergeStats; import org.elasticsearch.index.merge.OnGoingMerge; import org.elasticsearch.index.seqno.LocalCheckpointTracker; import org.elasticsearch.index.seqno.SequenceNumbers; @@ -87,6 +90,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -975,8 +979,114 @@ protected ElasticsearchMergeScheduler createMergeScheduler( this::estimateMergeBytes, mergeMetrics ); - } else { - return super.createMergeScheduler(shardId, indexSettings, threadPoolMergeExecutorService, mergeMetrics); + } + return new AbortOnCloseElasticsearchMergeScheduler( + super.createMergeScheduler(shardId, indexSettings, threadPoolMergeExecutorService, mergeMetrics), + this::signalAbortMergeReads + ); + } + + /** + * Signal the abort before the IndexWriter rollback starts. Lucene 10 calls {@code abortMerges()} + * before {@code mergeScheduler.close()}, so without this early signal the abort would only fire + * after all running merges finish — defeating the purpose of the interrupt. + * + *

All queued and running merges are also marked as aborted + * ({@link MergePolicy.OneMerge#setAborted()}) before the signal fires. Without this, a merge in + * the compound-file creation phase can receive {@link MergePolicy.MergeAbortedException} while + * {@code merge.isAborted()} is still {@code false}, causing Lucene's {@code mergeMiddle()} to + * skip its abort guard and commit a segment with {@code useCompoundFile=true} but without the + * corresponding {@code .cfs} file. + */ + @Override + public void close() throws IOException { + getMergeScheduler().abortAllMerges(); + signalAbortMergeReads(); + super.close(); + } + + private void signalAbortMergeReads() { + IndexDirectory.tryUnwrapDirectory(store.directory()).ifPresent(IndexDirectory::abortMergeReads); + } + + /** + * Wraps the Lucene {@link MergeScheduler} returned by {@link #getMergeScheduler()} so that + * {@link IndexDirectory#abortMergeReads()} is signalled before the delegate closes. + */ + private static final class AbortOnCloseElasticsearchMergeScheduler implements ElasticsearchMergeScheduler { + + private final ElasticsearchMergeScheduler delegate; + private final Runnable abortAction; + private MergeScheduler mergeSchedulerWrapper; + + private AbortOnCloseElasticsearchMergeScheduler(ElasticsearchMergeScheduler delegate, Runnable abortAction) { + this.delegate = delegate; + this.abortAction = abortAction; + } + + @Override + public Set onGoingMerges() { + return delegate.onGoingMerges(); + } + + @Override + public MergeStats stats() { + return delegate.stats(); + } + + @Override + public void refreshConfig() { + delegate.refreshConfig(); + } + + @Override + public MergeScheduler getMergeScheduler() { + if (mergeSchedulerWrapper == null) { + mergeSchedulerWrapper = new AbortOnCloseMergeScheduler(delegate.getMergeScheduler(), abortAction); + } + return mergeSchedulerWrapper; + } + + /** + * Marks all queued and running merges as aborted so that {@code merge.isAborted()} is + * {@code true} before the abort signal fires. This prevents the compound-file creation + * race in Lucene's {@code mergeMiddle()}. + */ + @Override + public void abortAllMerges() { + delegate.abortAllMerges(); + } + } + + private static final class AbortOnCloseMergeScheduler extends MergeScheduler { + + private final MergeScheduler delegate; + private final Runnable abortAction; + + private AbortOnCloseMergeScheduler(MergeScheduler delegate, Runnable abortAction) { + this.delegate = delegate; + this.abortAction = abortAction; + } + + @Override + public void merge(MergeSource mergeSource, MergeTrigger trigger) throws IOException { + delegate.merge(mergeSource, trigger); + } + + @Override + public Directory wrapForMerge(MergePolicy.OneMerge merge, Directory in) { + return delegate.wrapForMerge(merge, in); + } + + @Override + public Executor getIntraMergeExecutor(MergePolicy.OneMerge merge) { + return delegate.getIntraMergeExecutor(merge); + } + + @Override + public void close() throws IOException { + abortAction.run(); + delegate.close(); } } @@ -1149,6 +1259,12 @@ public void refreshConfig() { // no-op } + @Override + public void close() throws IOException { + signalAbortMergeReads(); + super.close(); + } + @Override protected void handleMergeException(Throwable t) { mergeException(t); diff --git a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInput.java b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInput.java index 9a2cda0878343..24fbe04b9bbc7 100644 --- a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInput.java +++ b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInput.java @@ -10,6 +10,7 @@ import org.apache.logging.log4j.Level; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.MergePolicy; import org.apache.lucene.store.DataAccessHint; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; @@ -76,6 +77,12 @@ public BlobCacheIndexInput( this(name, context, cacheFileReader, releasable, length, offset, null); } + private void checkMergeReadAborted() throws IOException { + if (cacheFileReader.isMergeReadAborted()) { + throw new MergePolicy.MergeAbortedException("shard is closing"); + } + } + @Override protected void seekInternal(long pos) throws IOException { BlobCacheUtils.ensureSeek(pos, this); @@ -171,6 +178,7 @@ String getSliceDescription() { @Override public boolean withMemorySegmentSlice(long offset, long length, CheckedConsumer action) throws IOException { + checkMergeReadAborted(); return cacheFileReader.withMemorySegmentSlice(this.offset + offset, Math.toIntExact(length), action); } @@ -182,6 +190,7 @@ public boolean withSliceAddresses( MemorySegment addressesScratch, CheckedConsumer action ) throws IOException { + checkMergeReadAborted(); if (DirectAccessInput.checkSlicesArgs(offsets, count, addressesScratch)) { return false; } @@ -202,6 +211,7 @@ public void prefetch(long offset, long length) throws IOException { @Override protected void readInternal(ByteBuffer b) throws IOException { + checkMergeReadAborted(); try { doReadInternal(b); } catch (Exception e) { @@ -265,6 +275,7 @@ private void doReadInternal(ByteBuffer b) throws Exception { } private void readInternalSlow(ByteBuffer b, long position, int length) throws Exception { + checkMergeReadAborted(); cacheFileReader.read( this, b, diff --git a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobStoreCacheDirectory.java b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobStoreCacheDirectory.java index eb8df313000a7..014a80eae467d 100644 --- a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobStoreCacheDirectory.java +++ b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/BlobStoreCacheDirectory.java @@ -40,10 +40,12 @@ import java.util.Collection; import java.util.Collections; import java.util.Map; +import java.util.Objects; import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.atomic.LongAdder; +import java.util.function.BooleanSupplier; import java.util.function.LongFunction; import java.util.function.LongSupplier; import java.util.stream.Stream; @@ -68,6 +70,7 @@ public abstract class BlobStoreCacheDirectory extends ByteSizeDirectory { private final AtomicReference updatingCommitThread = Assertions.ENABLED ? new AtomicReference<>() : null;// only used in asserts protected volatile Map currentMetadata = Map.of(); protected volatile long currentDataSetSizeInBytes = 0L; + private BooleanSupplier mergeReadAbortSupplier = () -> false; private final PluggableDirectoryMetricsHolder metricsHolder; BlobStoreCacheDirectory(StatelessSharedBlobCacheService cacheService, ShardId shardId) { @@ -101,6 +104,10 @@ public final void setBlobContainer(LongFunction blobContainer) { this.blobContainer.compareAndSet(null, blobContainer); // set once } + public void setMergeReadAbortSupplier(BooleanSupplier mergeReadAbortSupplier) { + this.mergeReadAbortSupplier = Objects.requireNonNull(mergeReadAbortSupplier); + } + public boolean containsFile(String name) { if (currentMetadata.isEmpty()) { try { @@ -323,8 +330,8 @@ protected final IndexInput doOpenInput( cacheService.getRegionSize(), context, cacheService.hasSearchRole() - ); - return new BlobCacheIndexInput(name, context, reader, releasable, blobFileRanges.fileLength(), blobFileRanges.fileOffset()); + ).withMergeReadAbortSupplier(mergeReadAbortSupplier); + return new BlobCacheIndexInput(name, context, reader, releasable, blobFileRanges.fileLength(), blobFileRanges.fileOffset(), null); } private SharedBlobCacheService.CacheFile getCacheFile(BlobFileRanges blobFileRanges) { diff --git a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectory.java b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectory.java index 7be15859ab717..8bc4ab55fde5f 100644 --- a/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectory.java +++ b/x-pack/plugin/stateless/src/main/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectory.java @@ -8,6 +8,7 @@ package org.elasticsearch.xpack.stateless.lucene; import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.MergePolicy; import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.store.ByteBuffersDirectory; @@ -108,6 +109,9 @@ public class IndexDirectory extends ByteSizeDirectory { */ private final AtomicLong estimatedSize = new AtomicLong(); + private final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + private volatile Runnable mergeReadCallback; + private final SetOnce recoveryCommitMetadataNodeEphemeralId = new SetOnce<>(); private final SetOnce recoveryCommitTranslogRecoveryStartFile = new SetOnce<>(); @@ -121,6 +125,7 @@ public IndexDirectory( ) { super(in); this.cacheDirectory = Objects.requireNonNull(cacheDirectory); + this.cacheDirectory.setMergeReadAbortSupplier(abortMergeReads::get); this.onGenerationalFileDeletion = onGenerationalFileDeletion; this.readSiFromMemoryIfPossible = readSiFromMemoryIfPossible; } @@ -483,20 +488,63 @@ public long getTranslogRecoveryStartFile() { return translogRecoveryStartFile != null ? translogRecoveryStartFile : 0; } + /** + * One-way latch set when the shard is closing. All reads on {@link BlobCacheIndexInput} + * wired through this directory throw {@link org.apache.lucene.index.MergePolicy.MergeAbortedException} + * once set, regardless of {@link IOContext}. Reads on {@link ReopeningIndexInput} are also + * aborted, but only when the read carries an explicit {@link IOContext.Context#MERGE} context + * (unuploaded segments always receive the real merge context via {@code openInput}). + *

+ * The flag is set during engine close, after all engine operations have been drained. + */ + public void abortMergeReads() { + abortMergeReads.set(true); + } + + private void checkMergeReadAborted(IOContext context) throws IOException { + if (context.context() == IOContext.Context.MERGE) { + Runnable callback = mergeReadCallback; + if (callback != null) { + callback.run(); + } + if (shouldAbortMergeReads()) { + throw new MergePolicy.MergeAbortedException("shard is closing"); + } + } + } + + /** Package-private for tests only — fires {@code callback} at the start of every merge read. */ + void setMergeReadCallback(Runnable callback) { + this.mergeReadCallback = callback; + } + + /** + * Package-private for tests and merge-scheduler close wiring only. + */ + boolean shouldAbortMergeReads() { + return abortMergeReads.get(); + } + public static IndexDirectory unwrapDirectory(final Directory directory) { + return tryUnwrapDirectory(directory).orElseThrow(() -> { + var e = new IllegalStateException(directory.getClass() + " cannot be unwrapped as " + IndexDirectory.class); + assert false : e; + return e; + }); + } + + public static Optional tryUnwrapDirectory(final Directory directory) { Directory dir = directory; while (dir != null) { if (dir instanceof IndexDirectory indexDirectory) { - return indexDirectory; + return Optional.of(indexDirectory); } else if (dir instanceof FilterDirectory) { dir = ((FilterDirectory) dir).getDelegate(); } else { dir = null; } } - var e = new IllegalStateException(directory.getClass() + " cannot be unwrapped as " + IndexDirectory.class); - assert false : e; - throw e; + return Optional.empty(); } class LocalFileRef extends AbstractRefCounted { @@ -1073,6 +1121,7 @@ public IndexInput slice(String sliceDescription, long sliceOffset, long sliceLen @Override protected void readInternal(ByteBuffer b) throws IOException { + IndexDirectory.this.checkMergeReadAborted(context); executeLocallyOrReopen(current -> { assert assertPositionMatchesFilePointer(current); var len = b.remaining(); diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/AbstractEngineTestCase.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/AbstractEngineTestCase.java index ecf6366336c87..1e7925a4b4740 100644 --- a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/AbstractEngineTestCase.java +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/AbstractEngineTestCase.java @@ -12,8 +12,10 @@ import org.apache.lucene.document.StoredField; import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexCommit; +import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.MergePolicy; import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.util.SetOnce; import org.elasticsearch.action.ActionListener; import org.elasticsearch.blobcache.BlobCacheMetrics; @@ -31,6 +33,7 @@ import org.elasticsearch.common.lucene.Lucene; import org.elasticsearch.common.settings.ClusterSettings; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.concurrent.ConcurrentCollections; import org.elasticsearch.common.util.concurrent.DeterministicTaskQueue; @@ -39,9 +42,12 @@ import org.elasticsearch.core.IOUtils; import org.elasticsearch.core.PathUtils; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.env.Environment; import org.elasticsearch.env.NodeEnvironment; +import org.elasticsearch.env.TestEnvironment; import org.elasticsearch.index.Index; import org.elasticsearch.index.IndexModule; +import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.codec.CodecService; import org.elasticsearch.index.engine.Engine; import org.elasticsearch.index.engine.EngineConfig; @@ -100,6 +106,8 @@ import org.elasticsearch.xpack.stateless.engine.translog.TranslogRecoveryMetrics; import org.elasticsearch.xpack.stateless.engine.translog.TranslogReplicator; import org.elasticsearch.xpack.stateless.lucene.BlobStoreCacheDirectoryMetrics; +import org.elasticsearch.xpack.stateless.lucene.IndexBlobStoreCacheDirectory; +import org.elasticsearch.xpack.stateless.lucene.IndexDirectory; import org.elasticsearch.xpack.stateless.lucene.SearchDirectory; import org.elasticsearch.xpack.stateless.lucene.StatelessCommitRef; import org.elasticsearch.xpack.stateless.objectstore.ObjectStoreService; @@ -128,10 +136,12 @@ import java.util.function.LongSupplier; import static org.elasticsearch.blobcache.shared.SharedBlobCacheService.SHARED_CACHE_REGION_SIZE_SETTING; +import static org.elasticsearch.blobcache.shared.SharedBlobCacheService.SHARED_CACHE_SIZE_SETTING; import static org.elasticsearch.common.util.concurrent.EsExecutors.DIRECT_EXECUTOR_SERVICE; import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; import static org.elasticsearch.xpack.stateless.StatelessPlugin.SHARD_READ_THREAD_POOL; import static org.elasticsearch.xpack.stateless.StatelessPlugin.SHARD_READ_THREAD_POOL_SETTING; +import static org.elasticsearch.xpack.stateless.TestUtils.newCacheService; import static org.hamcrest.Matchers.instanceOf; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; @@ -149,6 +159,8 @@ public abstract class AbstractEngineTestCase extends ESTestCase { private BlobContainer blobContainer; protected StatelessSharedBlobCacheService sharedBlobCacheService; protected NodeEnvironment nodeEnvironment; + private final List indexDirectoryNodeEnvironments = new ArrayList<>(); + private final List indexDirectoryCacheServices = new ArrayList<>(); @SuppressWarnings("unchecked") @Override @@ -180,6 +192,12 @@ public void tearDown() throws Exception { assert threadPools.isEmpty() : threadPools; IOUtils.rm(blobStorePath); nodeEnvironment.close(); + for (NodeEnvironment extraNodeEnvironment : indexDirectoryNodeEnvironments) { + extraNodeEnvironment.close(); + } + indexDirectoryNodeEnvironments.clear(); + IOUtils.close(indexDirectoryCacheServices); + indexDirectoryCacheServices.clear(); for (CircuitBreaker breaker : readerHeapBreakers) { assertEquals( "reader-heap breaker leaked bytes after test — every reservation must be released by reader close", @@ -378,6 +396,34 @@ protected EngineConfig indexConfig( ); store.associateIndexWithNewTranslog(translogUuid); + return buildEngineConfig( + shardId, + indexSettings, + translogConfig, + threadPool, + threadPoolMergeExecutorService, + store, + mergePolicy, + indexWriterConfig, + mapperService, + primaryTermSupplier, + new CapturingEngineEventListener() + ); + } + + private EngineConfig buildEngineConfig( + ShardId shardId, + IndexSettings indexSettings, + TranslogConfig translogConfig, + ThreadPool threadPool, + ThreadPoolMergeExecutorService threadPoolMergeExecutorService, + Store store, + MergePolicy mergePolicy, + IndexWriterConfig indexWriterConfig, + MapperService mapperService, + LongSupplier primaryTermSupplier, + Engine.EventListener eventListener + ) { return EngineConfig.builder() .shardId(shardId) .threadPool(threadPool) @@ -388,7 +434,7 @@ protected EngineConfig indexConfig( .analyzer(indexWriterConfig.getAnalyzer()) .similarity(indexWriterConfig.getSimilarity()) .codecProvider(new CodecService(null, BigArrays.NON_RECYCLING_INSTANCE, threadPool)) - .eventListener(new CapturingEngineEventListener()) + .eventListener(eventListener) .queryCache(IndexSearcher.getDefaultQueryCache()) .queryCachingPolicy(IndexSearcher.getDefaultQueryCachingPolicy()) .translogConfig(translogConfig) @@ -408,6 +454,98 @@ protected EngineConfig indexConfig( .build(); } + protected EngineConfig indexConfigWithIndexDirectory(Settings settings, Settings nodeSettings) throws IOException { + var primaryTerm = new AtomicLong(1L); + return indexConfigWithIndexDirectory(settings, nodeSettings, primaryTerm::get); + } + + protected EngineConfig indexConfigWithIndexDirectory(Settings settings, Settings nodeSettings, LongSupplier primaryTermSupplier) + throws IOException { + // Mock is sufficient: tests only index Lucene fields directly and do not require bbq_hnsw mapping or codec support. + MapperService mapperService = Mockito.mock(MapperService.class); + when(mapperService.mappingLookup()).thenReturn(MappingLookup.EMPTY); + return indexConfigWithIndexDirectory(settings, nodeSettings, primaryTermSupplier, newMergePolicy(), mapperService); + } + + protected EngineConfig indexConfigWithIndexDirectory( + Settings settings, + Settings nodeSettings, + LongSupplier primaryTermSupplier, + MergePolicy mergePolicy, + MapperService mapperService + ) throws IOException { + return indexConfigWithIndexDirectory(settings, nodeSettings, primaryTermSupplier, mergePolicy, mapperService, null); + } + + protected EngineConfig indexConfigWithIndexDirectory( + Settings settings, + Settings nodeSettings, + LongSupplier primaryTermSupplier, + MergePolicy mergePolicy, + MapperService mapperService, + CapturingEngineEventListener eventListener + ) throws IOException { + Settings.Builder mergedNodeSettingsBuilder = Settings.builder().put(nodeSettings); + if (nodeSettings.hasValue(ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.getKey()) == false) { + mergedNodeSettingsBuilder.put(ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.getKey(), true); + } + Settings mergedNodeSettings = mergedNodeSettingsBuilder.build(); + + var shardId = new ShardId(new Index(randomAlphaOfLengthBetween(5, 10), UUIDs.randomBase64UUID(random())), randomInt(10)); + var indexSettings = IndexSettingsModule.newIndexSettings(shardId.getIndex(), settings, mergedNodeSettings); + var translogConfig = new TranslogConfig(shardId, createTempDir(), indexSettings, BigArrays.NON_RECYCLING_INSTANCE); + var indexWriterConfig = newIndexWriterConfig(); + var threadPool = registerThreadPool( + new TestThreadPool( + getTestName() + "[" + shardId + "][index-directory]", + StatelessPlugin.statelessExecutorBuilders(Settings.EMPTY, true) + ) + ); + var threadPoolMergeExecutorService = ThreadPoolMergeExecutorService.maybeCreateThreadPoolMergeExecutorService( + threadPool, + ClusterSettings.createBuiltInClusterSettings(mergedNodeSettings), + nodeEnvironment + ); + Settings envSettings = Settings.builder() + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir().toAbsolutePath()) + .putList(Environment.PATH_DATA_SETTING.getKey(), createTempDir().toAbsolutePath().toString()) + .put(SHARED_CACHE_SIZE_SETTING.getKey(), ByteSizeValue.ofMb(10)) + .put(SHARED_CACHE_REGION_SIZE_SETTING.getKey(), ByteSizeValue.ofKb(256)) + .build(); + NodeEnvironment cacheNodeEnvironment = new NodeEnvironment(envSettings, TestEnvironment.newEnvironment(envSettings)); + indexDirectoryNodeEnvironments.add(cacheNodeEnvironment); + StatelessSharedBlobCacheService cacheService = newCacheService(cacheNodeEnvironment, envSettings, threadPool); + indexDirectoryCacheServices.add(cacheService); + IndexBlobStoreCacheDirectory cacheDirectory = new IndexBlobStoreCacheDirectory(cacheService, shardId); + cacheDirectory.setBlobContainer(term -> blobContainer); + + var underlying = new ByteBuffersDirectory(); + IndexDirectory indexDirectory = new IndexDirectory(underlying, cacheDirectory, null, true); + var store = new Store(shardId, indexSettings, indexDirectory, new DummyShardLock(shardId)); + store.createEmpty(); + final String translogUuid = Translog.createEmptyTranslog( + translogConfig.getTranslogPath(), + SequenceNumbers.NO_OPS_PERFORMED, + shardId, + primaryTermSupplier.getAsLong() + ); + store.associateIndexWithNewTranslog(translogUuid); + + return buildEngineConfig( + shardId, + indexSettings, + translogConfig, + threadPool, + threadPoolMergeExecutorService, + store, + mergePolicy, + indexWriterConfig, + mapperService, + primaryTermSupplier, + eventListener != null ? eventListener : new CapturingEngineEventListener() + ); + } + protected SearchEngine newSearchEngine() { return newSearchEngineFromIndexEngine(searchConfig()); } diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/IndexEngineTests.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/IndexEngineTests.java index df7a91c02ee80..26c58da384781 100644 --- a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/IndexEngineTests.java +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/engine/IndexEngineTests.java @@ -12,6 +12,7 @@ import org.apache.lucene.index.SegmentInfos; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.support.PlainActionFuture; +import org.elasticsearch.cluster.ClusterModule; import org.elasticsearch.cluster.metadata.IndexMetadata; import org.elasticsearch.common.CheckedBiFunction; import org.elasticsearch.common.UUIDs; @@ -22,6 +23,8 @@ import org.elasticsearch.escf.EscfBatch; import org.elasticsearch.escf.EscfEncoder; import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.IndexVersion; +import org.elasticsearch.index.MapperTestUtils; import org.elasticsearch.index.VersionType; import org.elasticsearch.index.engine.Engine; import org.elasticsearch.index.engine.EngineBatch; @@ -30,8 +33,12 @@ import org.elasticsearch.index.engine.MergeMemoryEstimator; import org.elasticsearch.index.engine.MergeMetrics; import org.elasticsearch.index.engine.ThreadPoolMergeScheduler; +import org.elasticsearch.index.mapper.DocumentMapper; import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.MappingLookup; +import org.elasticsearch.index.mapper.ParsedDocument; +import org.elasticsearch.index.mapper.SourceToParse; +import org.elasticsearch.index.mapper.Uid; import org.elasticsearch.index.merge.OnGoingMerge; import org.elasticsearch.index.shard.ShardSplittingQuery; import org.elasticsearch.index.translog.Translog; @@ -41,6 +48,9 @@ import org.elasticsearch.sourcebatch.SourceBatch; import org.elasticsearch.telemetry.TelemetryProvider; import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.xcontent.NamedXContentRegistry; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentFactory; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.stateless.StatelessPlugin; import org.elasticsearch.xpack.stateless.cache.SharedBlobCacheWarmingService; @@ -62,11 +72,13 @@ import java.util.Collections; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.NavigableMap; import java.util.Set; import java.util.TreeMap; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.stream.LongStream; @@ -80,8 +92,11 @@ import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.lessThan; import static org.hamcrest.Matchers.lessThanOrEqualTo; import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.nullValue; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.any; import static org.mockito.Mockito.doAnswer; @@ -870,6 +885,135 @@ public void moveQueuedMergeBytesToRunning(OnGoingMerge currentMerge, long estima } } + public void testCloseDuringVectorMergeCompletesQuickly() throws Exception { + assertCloseDuringForceMergeCompletesQuickly(VectorDocType.FLOAT, true, true); + } + + public void testCloseDuringBbqHnswMergeCompletesQuickly() throws Exception { + assertCloseDuringForceMergeCompletesQuickly(VectorDocType.BBQ_HNSW, true, true); + } + + private enum VectorDocType { + FLOAT, + BBQ_HNSW + } + + private void assertCloseDuringForceMergeCompletesQuickly( + VectorDocType vectorDocType, + boolean periodicFlush, + boolean finalFlushBeforeMerge + ) throws Exception { + Settings nodeSettings = Settings.builder().put(StatelessPlugin.STATELESS_ENABLED.getKey(), true).build(); + CapturingEngineEventListener eventListener = new CapturingEngineEventListener(); + EngineConfig config; + if (vectorDocType == VectorDocType.BBQ_HNSW) { + int numDimensions = randomIntBetween(64, 128); + MapperService mapperService = createBbqHnswMapperService(numDimensions); + config = indexConfigWithIndexDirectory(Settings.EMPTY, nodeSettings, () -> 1L, newMergePolicy(), mapperService, eventListener); + assertCloseDuringForceMerge(config, eventListener, numDimensions, vectorDocType, periodicFlush, finalFlushBeforeMerge); + } else { + // Mock is sufficient: float vector tests index KnnFloatVectorField on the Lucene document directly + // without requiring bbq_hnsw mapping or codec support. + config = indexConfigWithIndexDirectory( + Settings.EMPTY, + nodeSettings, + () -> 1L, + newMergePolicy(), + mock(MapperService.class), + eventListener + ); + when(config.getMapperService().mappingLookup()).thenReturn(MappingLookup.EMPTY); + int numDimensions = randomIntBetween(16, 64); + assertCloseDuringForceMerge(config, eventListener, numDimensions, vectorDocType, periodicFlush, finalFlushBeforeMerge); + } + } + + private void assertCloseDuringForceMerge( + EngineConfig config, + CapturingEngineEventListener eventListener, + int numDimensions, + VectorDocType vectorDocType, + boolean periodicFlush, + boolean finalFlushBeforeMerge + ) throws Exception { + final Thread[] mergeThread = new Thread[1]; + long startNanos; + try (IndexEngine engine = newIndexEngine(config)) { + int numDocs = finalFlushBeforeMerge ? randomIntBetween(40, 100) : randomIntBetween(80, 150); + for (int i = 0; i < numDocs; i++) { + indexVectorDoc(engine, config, vectorDocType, numDimensions, String.valueOf(i)); + if (periodicFlush && i % 10 == 0) { + engine.flush(); + } + } + if (finalFlushBeforeMerge) { + engine.flush(); + } + + mergeThread[0] = new Thread(() -> { + try { + engine.forceMerge(false, 1, false, UUIDs.randomBase64UUID()); + } catch (Throwable e) { + // merge may be aborted during close + } + }); + mergeThread[0].start(); + assertBusy(() -> assertTrue("merge was not queued or running before engine close", engine.hasQueuedOrRunningMerges())); + startNanos = System.nanoTime(); + } + mergeThread[0].join(TimeUnit.SECONDS.toMillis(5)); + assertFalse("merge thread should finish after engine close", mergeThread[0].isAlive()); + assertThat(System.nanoTime() - startNanos, lessThan(TimeValue.timeValueSeconds(5).nanos())); + assertThat(eventListener.reason.get(), is(nullValue())); + assertThat(eventListener.exception.get(), is(nullValue())); + } + + private void indexVectorDoc(IndexEngine engine, EngineConfig config, VectorDocType vectorDocType, int numDimensions, String id) + throws IOException { + if (vectorDocType == VectorDocType.BBQ_HNSW) { + engine.index(bbqHnswDoc(config.getMapperService(), id, randomVector(numDimensions))); + } else { + float[] vector = randomVector(numDimensions); + engine.index(randomDoc(id, (builder, doc) -> doc.add(new KnnFloatVectorField("vector", vector)))); + } + } + + private static Engine.Index bbqHnswDoc(MapperService mapperService, String id, float[] vector) throws IOException { + try (XContentBuilder builder = XContentFactory.jsonBuilder().startObject().array("vector", vector).endObject()) { + DocumentMapper documentMapper = mapperService.documentMapper(); + ParsedDocument parsed = documentMapper.parse(new SourceToParse(id, BytesReference.bytes(builder), XContentType.JSON)); + return new Engine.Index(Uid.encodeId(id), 1L, parsed); + } + } + + private MapperService createBbqHnswMapperService(int dims) throws IOException { + String mapping = String.format(Locale.ROOT, """ + { + "properties": { + "vector": { + "type": "dense_vector", + "dims": %d, + "index": true, + "similarity": "cosine", + "index_options": { "type": "bbq_hnsw" } + } + } + } + """, dims); + IndexMetadata indexMetadata = IndexMetadata.builder("index") + .settings(indexSettings(1, 1).put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current())) + .putMapping(mapping) + .build(); + MapperService mapperService = MapperTestUtils.newMapperService( + new NamedXContentRegistry(ClusterModule.getNamedXWriteables()), + createTempDir(), + indexMetadata.getSettings(), + "index" + ); + mapperService.merge(indexMetadata, MapperService.MergeReason.MAPPING_UPDATE); + return mapperService; + } + public void testDeleteUnownedDocumentsBumpsForceMergeUUID() throws Exception { Settings settings = Settings.builder() .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 2) diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInputTests.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInputTests.java index 0d66ba3ed862f..c508914187228 100644 --- a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInputTests.java +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/BlobCacheIndexInputTests.java @@ -7,8 +7,11 @@ package org.elasticsearch.xpack.stateless.lucene; +import org.apache.lucene.index.MergePolicy; import org.apache.lucene.store.AlreadyClosedException; +import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.MergeInfo; import org.apache.lucene.store.RandomAccessInput; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.ResourceAlreadyUploadedException; @@ -1693,6 +1696,340 @@ public void getRangeInputStream(long position, int length, ActionListener { + try ( + BlobCacheIndexInput abortedInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.merge(new MergeInfo(100, 1024L, false, -1)) + ) + ) { + abortedInput.readByte(); + } + }); + } + } + + public void testAbortMergeReadsThrowsOnDefaultContextRead() throws IOException { + final AtomicBoolean abortMergeReads = new AtomicBoolean(true); + final byte[] input = randomByteArrayOfLength(randomIntBetween(64, 512)); + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(64)); + final var settings = sharedCacheSettings(ByteSizeValue.ofBytes(regionSize.getBytes() * 10), regionSize); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.DEFAULT + ) + ) { + byte[] actual = new byte[16]; + expectThrows(MergePolicy.MergeAbortedException.class, () -> indexInput.readBytes(actual, 0, actual.length)); + } + } + + public void testAbortMergeReadsBeforeWithMemorySegmentSlice() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(64, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput(sharedBlobCacheService, input, abortMergeReads) + ) { + populateCache(indexInput, input); + boolean available = indexInput.withMemorySegmentSlice(0, input.length, seg -> { assertEquals(input.length, seg.byteSize()); }); + assertTrue(available); + } + } + + public void testAbortMergeReadsThrowsOnWithMemorySegmentSlice() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(64, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.merge(new MergeInfo(100, 1024L, false, -1)) + ) + ) { + populateCache(indexInput, input); + abortMergeReads.set(true); + expectThrows(MergePolicy.MergeAbortedException.class, () -> indexInput.withMemorySegmentSlice(0, input.length, seg -> { + throw new AssertionError("should not run"); + })); + } + } + + // Verifies the production HNSW path: abort fires on DEFAULT-context reads (the IOContext guard was a bug). + public void testAbortMergeReadsThrowsOnDefaultContextWithMemorySegmentSlice() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(64, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.DEFAULT + ) + ) { + populateCache(indexInput, input); + abortMergeReads.set(true); + expectThrows(MergePolicy.MergeAbortedException.class, () -> indexInput.withMemorySegmentSlice(0, input.length, seg -> { + throw new AssertionError("should not run"); + })); + } + } + + public void testAbortMergeReadsBeforeWithMemorySegmentSlices() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(200, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput(sharedBlobCacheService, input, abortMergeReads) + ) { + populateCache(indexInput, input); + int sliceLen = randomIntBetween(1, input.length / 4); + long[] offsets = new long[] { 0, randomIntBetween(1, input.length / 2 - sliceLen), input.length - sliceLen }; + MemorySegment addrsOut = MemorySegment.ofArray(new long[3]); + boolean available = indexInput.withSliceAddresses(offsets, sliceLen, 3, addrsOut, addrs -> {}); + assertTrue(available); + } + } + + public void testAbortMergeReadsThrowsOnWithMemorySegmentSlices() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(200, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.merge(new MergeInfo(100, 1024L, false, -1)) + ) + ) { + populateCache(indexInput, input); + int sliceLen = randomIntBetween(1, input.length / 4); + long[] offsets = new long[] { 0, randomIntBetween(1, input.length / 2 - sliceLen), input.length - sliceLen }; + abortMergeReads.set(true); + MemorySegment addrsOut = MemorySegment.ofArray(new long[3]); + expectThrows( + MergePolicy.MergeAbortedException.class, + () -> indexInput.withSliceAddresses(offsets, sliceLen, 3, addrsOut, addrs -> { + throw new AssertionError("should not run"); + }) + ); + } + } + + // Verifies the production HNSW path: abort fires on DEFAULT-context reads (the IOContext guard was a bug). + public void testAbortMergeReadsThrowsOnDefaultContextWithSliceAddresses() throws IOException { + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(randomIntBetween(4, 64))); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(regionSize.getBytes() * 10); + final var settings = Settings.builder() + .put(sharedCacheSettings(cacheSize, regionSize)) + .put(SharedBlobCacheService.SHARED_CACHE_MMAP.getKey(), true) + .build(); + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(200, (int) regionSize.getBytes() / 2)); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + input, + abortMergeReads, + IOContext.DEFAULT + ) + ) { + populateCache(indexInput, input); + int sliceLen = randomIntBetween(1, input.length / 4); + long[] offsets = new long[] { 0, randomIntBetween(1, input.length / 2 - sliceLen), input.length - sliceLen }; + abortMergeReads.set(true); + MemorySegment addrsOut = MemorySegment.ofArray(new long[3]); + expectThrows( + MergePolicy.MergeAbortedException.class, + () -> indexInput.withSliceAddresses(offsets, sliceLen, 3, addrsOut, addrs -> { + throw new AssertionError("should not run"); + }) + ); + } + } + + /** + * Regression test for the bug where {@code doSlice(sliceDescription, offset, length, IOContext)} — the overload used + * by Lucene's compound-file reader to open sub-files with a per-sub-file context — creates a new + * {@code BlobCacheIndexInput} without forwarding the {@code mergeReadAbortSupplier}. The slice therefore never + * checks for abort, even when the parent had an active abort supplier. Since compound-file reads account for the + * majority of HNSW merge I/O, this gap means merge reads in {@code .cfs} blobs are not aborted on shard close. + */ + public void testAbortMergeReadsThrowsOnIoContextSlice() throws IOException { + final AtomicBoolean abortMergeReads = new AtomicBoolean(false); + final byte[] input = randomByteArrayOfLength(randomIntBetween(64, 512)); + final ByteSizeValue regionSize = pageAligned(ByteSizeValue.ofKb(64)); + final var settings = sharedCacheSettings(ByteSizeValue.ofBytes(regionSize.getBytes() * 10), regionSize); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + BlobCacheIndexInput indexInput = newAbortableBlobCacheIndexInput(sharedBlobCacheService, input, abortMergeReads) + ) { + // Populate the cache before enabling abort so the read itself doesn't fail for other reasons. + populateCache(indexInput, input); + abortMergeReads.set(true); + + // doSlice with an explicit IOContext is the path Lucene's compound-file reader uses to open sub-files. + // It should propagate the parent's abort supplier to the resulting slice. + long sliceLength = input.length / 2; + BlobCacheIndexInput slice = asInstanceOf( + BlobCacheIndexInput.class, + indexInput.doSlice("sub-file", 0, sliceLength, IOContext.merge(new MergeInfo(100, 1024L, false, -1))) + ); + expectThrows(MergePolicy.MergeAbortedException.class, slice::readByte); + } + } + + private BlobCacheIndexInput newAbortableBlobCacheIndexInput( + StatelessSharedBlobCacheService sharedBlobCacheService, + byte[] input, + AtomicBoolean abortMergeReads + ) throws IOException { + return newAbortableBlobCacheIndexInput(sharedBlobCacheService, input, abortMergeReads, randomIOContext()); + } + + private BlobCacheIndexInput newAbortableBlobCacheIndexInput( + StatelessSharedBlobCacheService sharedBlobCacheService, + byte[] input, + AtomicBoolean abortMergeReads, + IOContext context + ) throws IOException { + final ShardId shardId = new ShardId(new Index("_index_name", "_index_id"), 0); + final String fileName = randomAlphaOfLength(5) + randomFileExtension(); + final long primaryTerm = randomNonNegativeLong(); + return newAbortableBlobCacheIndexInput(sharedBlobCacheService, shardId, fileName, primaryTerm, input, abortMergeReads, context); + } + + private BlobCacheIndexInput newAbortableBlobCacheIndexInput( + StatelessSharedBlobCacheService sharedBlobCacheService, + ShardId shardId, + String fileName, + long primaryTerm, + byte[] input, + AtomicBoolean abortMergeReads + ) throws IOException { + return newAbortableBlobCacheIndexInput( + sharedBlobCacheService, + shardId, + fileName, + primaryTerm, + input, + abortMergeReads, + randomIOContext() + ); + } + + private BlobCacheIndexInput newAbortableBlobCacheIndexInput( + StatelessSharedBlobCacheService sharedBlobCacheService, + ShardId shardId, + String fileName, + long primaryTerm, + byte[] input, + AtomicBoolean abortMergeReads, + IOContext context + ) throws IOException { + return new BlobCacheIndexInput( + fileName, + context, + new CacheFileReader( + sharedBlobCacheService.getCacheFile( + new FileCacheKey(shardId, primaryTerm, fileName), + input.length, + SharedBlobCacheService.CacheMissHandler.NOOP + ), + createBlobReader(fileName, input, sharedBlobCacheService), + createBlobFileRanges(primaryTerm, 0L, 0, input.length), + BlobCacheMetrics.NOOP, + System::currentTimeMillis + ).withMergeReadAbortSupplier(abortMergeReads::get), + null, + input.length, + 0, + null + ); + } + + private void populateCache(BlobCacheIndexInput indexInput, byte[] input) throws IOException { + byte[] output = randomReadAndSlice(indexInput, input.length); + assertArrayEquals(input, output); + } + private static Settings sharedCacheSettings(ByteSizeValue cacheSize) { return sharedCacheSettings(cacheSize, pageAligned(ByteSizeValue.of(randomIntBetween(4, 1024), ByteSizeUnit.KB))); } diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectoryTests.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectoryTests.java index 8d9455e351199..ef7c4a952c8de 100644 --- a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectoryTests.java +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexDirectoryTests.java @@ -11,6 +11,7 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.MergePolicy; import org.apache.lucene.index.SegmentCommitInfo; import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.store.Directory; @@ -18,14 +19,17 @@ import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.store.MergeInfo; import org.apache.lucene.tests.mockfile.FilterFileChannel; import org.apache.lucene.tests.mockfile.FilterFileSystemProvider; import org.elasticsearch.blobcache.common.BlobCacheBufferedIndexInput; import org.elasticsearch.blobcache.shared.SharedBlobCacheService; import org.elasticsearch.common.blobstore.BlobContainer; import org.elasticsearch.common.blobstore.BlobPath; +import org.elasticsearch.common.blobstore.OperationPurpose; import org.elasticsearch.common.blobstore.fs.FsBlobContainer; import org.elasticsearch.common.blobstore.fs.FsBlobStore; +import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.util.Maps; @@ -54,6 +58,7 @@ import java.io.FileNotFoundException; import java.io.IOException; +import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.FileSystem; import java.nio.file.NoSuchFileException; @@ -77,6 +82,7 @@ import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.is; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -457,6 +463,123 @@ public void testUpdateRecoveryCommit() throws IOException { } } + public void testAbortMergeReadsOnCacheBackedOpenInput() throws Exception { + final Path dataPath = createTempDir(); + final ShardId shardId = new ShardId(new Index("_index_name", "_index_id"), 0); + final ThreadPool threadPool = getThreadPool("testAbortMergeReadsOnCacheBackedOpenInput"); + final var settings = Settings.builder() + .put(SharedBlobCacheService.SHARED_CACHE_SIZE_SETTING.getKey(), ByteSizeValue.ofKb(4L)) + .put(SharedBlobCacheService.SHARED_CACHE_REGION_SIZE_SETTING.getKey(), ByteSizeValue.ofKb(4L)) + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir()) + .putList(Environment.PATH_DATA_SETTING.getKey(), dataPath.toAbsolutePath().toString()) + .build(); + final Path indexDataPath = dataPath.resolve("index"); + final Path blobStorePath = PathUtils.get(createTempDir().toString()); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + FsBlobStore blobStore = new FsBlobStore(randomIntBetween(1, 8) * 1024, blobStorePath, false); + IndexDirectory directory = new IndexDirectory( + newFSDirectory(indexDataPath), + new IndexBlobStoreCacheDirectory(sharedBlobCacheService, shardId), + null, + true + ) + ) { + final FsBlobContainer blobContainer = new FsBlobContainer(blobStore, BlobPath.EMPTY, blobStorePath); + directory.getBlobStoreCacheDirectory().setBlobContainer(value -> blobContainer); + + final int fileLength = randomIntBetween(1024, 10240); + final byte[] bytes = randomByteArrayOfLength(fileLength); + final String fileName = "file." + randomFrom(LuceneFilesExtensions.values()).getExtension(); + final long primaryTerm = 1L; + try (IndexOutput output = directory.createOutput(fileName, IOContext.DEFAULT)) { + output.writeBytes(bytes, fileLength); + } + blobContainer.writeBlob( + OperationPurpose.INDICES, + StatelessCompoundCommit.PREFIX + primaryTerm, + BytesReference.fromByteBuffer(ByteBuffer.wrap(bytes)), + true + ); + directory.updateCommit( + primaryTerm, + fileLength, + Set.of(fileName), + Map.of(fileName, createBlobFileRanges(primaryTerm, primaryTerm, 0L, fileLength)) + ); + + try (IndexInput input = directory.openInput(fileName, IOContext.DEFAULT)) { + assertThat(input, instanceOf(BlobCacheIndexInput.class)); + input.readBytes(new byte[16], 0, 16); + } + directory.abortMergeReads(); + expectThrows(MergePolicy.MergeAbortedException.class, () -> { + try (IndexInput input = directory.openInput(fileName, IOContext.merge(new MergeInfo(100, 1024L, false, -1)))) { + input.readByte(); + } + }); + } finally { + assertTrue(ThreadPool.terminate(threadPool, 10L, TimeUnit.SECONDS)); + } + } + + public void testAbortMergeReadsThrowsOnDefaultContextRead() throws Exception { + final Path dataPath = createTempDir(); + final ShardId shardId = new ShardId(new Index("_index_name", "_index_id"), 0); + final ThreadPool threadPool = getThreadPool("testAbortMergeReadsThrowsOnDefaultContextRead"); + final var settings = Settings.builder() + .put(SharedBlobCacheService.SHARED_CACHE_SIZE_SETTING.getKey(), ByteSizeValue.ofKb(4L)) + .put(SharedBlobCacheService.SHARED_CACHE_REGION_SIZE_SETTING.getKey(), ByteSizeValue.ofKb(4L)) + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir()) + .putList(Environment.PATH_DATA_SETTING.getKey(), dataPath.toAbsolutePath().toString()) + .build(); + final Path indexDataPath = dataPath.resolve("index"); + final Path blobStorePath = PathUtils.get(createTempDir().toString()); + try ( + NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + FsBlobStore blobStore = new FsBlobStore(randomIntBetween(1, 8) * 1024, blobStorePath, false); + IndexDirectory directory = new IndexDirectory( + newFSDirectory(indexDataPath), + new IndexBlobStoreCacheDirectory(sharedBlobCacheService, shardId), + null, + true + ) + ) { + final FsBlobContainer blobContainer = new FsBlobContainer(blobStore, BlobPath.EMPTY, blobStorePath); + directory.getBlobStoreCacheDirectory().setBlobContainer(value -> blobContainer); + + final int fileLength = randomIntBetween(1024, 10240); + final byte[] bytes = randomByteArrayOfLength(fileLength); + final String fileName = "file." + randomFrom(LuceneFilesExtensions.values()).getExtension(); + final long primaryTerm = 1L; + try (IndexOutput output = directory.createOutput(fileName, IOContext.DEFAULT)) { + output.writeBytes(bytes, fileLength); + } + blobContainer.writeBlob( + OperationPurpose.INDICES, + StatelessCompoundCommit.PREFIX + primaryTerm, + BytesReference.fromByteBuffer(ByteBuffer.wrap(bytes)), + true + ); + directory.updateCommit( + primaryTerm, + fileLength, + Set.of(fileName), + Map.of(fileName, createBlobFileRanges(primaryTerm, primaryTerm, 0L, fileLength)) + ); + + directory.abortMergeReads(); + try (IndexInput input = directory.openInput(fileName, IOContext.DEFAULT)) { + assertThat(input, instanceOf(BlobCacheIndexInput.class)); + expectThrows(MergePolicy.MergeAbortedException.class, () -> input.readBytes(new byte[16], 0, 16)); + } + } finally { + assertTrue(ThreadPool.terminate(threadPool, 10L, TimeUnit.SECONDS)); + } + } + public void testReadFromMemory() throws Exception { final Path dataPath = createTempDir(); final Path indexDataPath = dataPath.resolve("index"); diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexEngineMergeAbortTests.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexEngineMergeAbortTests.java new file mode 100644 index 0000000000000..390584137a586 --- /dev/null +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/IndexEngineMergeAbortTests.java @@ -0,0 +1,56 @@ +/* + * 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.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.engine.EngineConfig; +import org.elasticsearch.index.engine.ThreadPoolMergeScheduler; +import org.elasticsearch.xpack.stateless.StatelessPlugin; +import org.elasticsearch.xpack.stateless.engine.AbstractEngineTestCase; +import org.elasticsearch.xpack.stateless.engine.IndexEngine; +import org.junit.After; + +/** + * Verifies merge-read abort wiring on {@link IndexDirectory}, including the path where + * {@link IndexEngine}'s merge scheduler close signals abort before the directory is closed. + */ +public class IndexEngineMergeAbortTests extends AbstractEngineTestCase { + + @After + public void assertWarningHeaders() { + assertWarnings( + "[indices.merge.scheduler.use_thread_pool] setting was deprecated in Elasticsearch and will be removed in a future release. " + + "See the breaking changes documentation for the next major version." + ); + } + + public void testMergeSchedulerCloseSignalsAbortMergeReads() throws Exception { + Settings nodeSettings = Settings.builder().put(StatelessPlugin.STATELESS_ENABLED.getKey(), true).build(); + EngineConfig config = indexConfigWithIndexDirectory(Settings.EMPTY, nodeSettings); + IndexDirectory indexDirectory = IndexDirectory.unwrapDirectory(config.getStore().directory()); + try (IndexEngine engine = newIndexEngine(config)) { + engine.index(randomDoc(String.valueOf(0))); + engine.flush(); + } + assertTrue(indexDirectory.shouldAbortMergeReads()); + } + + public void testConcurrentMergeSchedulerCloseSignalsAbortMergeReads() throws Exception { + Settings nodeSettings = Settings.builder() + .put(StatelessPlugin.STATELESS_ENABLED.getKey(), true) + .put(ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.getKey(), false) + .build(); + EngineConfig config = indexConfigWithIndexDirectory(Settings.EMPTY, nodeSettings); + IndexDirectory indexDirectory = IndexDirectory.unwrapDirectory(config.getStore().directory()); + try (IndexEngine engine = newIndexEngine(config)) { + engine.index(randomDoc(String.valueOf(0))); + engine.flush(); + } + assertTrue(indexDirectory.shouldAbortMergeReads()); + } +} diff --git a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/ReopeningIndexInputTests.java b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/ReopeningIndexInputTests.java index 587d640755483..7fe6e19e5d80b 100644 --- a/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/ReopeningIndexInputTests.java +++ b/x-pack/plugin/stateless/src/test/java/org/elasticsearch/xpack/stateless/lucene/ReopeningIndexInputTests.java @@ -7,10 +7,12 @@ package org.elasticsearch.xpack.stateless.lucene; +import org.apache.lucene.index.MergePolicy; import org.apache.lucene.store.FilterIndexInput; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.store.MergeInfo; import org.apache.lucene.tests.mockfile.FilterFileSystemProvider; import org.apache.lucene.tests.mockfile.HandleTrackingFS; import org.elasticsearch.action.ActionListener; @@ -44,6 +46,7 @@ import org.elasticsearch.xpack.stateless.cache.StatelessSharedBlobCacheService; import org.elasticsearch.xpack.stateless.test.FakeStatelessNode; +import java.io.Closeable; import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; @@ -315,6 +318,117 @@ protected void onClose(Path path, Object stream) { } } + public void testAbortMergeReadsBeforeRead() throws IOException { + try (LocalReopeningIndexDirectory fixture = newLocalReopeningIndexDirectory("testAbortMergeReadsBeforeRead")) { + final byte[] bytes = writeLocalFile(fixture.indexDirectory); + final String fileName = fixture.indexDirectory.listAll()[0]; + try (IndexInput input = fixture.indexDirectory.openInput(fileName, IOContext.DEFAULT)) { + assertThat(input, instanceOf(IndexDirectory.ReopeningIndexInput.class)); + assertFalse(fixture.indexDirectory.shouldAbortMergeReads()); + readNBytes(input, bytes, 0L, 16); + } + } + } + + public void testAbortMergeReadsThrowsOnRead() throws IOException { + try (LocalReopeningIndexDirectory fixture = newLocalReopeningIndexDirectory("testAbortMergeReadsThrowsOnRead")) { + final byte[] bytes = writeLocalFile(fixture.indexDirectory); + final String fileName = fixture.indexDirectory.listAll()[0]; + try (IndexInput input = fixture.indexDirectory.openInput(fileName, IOContext.DEFAULT)) { + assertThat(input, instanceOf(IndexDirectory.ReopeningIndexInput.class)); + readNBytes(input, bytes, 0L, 16); + } + fixture.indexDirectory.abortMergeReads(); + expectThrows(MergePolicy.MergeAbortedException.class, () -> { + try (IndexInput input = fixture.indexDirectory.openInput(fileName, IOContext.merge(new MergeInfo(100, 1024L, false, -1)))) { + input.readByte(); + } + }); + } + } + + public void testAbortMergeReadsIgnoresNonMergeReadAfterAbort() throws IOException { + try (LocalReopeningIndexDirectory fixture = newLocalReopeningIndexDirectory("testAbortMergeReadsIgnoresNonMergeReadAfterAbort")) { + final byte[] bytes = writeLocalFile(fixture.indexDirectory); + final String fileName = fixture.indexDirectory.listAll()[0]; + fixture.indexDirectory.abortMergeReads(); + try (IndexInput input = fixture.indexDirectory.openInput(fileName, IOContext.DEFAULT)) { + assertThat(input, instanceOf(IndexDirectory.ReopeningIndexInput.class)); + readNBytes(input, bytes, 0L, 16); + } + } + } + + private LocalReopeningIndexDirectory newLocalReopeningIndexDirectory(String threadPoolName) throws IOException { + final Path dataPath = createTempDir(); + final ShardId shardId = new ShardId(new Index("_index_name", "_index_id"), 0); + final ThreadPool threadPool = getThreadPool(threadPoolName); + final ByteSizeValue cacheSize = ByteSizeValue.ofBytes(randomLongBetween(0, 10_000_000)); + final var settings = Settings.builder() + .put(SharedBlobCacheService.SHARED_CACHE_SIZE_SETTING.getKey(), cacheSize) + .put( + SharedBlobCacheService.SHARED_CACHE_REGION_SIZE_SETTING.getKey(), + pageAligned(ByteSizeValue.of(randomIntBetween(4, 1024), ByteSizeUnit.KB)) + ) + .put(Environment.PATH_HOME_SETTING.getKey(), createTempDir()) + .putList(Environment.PATH_DATA_SETTING.getKey(), dataPath.toAbsolutePath().toString()) + .build(); + final Path indexDataPath = dataPath.resolve("index"); + final Path blobStorePath = createTempDir(); + final NodeEnvironment nodeEnvironment = new NodeEnvironment(settings, TestEnvironment.newEnvironment(settings)); + final StatelessSharedBlobCacheService sharedBlobCacheService = newCacheService(nodeEnvironment, settings, threadPool); + final IndexDirectory indexDirectory = new IndexDirectory( + newFSDirectory(indexDataPath), + new IndexBlobStoreCacheDirectory(sharedBlobCacheService, shardId), + null, + true + ); + final FsBlobContainer blobContainer = new FsBlobContainer( + new FsBlobStore(1024, blobStorePath, false), + BlobPath.EMPTY, + blobStorePath + ); + indexDirectory.getBlobStoreCacheDirectory().setBlobContainer(value -> blobContainer); + return new LocalReopeningIndexDirectory(nodeEnvironment, sharedBlobCacheService, indexDirectory, threadPool); + } + + private static final class LocalReopeningIndexDirectory implements Closeable { + private final NodeEnvironment nodeEnvironment; + private final StatelessSharedBlobCacheService sharedBlobCacheService; + private final IndexDirectory indexDirectory; + private final ThreadPool threadPool; + + private LocalReopeningIndexDirectory( + NodeEnvironment nodeEnvironment, + StatelessSharedBlobCacheService sharedBlobCacheService, + IndexDirectory indexDirectory, + ThreadPool threadPool + ) { + this.nodeEnvironment = nodeEnvironment; + this.sharedBlobCacheService = sharedBlobCacheService; + this.indexDirectory = indexDirectory; + this.threadPool = threadPool; + } + + @Override + public void close() throws IOException { + indexDirectory.close(); + sharedBlobCacheService.close(); + nodeEnvironment.close(); + assertTrue(ThreadPool.terminate(threadPool, 10L, TimeUnit.SECONDS)); + } + } + + private byte[] writeLocalFile(IndexDirectory indexDirectory) throws IOException { + final int fileLength = randomIntBetween(1024, 10240); + final byte[] bytes = randomByteArrayOfLength(fileLength); + final String fileName = "file." + randomFrom(LuceneFilesExtensions.values()).getExtension(); + try (IndexOutput output = indexDirectory.createOutput(fileName, IOContext.DEFAULT)) { + output.writeBytes(bytes, fileLength); + } + return bytes; + } + /** * See Lucene change https://github.com/apache/lucene/pull/13535 * @throws IOException