Make things testable

Author: carlosdelest

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

@@ -2850,7 +2850,7 @@ public synchronized void beforeMerge(OnGoingMerge merge) {
         @Override
         protected long estimateMergeMemory(MergePolicy.OneMerge merge) {
             try (Searcher searcher = acquireSearcher("merge_memory_estimation", SearcherScope.INTERNAL)) {
-                return SegmentMergeMemoryEstimator.estimateSegmentMemory(merge, searcher.getIndexReader());
+                return SegmentMergeMemoryEstimator.estimateMergeMemory(merge, searcher.getIndexReader());
             } catch (AlreadyClosedException e) {
                 failOnTragicEvent(e);
                 return 0L;

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

@@ -25,28 +25,28 @@
 
 public class SegmentMergeMemoryEstimator {
 
-    public static long estimateSegmentMemory(MergePolicy.OneMerge merge, IndexReader indexReader) {
+    public static long estimateMergeMemory(MergePolicy.OneMerge merge, IndexReader indexReader) {
         long memoryNeeded = 0;
         for (SegmentCommitInfo mergedSegment : merge.segments) {
-            memoryNeeded += estimateSegmentMemory(mergedSegment.info.name, indexReader);
+            memoryNeeded += estimateMergeMemory(mergedSegment.info.name, indexReader);
         }
         return memoryNeeded;
     }
 
-    private static long estimateSegmentMemory(String segmentName, IndexReader indexReader) {
+    private static long estimateMergeMemory(String segmentName, IndexReader indexReader) {
         List<LeafReaderContext> leaves = indexReader.leaves();
         for (LeafReaderContext leafReaderContext : leaves) {
             SegmentReader segmentReader = Lucene.segmentReader(leafReaderContext.reader());
             if (segmentReader.getSegmentName().equals(segmentName)) {
-                return estimateSegmentMemory(segmentReader);
+                return estimateMergeMemory(segmentReader);
             }
         }
 
         // Segment not found, we can't estimate it
         return 0;
     }
 
-    private static long estimateSegmentMemory(SegmentReader reader) {
+    private static long estimateMergeMemory(SegmentReader reader) {
         long maxMem = 0;
         for (FieldInfo fieldInfo : reader.getFieldInfos()) {
             maxMem = Math.max(maxMem, estimateFieldMemory(fieldInfo, reader));

server/src/main/java/org/elasticsearch/index/merge/OnGoingMerge.java

@@ -56,4 +56,8 @@ public List<SegmentCommitInfo> getMergedSegments() {
     public MergePolicy.OneMerge getMerge() {
         return oneMerge;
     }
+
+    public long getMemoryBytesNeeded() {
+        return memoryBytesNeeded;
+    }
 }

server/src/test/java/org/elasticsearch/index/engine/InternalEngineTests.java

@@ -114,6 +114,7 @@
 import org.elasticsearch.index.mapper.SourceFieldMapper;
 import org.elasticsearch.index.mapper.Uid;
 import org.elasticsearch.index.mapper.VersionFieldMapper;
+import org.elasticsearch.index.merge.OnGoingMerge;
 import org.elasticsearch.index.seqno.LocalCheckpointTracker;
 import org.elasticsearch.index.seqno.ReplicationTracker;
 import org.elasticsearch.index.seqno.RetentionLease;
@@ -2639,6 +2640,85 @@ public void testMergeThreadLogging() throws Exception {
         }
     }
 
+    public void testMergeEstimatesMemorySize() throws Exception {
+
+        try (Store store = createStore()) {
+            LogMergePolicy lmp = newLogMergePolicy();
+            lmp.setMergeFactor(2);
+
+            EngineConfig config = config(defaultSettings, store, createTempDir(), lmp, null);
+
+            store.createEmpty();
+            final String translogUuid = Translog.createEmptyTranslog(
+                config.getTranslogConfig().getTranslogPath(),
+                SequenceNumbers.NO_OPS_PERFORMED,
+                shardId,
+                primaryTerm.get()
+            );
+            store.associateIndexWithNewTranslog(translogUuid);
+
+            AtomicInteger estimatedMerges = new AtomicInteger();
+            AtomicInteger actualMerges = new AtomicInteger();
+            Set<Long> estimatedMemorySizes = Collections.synchronizedSet(new HashSet<>());
+
+            InternalEngine engine = new InternalTestEngine(config) {
+                @Override
+                protected IndexWriter createWriter(Directory directory, IndexWriterConfig iwc) throws IOException {
+                    return new IndexWriter(directory, iwc) {
+                        @Override
+                        public void merge(MergePolicy.OneMerge merge) throws IOException {
+                            actualMerges.addAndGet(1);
+                            super.merge(merge);
+                        }
+                    };
+                }
+
+                @Override
+                protected ElasticsearchMergeScheduler createMergeScheduler(ShardId shardId, IndexSettings indexSettings) {
+                    return new ElasticsearchConcurrentMergeScheduler(shardId, indexSettings) {
+                        @Override
+                        protected long estimateMergeMemory(MergePolicy.OneMerge merge) {
+                            // Estimate merge memory randomly, and record the estimated merges and memory sizes for later comparison
+                            long estimation = randomLongBetween(1, 100000000L);
+                            estimatedMerges.addAndGet(1);
+                            estimatedMemorySizes.add(estimation);
+                            return estimation;
+                        }
+
+                        @Override
+                        protected void beforeMerge(OnGoingMerge merge) {
+                            // Checks the estimation is available before merging, and that it has been done previously
+                            assertThat(merge.getMemoryBytesNeeded(), Matchers.greaterThan(0L));
+                            assertTrue(estimatedMemorySizes.remove(merge.getMemoryBytesNeeded()));
+                        }
+
+                        @Override
+                        protected void afterMerge(OnGoingMerge merge) {}
+                    };
+                }
+            };
+            recoverFromTranslog(engine, translogHandler, Long.MAX_VALUE);
+
+            int numDocs = randomIntBetween(10, 100);
+            for (int i = 0; i < numDocs; i++) {
+                engine.index(indexForDoc(testParsedDocument(randomIdentifier(), null, testDocument(), B_1, null)));
+                if (randomBoolean()) {
+                    engine.flush();
+                    engine.forceMerge(true, 1, false, UUIDs.randomBase64UUID());
+                }
+            }
+
+            assertBusy(() -> {
+                // All scheduled merges have been done
+                assertThat(estimatedMerges.get(), is(actualMerges.get()));
+                // All estimations have been accounted for
+                assertTrue(estimatedMemorySizes.isEmpty());
+            });
+
+            engine.close();
+        }
+    }
+
     public void testSeqNoAndCheckpoints() throws IOException, InterruptedException {
         final int opCount = randomIntBetween(1, 256);
         long primarySeqNo = SequenceNumbers.NO_OPS_PERFORMED;