From 3ff67c48ea62550bdccd383014e8f9333ac50417 Mon Sep 17 00:00:00 2001 From: Andrew Ross Date: Wed, 29 Jul 2026 16:21:12 +0000 Subject: [PATCH] Speed up reading history operations from Lucene MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously when reading history, every document would fetch a fresh StoredFields instance, which decompress that compression block for every document. This optimization detects a contiguous gapless range of docIDs and serves them from the same (sequential) stored fields reader. The optimization is gated on the same thread that created the StoredFields instance must read it, as Lucene documents that instances are not safe to hand across threads. Additionally, a min batch size threshold must be met to avoid decompress the block eagerly in the case where it does not matter. Benchmark results: ┌────────────────────┬──────────┬─────────┐ │ layout │ baseline │ change │ ├────────────────────┼──────────┼─────────┤ │ CONTIGUOUS │ 331 ms │ 19.2 ms │ ├────────────────────┼──────────┼─────────┤ │ STRIDED (stride 2) │ 379 ms │ 324 ms │ ├────────────────────┼──────────┼─────────┤ │ INTERLEAVED │ 334 ms │ 312 ms │ └────────────────────┴──────────┴─────────┘ Signed-off-by: Andrew Ross --- .../LuceneChangesSnapshotBenchmark.java | 275 ++++++++++++++++++ .../index/engine/LuceneChangesSnapshot.java | 99 ++++++- .../engine/LuceneChangesSnapshotTests.java | 241 +++++++++++++++ 3 files changed, 609 insertions(+), 6 deletions(-) create mode 100644 benchmarks/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshotBenchmark.java diff --git a/benchmarks/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshotBenchmark.java b/benchmarks/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshotBenchmark.java new file mode 100644 index 0000000000000..b5f921d3ef0ca --- /dev/null +++ b/benchmarks/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshotBenchmark.java @@ -0,0 +1,275 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine; + +import org.apache.lucene.codecs.lucene104.Lucene104Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.Sort; +import org.apache.lucene.search.SortField; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.TopFieldCollectorManager; +import org.apache.lucene.search.UsageTrackingQueryCachingPolicy; +import org.apache.lucene.search.similarities.BM25Similarity; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.NIOFSDirectory; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.IOUtils; +import org.opensearch.common.lucene.index.OpenSearchDirectoryReader; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.mapper.IdFieldMapper; +import org.opensearch.index.mapper.SeqNoFieldMapper; +import org.opensearch.index.mapper.SourceFieldMapper; +import org.opensearch.index.mapper.Uid; +import org.opensearch.index.mapper.VersionFieldMapper; +import org.opensearch.index.translog.Translog; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Random; +import java.util.concurrent.TimeUnit; + +/** + * Benchmark for reading history from LuceneChangesSnapshot. + * + *

The {@link Layout} parameter controls how the operations that get read are placed across docIDs and segments, + * which is what determines whether a sequential read is possible at all: + * + *

+ * + * Run a single configuration with, for example: + * 
+ * ./gradlew -p benchmarks run --args 'LuceneChangesSnapshotBenchmark -p layout=CONTIGUOUS -p docSizeBytes=256'
+ *
+ */ +@Fork(value = 1, jvmArgsAppend = { "-Xms2g", "-Xmx2g" }) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@State(Scope.Benchmark) +public class LuceneChangesSnapshotBenchmark { + + /** + * How the operations that will be read are laid out across docIDs and segments. See the class javadoc. + */ + public enum Layout { + CONTIGUOUS, + STRIDED, + INTERLEAVED + } + + /** Operations read per invocation. Every layout reads this many, only their placement differs. */ + @Param({ "20000" }) + private int numOps; + + /** + * For {@code STRIDED}, how many docIDs apart consecutive reads are; for {@code INTERLEAVED}, how many segments the + * seqNo range is interleaved across (the flush fan-out). {@code 1} makes both layouts contiguous. + */ + @Param({ "2" }) + private int stride; + + @Param({ "CONTIGUOUS", "STRIDED", "INTERLEAVED" }) + private Layout layout; + + @Param({ "256", "2048" }) + private int docSizeBytes; + + @Param({ "BEST_SPEED", "BEST_COMPRESSION" }) + private String codecMode; + + private static final ShardId SHARD_ID = new ShardId("index", "_na_", 0); + + private Path indexPath; + private Directory directory; + private DirectoryReader reader; + + @Setup(Level.Trial) + public void setup() throws IOException { + indexPath = Files.createTempDirectory("lucene-changes-snapshot-benchmark"); + // OpenSearch's HybridDirectory reads stored fields through NIOFS rather than mmap by default; see + // IndexModule#INDEX_STORE_HYBRID_NIO_EXTENSIONS, which lists fdt and fdx. + directory = new NIOFSDirectory(indexPath); + buildIndex(); + reader = OpenSearchDirectoryReader.wrap(DirectoryReader.open(directory), SHARD_ID); + verifyLayout(); + } + + @TearDown(Level.Trial) + public void tearDown() throws IOException { + IOUtils.close(reader, directory); + IOUtils.rm(indexPath); + } + + /** + * Drains the whole snapshot, the way the get-changes action does: the seqNo range query, the doc values reads and + * the stored fields reads for every operation. + */ + @Benchmark + public void drainSnapshot(Blackhole bh) throws IOException { + try (LuceneChangesSnapshot snapshot = newSnapshot()) { + Translog.Operation op; + while ((op = snapshot.next()) != null) { + bh.consume(op); + } + } + } + + private LuceneChangesSnapshot newSnapshot() throws IOException { + final Engine.Searcher searcher = new Engine.Searcher( + "benchmark", + reader, + new BM25Similarity(), + null, + new UsageTrackingQueryCachingPolicy(), + () -> {} // the reader outlives the snapshot and is closed by tearDown + ); + return new LuceneChangesSnapshot(searcher, LuceneChangesSnapshot.DEFAULT_BATCH_SIZE, 0, numOps - 1, true, true); + } + + private void buildIndex() throws IOException { + final Lucene104Codec.Mode mode = Lucene104Codec.Mode.valueOf(codecMode); + final IndexWriterConfig config = new IndexWriterConfig().setCodec(new Lucene104Codec(mode)) + // one segment per commit, so the layouts below are exactly what gets read + .setMergePolicy(NoMergePolicy.INSTANCE) + .setRAMBufferSizeMB(512); + try (IndexWriter writer = new IndexWriter(directory, config)) { + switch (layout) { + case CONTIGUOUS: + // seqNo == docID, one segment + for (int i = 0; i < numOps; i++) { + writer.addDocument(newDocument(i)); + } + break; + case STRIDED: + // one segment; the range that gets read covers every stride-th docID, the rest hold seqNos above it + long filler = numOps; + for (int i = 0; i < numOps; i++) { + writer.addDocument(newDocument(i)); + for (int gap = 1; gap < stride; gap++) { + writer.addDocument(newDocument(filler++)); + } + } + break; + case INTERLEAVED: + // segments with interleaved seqNo ranges, as concurrently flushed segments have + for (int segment = 0; segment < stride; segment++) { + for (int seqNo = segment; seqNo < numOps; seqNo += stride) { + writer.addDocument(newDocument(seqNo)); + } + writer.commit(); + } + break; + default: + throw new AssertionError(layout); + } + writer.commit(); + } + } + + private Document newDocument(long seqNo) { + final Document doc = new Document(); + doc.add(new StoredField(IdFieldMapper.NAME, Uid.encodeId(Long.toString(seqNo)))); + doc.add(new StoredField(SourceFieldMapper.NAME, new BytesRef(source(seqNo)))); + doc.add(new LongPoint(SeqNoFieldMapper.NAME, seqNo)); + doc.add(new NumericDocValuesField(SeqNoFieldMapper.NAME, seqNo)); + doc.add(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, 1L)); + doc.add(new NumericDocValuesField(VersionFieldMapper.NAME, 1L)); + return doc; + } + + /** + * A pseudo-random payload of the requested size. Random bytes would not compress at all and identical bytes would + * compress perfectly, so this mixes a fixed pattern with random content to land somewhere realistic. + */ + private byte[] source(long seqNo) { + final byte[] bytes = new byte[docSizeBytes]; + final Random random = new Random(seqNo); + for (int i = 0; i < bytes.length; i++) { + bytes[i] = (byte) (i % 8 == 0 ? random.nextInt(26) + 'a' : 'a' + (i % 26)); + } + return bytes; + } + + /** + * Fails fast if the index does not have the shape the layout intends, so a run cannot quietly measure something + * else. Checks that the expected operations are all present and readable, and that the docID sequence in read order + * is gapless and ascending only for {@code CONTIGUOUS} - the property that decides whether a sequential read of + * stored fields is possible. Deliberately asserts on the index rather than on the snapshot's internals, so this + * benchmark also compiles and runs against a baseline checkout. + */ + private void verifyLayout() throws IOException { + final IndexSearcher searcher = new IndexSearcher(reader); + searcher.setQueryCache(null); + final TopDocs topDocs = searcher.search( + LongPoint.newRangeQuery(SeqNoFieldMapper.NAME, 0, numOps - 1), + new TopFieldCollectorManager(new Sort(new SortField(SeqNoFieldMapper.NAME, SortField.Type.LONG)), numOps, null, 0, false) + ); + if (topDocs.scoreDocs.length != numOps) { + throw new IllegalStateException("expected " + numOps + " hits but got " + topDocs.scoreDocs.length); + } + boolean contiguous = true; + for (int i = 0; i < topDocs.scoreDocs.length - 1; i++) { + final ScoreDoc current = topDocs.scoreDocs[i]; + final ScoreDoc next = topDocs.scoreDocs[i + 1]; + if (current.doc + 1 != next.doc) { + contiguous = false; + break; + } + } + if (contiguous != (layout == Layout.CONTIGUOUS)) { + throw new IllegalStateException( + "layout [" + layout + "] with stride [" + stride + "] produced a contiguous read order of [" + contiguous + "]" + ); + } + int read = 0; + try (LuceneChangesSnapshot snapshot = newSnapshot()) { + while (snapshot.next() != null) { + read++; + } + } + if (read != numOps) { + throw new IllegalStateException("expected " + numOps + " operations but read " + read); + } + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshot.java b/server/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshot.java index 72ccc097e20d6..2744ec76ff80c 100644 --- a/server/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/LuceneChangesSnapshot.java @@ -33,9 +33,11 @@ package org.opensearch.index.engine; import org.apache.lucene.document.LongPoint; +import org.apache.lucene.index.FilterLeafReader; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.StoredFields; import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.FieldDoc; @@ -49,6 +51,7 @@ import org.apache.lucene.search.TopFieldCollectorManager; import org.apache.lucene.util.ArrayUtil; import org.opensearch.common.lucene.Lucene; +import org.opensearch.common.lucene.index.SequentialStoredFieldsLeafReader; import org.opensearch.common.lucene.search.Queries; import org.opensearch.common.util.io.IOUtils; import org.opensearch.core.common.bytes.BytesReference; @@ -71,6 +74,12 @@ final class LuceneChangesSnapshot implements Translog.Snapshot { static final int DEFAULT_BATCH_SIZE = 1024; + /** + * Batches smaller than this are read with the default stored fields reader: acquiring a sequential reader costs a + * reader clone plus the decompression of a whole block, which only pays for itself over enough documents. + */ + static final int MIN_SEQUENTIAL_ACCESS_BATCH_SIZE = 10; + private final int searchBatchSize; private final long fromSeqNo, toSeqNo; private long lastSeenSeqNo; @@ -84,6 +93,18 @@ final class LuceneChangesSnapshot implements Translog.Snapshot { private final ParallelArray parallelArray; private final Closeable onClose; + /** + * Whether to use a sequential stored fields reader. Only used when there is a gapless run of docIDs + * matching sequence number order. This allows for an optimization to decompress a block once and + * read multiple stored fields from it, as opposed to the default path that does a decompression on + * every stored fields read. + */ + private boolean useSequentialStoredFieldsReader; + + private StoredFields sequentialStoredFields; + private int sequentialStoredFieldsLeafOrd = -1; + private Thread sequentialStoredFieldsOwner; + /** * Creates a new "translog" snapshot from Lucene for reading operations whose seq# in the specified range. * @@ -130,6 +151,7 @@ final class LuceneChangesSnapshot implements Translog.Snapshot { @Override public void close() throws IOException { + releaseSequentialStoredFieldsReader(); onClose.close(); } @@ -215,9 +237,13 @@ private void fillParallelArray(ScoreDoc[] scoreDocs, ParallelArray parallelArray for (int i = 0; i < scoreDocs.length; i++) { scoreDocs[i].shardIndex = i; } - // for better loading performance we sort the array by docID and - // then visit all leaves in order. - ArrayUtil.introSort(scoreDocs, Comparator.comparingInt(i -> i.doc)); + useSequentialStoredFieldsReader = scoreDocs.length >= MIN_SEQUENTIAL_ACCESS_BATCH_SIZE && hasSequentialAccess(scoreDocs); + if (useSequentialStoredFieldsReader == false) { + releaseSequentialStoredFieldsReader(); + // for better loading performance we sort the array by docID and + // then visit all leaves in order. + ArrayUtil.introSort(scoreDocs, Comparator.comparingInt(i -> i.doc)); + } int docBase = -1; int maxDoc = 0; List leaves = indexSearcher.getIndexReader().leaves(); @@ -242,9 +268,23 @@ private void fillParallelArray(ScoreDoc[] scoreDocs, ParallelArray parallelArray parallelArray.isTombStone[index] = combinedDocValues.isTombstone(segmentDocID); parallelArray.hasRecoverySource[index] = combinedDocValues.hasRecoverySource(segmentDocID); } - // now sort back based on the shardIndex. we use this to store the previous index - ArrayUtil.introSort(scoreDocs, Comparator.comparingInt(i -> i.shardIndex)); + if (useSequentialStoredFieldsReader == false) { + // now sort back based on the shardIndex. we use this to store the previous index + ArrayUtil.introSort(scoreDocs, Comparator.comparingInt(i -> i.shardIndex)); + } + } + } + + /** + * Returns whether the batch, in the order it will be read, is a strictly increasing and gapless run of docIDs + */ + private static boolean hasSequentialAccess(ScoreDoc[] scoreDocs) { + for (int i = 0; i < scoreDocs.length - 1; i++) { + if (scoreDocs[i].doc + 1 != scoreDocs[i + 1].doc) { + return false; + } } + return true; } private static Query operationsRangeQuery(long fromSeqNo, long toSeqNo) { @@ -291,7 +331,7 @@ private Translog.Operation readDocAsOp(int docIndex) throws IOException { ? SourceFieldMapper.RECOVERY_SOURCE_NAME : SourceFieldMapper.NAME; final FieldsVisitor fields = new FieldsVisitor(true, sourceField); - leaf.reader().storedFields().document(segmentDocID, fields); + storedFieldsFor(leaf).document(segmentDocID, fields); final Translog.Operation op; final boolean isTombstone = parallelArray.isTombStone[docIndex]; @@ -344,6 +384,53 @@ private Translog.Operation readDocAsOp(int docIndex) throws IOException { return op; } + private StoredFields storedFieldsFor(LeafReaderContext leaf) throws IOException { + if (useSequentialStoredFieldsReader) { + // A StoredFields instance may only be consumed in the thread that acquired it, + // so track the owning thread and recreate it if consumed by another thread. + final Thread currentThread = Thread.currentThread(); + if (sequentialStoredFields == null + || sequentialStoredFieldsLeafOrd != leaf.ord + || sequentialStoredFieldsOwner != currentThread) { + sequentialStoredFields = sequentialStoredFields(leaf.reader()); + sequentialStoredFieldsLeafOrd = leaf.ord; + sequentialStoredFieldsOwner = currentThread; + } + return sequentialStoredFields; + } + return leaf.reader().storedFields(); + } + + /** + * Returns a merge-optimized (sequential) stored fields reader for the leaf + */ + private static StoredFields sequentialStoredFields(LeafReader leaf) throws IOException { + LeafReader reader = leaf; + while (true) { + if (reader instanceof SequentialStoredFieldsLeafReader) { + return ((SequentialStoredFieldsLeafReader) reader).getSequentialStoredFieldsReader(); + } + if (reader instanceof FilterLeafReader == false) { + return leaf.storedFields(); + } + reader = ((FilterLeafReader) reader).getDelegate(); + } + } + + /** + * Drops the cached reader, if any + */ + private void releaseSequentialStoredFieldsReader() { + sequentialStoredFields = null; + sequentialStoredFieldsLeafOrd = -1; + sequentialStoredFieldsOwner = null; + } + + // for testing + boolean useSequentialStoredFieldsReader() { + return useSequentialStoredFieldsReader; + } + private boolean assertDocSoftDeleted(LeafReader leafReader, int segmentDocId) throws IOException { final NumericDocValues ndv = leafReader.getNumericDocValues(Lucene.SOFT_DELETES_FIELD); if (ndv == null || ndv.advanceExact(segmentDocId) == false) { diff --git a/server/src/test/java/org/opensearch/index/engine/LuceneChangesSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/LuceneChangesSnapshotTests.java index b9b0e64d8811f..dfc35e8e99d6f 100644 --- a/server/src/test/java/org/opensearch/index/engine/LuceneChangesSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/LuceneChangesSnapshotTests.java @@ -32,6 +32,13 @@ package org.opensearch.index.engine; +import org.apache.lucene.codecs.StoredFieldsReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FilterDirectoryReader; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.StoredFieldVisitor; +import org.opensearch.common.CheckedSupplier; +import org.opensearch.common.lucene.index.SequentialStoredFieldsLeafReader; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.io.IOUtils; import org.opensearch.index.IndexSettings; @@ -49,10 +56,14 @@ import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasSize; +import static org.hamcrest.Matchers.instanceOf; public class LuceneChangesSnapshotTests extends EngineTestCase { private MapperService mapperService; @@ -369,4 +380,234 @@ public void testOverFlow() throws Exception { ); } } + + public void testOutOfOrderSeqNoUsesDefaultStoredFieldsReader() throws Exception { + final int numOps = between(LuceneChangesSnapshot.MIN_SEQUENTIAL_ACCESS_BATCH_SIZE, 100); + final Map expectedDocs = new HashMap<>(); + for (int i = 0; i < numOps; i++) { + final long seqNo = numOps - 1 - i; + final ParsedDocument doc = createParsedDoc("id-" + seqNo, null); + engine.index(replicaIndexForDoc(doc, 1, seqNo, false)); + expectedDocs.put(seqNo, doc); + } + engine.refresh("test"); + Engine.Searcher searcher = engine.acquireSearcher("test", Engine.SearcherScope.INTERNAL); + try (LuceneChangesSnapshot snapshot = new LuceneChangesSnapshot(searcher, numOps, 0, numOps - 1, true, true)) { + searcher = null; + assertThat(snapshot.totalOperations(), equalTo(numOps)); + // the flag is computed for the first batch by the constructor; a single batch holds every op here + assertFalse("descending docIDs must not use the sequential reader", snapshot.useSequentialStoredFieldsReader()); + assertOpsMatch(drainAll(snapshot), expectedDocs); + } finally { + IOUtils.close(searcher); + } + } + + public void testInterleavedSeqNosAcrossSegmentsUseDefaultStoredFieldsReader() throws Exception { + final int numOps = 2 * between(LuceneChangesSnapshot.MIN_SEQUENTIAL_ACCESS_BATCH_SIZE, 50); + final Map expectedDocs = new HashMap<>(); + for (int step = 0; step < 2; step++) { + for (long seqNo = step; seqNo < numOps; seqNo += 2) { + final ParsedDocument doc = createParsedDoc("id-" + seqNo, null); + engine.index(replicaIndexForDoc(doc, 1, seqNo, false)); + expectedDocs.put(seqNo, doc); + } + engine.flush(); + } + engine.refresh("test"); + Engine.Searcher searcher = engine.acquireSearcher("test", Engine.SearcherScope.INTERNAL); + try (LuceneChangesSnapshot snapshot = new LuceneChangesSnapshot(searcher, numOps, 0, numOps - 1, true, true)) { + searcher = null; + assertThat(snapshot.totalOperations(), equalTo(numOps)); + assertFalse("interleaved segments must not use the sequential reader", snapshot.useSequentialStoredFieldsReader()); + assertOpsMatch(drainAll(snapshot), expectedDocs); + } finally { + IOUtils.close(searcher); + } + } + + public void testContiguousReadsReuseSequentialStoredFieldsReader() throws Exception { + final int batchSize = LuceneChangesSnapshot.MIN_SEQUENTIAL_ACCESS_BATCH_SIZE; + final int numOps = batchSize * between(3, 10); + final Map expectedDocs = indexAppendOnly(numOps); + final AtomicInteger acquisitions = new AtomicInteger(); + final AtomicInteger documentsRead = new AtomicInteger(); + Engine.Searcher searcher = engine.acquireSearcher("test", Engine.SearcherScope.INTERNAL); + try { + final int numLeaves = searcher.getDirectoryReader().leaves().size(); + // a batch size well below the op count forces multiple batches, exercising reader reuse across them + final Engine.Searcher countingSearcher = countingSearcher(searcher, acquisitions, documentsRead); + try (LuceneChangesSnapshot snapshot = new LuceneChangesSnapshot(countingSearcher, batchSize, 0, numOps - 1, true, true)) { + searcher = null; // closed by the snapshot through the counting searcher + assertTrue("contiguous docIDs must use the sequential reader", snapshot.useSequentialStoredFieldsReader()); + assertOpsMatch(drainAll(snapshot), expectedDocs); + assertThat("sequential reader must be acquired once per leaf and reused", acquisitions.get(), equalTo(numLeaves)); + assertThat("every read must go through the SequentialStoredFieldsLeafReader", documentsRead.get(), equalTo(numOps)); + } + } finally { + IOUtils.close(searcher); + } + } + + public void testSequentialStoredFieldsReaderIsReacquiredOnThreadChange() throws Exception { + final int numOps = between(LuceneChangesSnapshot.MIN_SEQUENTIAL_ACCESS_BATCH_SIZE, 50); + final Map expectedDocs = indexAppendOnly(numOps); + final AtomicInteger acquisitions = new AtomicInteger(); + final AtomicInteger documentsRead = new AtomicInteger(); + Engine.Searcher searcher = engine.acquireSearcher("test", Engine.SearcherScope.INTERNAL); + try { + final Engine.Searcher countingSearcher = countingSearcher(searcher, acquisitions, documentsRead); + try (LuceneChangesSnapshot snapshot = new LuceneChangesSnapshot(countingSearcher, numOps, 0, numOps - 1, true, true)) { + searcher = null; // closed by the snapshot through the counting searcher + assertTrue(snapshot.useSequentialStoredFieldsReader()); + final List ops = new ArrayList<>(); + // every next() runs on a fresh thread; join() gives the happens-before that consumers get from locking + Translog.Operation op = callOnNewThread(snapshot::next); + while (op != null) { + ops.add(op); + op = callOnNewThread(snapshot::next); + } + assertOpsMatch(ops, expectedDocs); + assertThat("reader must be re-acquired for every reading thread", acquisitions.get(), equalTo(numOps)); + assertThat(documentsRead.get(), equalTo(numOps)); + } + } finally { + IOUtils.close(searcher); + } + } + + private Map indexAppendOnly(int numOps) throws IOException { + final Map expectedDocs = new HashMap<>(); + for (int i = 0; i < numOps; i++) { + final ParsedDocument doc = createParsedDoc("id-" + i, null); + final Engine.IndexResult result = engine.index(indexForDoc(doc)); + expectedDocs.put(result.getSeqNo(), doc); + } + engine.refresh("test"); + return expectedDocs; + } + + private void assertOpsMatch(List ops, Map expectedDocs) { + assertThat(ops, hasSize(expectedDocs.size())); + for (Translog.Operation op : ops) { + assertThat(op.toString(), op, instanceOf(Translog.Index.class)); + final Translog.Index index = (Translog.Index) op; + final ParsedDocument expected = expectedDocs.get(op.seqNo()); + assertNotNull("unexpected seqNo [" + op.seqNo() + "]", expected); + assertThat(index.id(), equalTo(expected.id())); + assertThat(index.source(), equalTo(expected.source())); + } + } + + private static T callOnNewThread(CheckedSupplier supplier) throws Exception { + final AtomicReference result = new AtomicReference<>(); + final AtomicReference failure = new AtomicReference<>(); + final Thread thread = new Thread(() -> { + try { + result.set(supplier.get()); + } catch (Exception e) { + failure.set(e); + } + }); + thread.start(); + thread.join(); + if (failure.get() != null) { + throw failure.get(); + } + return result.get(); + } + + /** + * Wraps the searcher's reader so that every leaf sits behind a counting {@link SequentialStoredFieldsLeafReader}. + * The returned searcher takes ownership of {@code searcher}. + */ + private static Engine.Searcher countingSearcher(Engine.Searcher searcher, AtomicInteger acquisitions, AtomicInteger documentsRead) + throws IOException { + final DirectoryReader reader = new CountingSequentialDirectoryReader(searcher.getDirectoryReader(), acquisitions, documentsRead); + return new Engine.Searcher( + searcher.source(), + reader, + searcher.getSimilarity(), + searcher.getQueryCache(), + searcher.getQueryCachingPolicy(), + searcher + ); + } + + /** + * Wraps every leaf in a pass-through {@link SequentialStoredFieldsLeafReader} that counts how many times its + * sequential stored fields reader is acquired, and how many documents that reader serves. + */ + private static final class CountingSequentialDirectoryReader extends FilterDirectoryReader { + private final AtomicInteger acquisitions; + private final AtomicInteger documentsRead; + + CountingSequentialDirectoryReader(DirectoryReader in, AtomicInteger acquisitions, AtomicInteger documentsRead) throws IOException { + super(in, new SubReaderWrapper() { + @Override + public LeafReader wrap(LeafReader reader) { + return new SequentialStoredFieldsLeafReader(reader) { + @Override + protected StoredFieldsReader doGetSequentialStoredFieldsReader(StoredFieldsReader storedFieldsReader) { + acquisitions.incrementAndGet(); + return new CountingStoredFieldsReader(storedFieldsReader, documentsRead); + } + + @Override + public CacheHelper getCoreCacheHelper() { + return reader.getCoreCacheHelper(); + } + + @Override + public CacheHelper getReaderCacheHelper() { + return reader.getReaderCacheHelper(); + } + }; + } + }); + this.acquisitions = acquisitions; + this.documentsRead = documentsRead; + } + + @Override + protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) throws IOException { + return new CountingSequentialDirectoryReader(in, acquisitions, documentsRead); + } + + @Override + public CacheHelper getReaderCacheHelper() { + return in.getReaderCacheHelper(); + } + } + + private static final class CountingStoredFieldsReader extends StoredFieldsReader { + private final StoredFieldsReader delegate; + private final AtomicInteger documentsRead; + + CountingStoredFieldsReader(StoredFieldsReader delegate, AtomicInteger documentsRead) { + this.delegate = delegate; + this.documentsRead = documentsRead; + } + + @Override + public void document(int docID, StoredFieldVisitor visitor) throws IOException { + documentsRead.incrementAndGet(); + delegate.document(docID, visitor); + } + + @Override + public StoredFieldsReader clone() { + return new CountingStoredFieldsReader(delegate.clone(), documentsRead); + } + + @Override + public void checkIntegrity() throws IOException { + delegate.checkIntegrity(); + } + + @Override + public void close() throws IOException { + delegate.close(); + } + } }