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