Expose merge events and their memory usage estimate

pxsalehi · pxsalehi · commit 6fffd735286c · 2025-04-11T15:26:07.000+02:00
diff --git a/server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java b/server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java
@@ -2870,7 +2870,7 @@ private final class EngineThreadPoolMergeScheduler extends ThreadPoolMergeSchedu
             IndexSettings indexSettings,
             ThreadPoolMergeExecutorService threadPoolMergeExecutorService
         ) {
-            super(shardId, indexSettings, threadPoolMergeExecutorService);
+            super(shardId, indexSettings, threadPoolMergeExecutorService, InternalEngine.this::estimateMergeBytes);
         }
 
         @Override
diff --git a/server/src/main/java/org/elasticsearch/index/engine/MergeMemoryEstimateProvider.java b/server/src/main/java/org/elasticsearch/index/engine/MergeMemoryEstimateProvider.java
@@ -0,0 +1,21 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.index.MergePolicy;
+
+@FunctionalInterface
+public interface MergeMemoryEstimateProvider {
+
+    /**
+     * Returns an estimate of the memory needed to perform a merge
+     */
+    long estimateMergeMemoryBytes(MergePolicy.OneMerge merge);
+}
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorService.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorService.java
@@ -14,11 +14,12 @@
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.engine.ThreadPoolMergeScheduler.MergeTask;
-import org.elasticsearch.index.merge.OnGoingMerge;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.util.Comparator;
+import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.PriorityBlockingQueue;
 import java.util.concurrent.RejectedExecutionException;
@@ -74,7 +75,7 @@ public class ThreadPoolMergeExecutorService {
     private final int concurrentMergesFloorLimitForThrottling;
     private final int concurrentMergesCeilLimitForThrottling;
 
-    private volatile MergeEventConsumer mergeEventConsumer;
+    private final List<MergeEventListener> mergeEventListeners = new CopyOnWriteArrayList<>();
 
     public static @Nullable ThreadPoolMergeExecutorService maybeCreateThreadPoolMergeExecutorService(
         ThreadPool threadPool,
@@ -130,13 +131,18 @@ boolean submitMergeTask(MergeTask mergeTask) {
                 );
             }
             // then enqueue the merge task proper
-            queuedMergeTasks.add(mergeTask);
+            enqueueMergeTask(mergeTask);
             return true;
         }
     }
 
     void reEnqueueBackloggedMergeTask(MergeTask mergeTask) {
+        enqueueMergeTask(mergeTask);
+    }
+
+    private void enqueueMergeTask(MergeTask mergeTask) {
         queuedMergeTasks.add(mergeTask);
+        mergeEventListeners.forEach(l -> l.onMergeQueued(mergeTask.getOnGoingMerge(), mergeTask.getEstimateMergeMemoryBytes()));
     }
 
     public boolean allDone() {
@@ -204,6 +210,7 @@ private void runMergeTask(MergeTask mergeTask) {
             if (mergeTask.supportsIOThrottling()) {
                 ioThrottledMergeTasksCount.decrementAndGet();
             }
+            mergeEventListeners.forEach(l -> l.onMergeCompleted(mergeTask.getOnGoingMerge()));
         }
     }
 
@@ -216,6 +223,7 @@ private void abortMergeTask(MergeTask mergeTask) {
             if (mergeTask.supportsIOThrottling()) {
                 ioThrottledMergeTasksCount.decrementAndGet();
             }
+            mergeEventListeners.forEach(l -> l.onMergeAborted(mergeTask.getOnGoingMerge()));
         }
     }
 
@@ -281,15 +289,8 @@ public boolean usingMaxTargetIORateBytesPerSec() {
         return MAX_IO_RATE.getBytes() == targetIORateBytesPerSec.get();
     }
 
-    public void registerMergeEventConsumer(MergeEventConsumer consumer) {
-        assert this.mergeEventConsumer == null;
-        this.mergeEventConsumer = consumer;
-    }
-
-    public interface MergeEventConsumer {
-        void onMergeQueued(OnGoingMerge merge, long estimateMergeMemoryBytes);
-
-        void onMergeCompleted(OnGoingMerge merge);
+    public void registerMergeEventListener(MergeEventListener consumer) {
+        mergeEventListeners.add(consumer);
     }
 
     // exposed for tests
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java
@@ -65,12 +65,13 @@ public class ThreadPoolMergeScheduler extends MergeScheduler implements Elastics
     private final AtomicLong doneMergeTaskCount = new AtomicLong();
     private final CountDownLatch closedWithNoRunningMerges = new CountDownLatch(1);
     private volatile boolean closed = false;
-    private final MergeMemoryEstimator mergeMemoryEstimator;
+    private final MergeMemoryEstimateProvider mergeMemoryEstimateProvider;
 
     public ThreadPoolMergeScheduler(
         ShardId shardId,
         IndexSettings indexSettings,
-        ThreadPoolMergeExecutorService threadPoolMergeExecutorService
+        ThreadPoolMergeExecutorService threadPoolMergeExecutorService,
+        MergeMemoryEstimateProvider mergeMemoryEstimateProvider
     ) {
         this.shardId = shardId;
         this.config = indexSettings.getMergeSchedulerConfig();
@@ -82,6 +83,7 @@ public ThreadPoolMergeScheduler(
                 : Double.POSITIVE_INFINITY
         );
         this.threadPoolMergeExecutorService = threadPoolMergeExecutorService;
+        this.mergeMemoryEstimateProvider = mergeMemoryEstimateProvider;
     }
 
     @Override
@@ -177,11 +179,13 @@ MergeTask newMergeTask(MergeSource mergeSource, MergePolicy.OneMerge merge, Merg
         // forced merges, as well as merges triggered when closing a shard, always run un-IO-throttled
         boolean isAutoThrottle = mergeTrigger != MergeTrigger.CLOSING && merge.getStoreMergeInfo().mergeMaxNumSegments() == -1;
         // IO throttling cannot be toggled for existing merge tasks, only new merge tasks pick up the updated IO throttling setting
+        long estimateMergeMemoryBytes = mergeMemoryEstimateProvider.estimateMergeMemoryBytes(merge);
         return new MergeTask(
             mergeSource,
             merge,
             isAutoThrottle && config.isAutoThrottle(),
-            "Lucene Merge Task #" + submittedMergeTaskCount.incrementAndGet() + " for shard " + shardId
+            "Lucene Merge Task #" + submittedMergeTaskCount.incrementAndGet() + " for shard " + shardId,
+            estimateMergeMemoryBytes
         );
     }
 
@@ -313,14 +317,22 @@ class MergeTask implements Runnable {
         private final OnGoingMerge onGoingMerge;
         private final MergeRateLimiter rateLimiter;
         private final boolean supportsIOThrottling;
-
-        MergeTask(MergeSource mergeSource, MergePolicy.OneMerge merge, boolean supportsIOThrottling, String name) {
+        private final long estimateMergeMemoryBytes;
+
+        MergeTask(
+            MergeSource mergeSource,
+            MergePolicy.OneMerge merge,
+            boolean supportsIOThrottling,
+            String name,
+            long estimateMergeMemoryBytes
+        ) {
             this.name = name;
             this.mergeStartTimeNS = new AtomicLong();
             this.mergeSource = mergeSource;
             this.onGoingMerge = new OnGoingMerge(merge);
             this.rateLimiter = new MergeRateLimiter(merge.getMergeProgress());
             this.supportsIOThrottling = supportsIOThrottling;
+            this.estimateMergeMemoryBytes = estimateMergeMemoryBytes;
         }
 
         Schedule schedule() {
@@ -451,7 +463,7 @@ long estimatedMergeSize() {
         }
 
         public long getEstimateMergeMemoryBytes() {
-            return mergeMemoryEstimator.estimateMergeMemoryBytes(onGoingMerge.getMerge());
+            return estimateMergeMemoryBytes;
         }
 
         public OnGoingMerge getOnGoingMerge() {
diff --git a/server/src/test/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorServiceTests.java b/server/src/test/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorServiceTests.java
@@ -660,6 +660,8 @@ public void testMergeTasksExecuteInSizeOrder() {
         }
     }
 
+    // todo: testMergeEventListeners
+
     static ThreadPoolMergeExecutorService getThreadPoolMergeExecutorService(ThreadPool threadPool) {
         ThreadPoolMergeExecutorService threadPoolMergeExecutorService = ThreadPoolMergeExecutorService
             .maybeCreateThreadPoolMergeExecutorService(
diff --git a/server/src/test/java/org/elasticsearch/index/engine/ThreadPoolMergeSchedulerTests.java b/server/src/test/java/org/elasticsearch/index/engine/ThreadPoolMergeSchedulerTests.java
@@ -61,7 +61,8 @@ public void testMergesExecuteInSizeOrder() throws IOException {
             ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                 new ShardId("index", "_na_", 1),
                 IndexSettingsModule.newIndexSettings("index", Settings.EMPTY),
-                threadPoolMergeExecutorService
+                threadPoolMergeExecutorService,
+                merge -> 0
             )
         ) {
             List<OneMerge> executedMergesList = new ArrayList<>();
@@ -103,7 +104,8 @@ public void testSimpleMergeTaskBacklogging() {
         ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
             new ShardId("index", "_na_", 1),
             IndexSettingsModule.newIndexSettings("index", mergeSchedulerSettings),
-            threadPoolMergeExecutorService
+            threadPoolMergeExecutorService,
+            merge -> 0
         );
         // more merge tasks than merge threads
         int mergeCount = mergeExecutorThreadCount + randomIntBetween(1, 5);
@@ -136,7 +138,8 @@ public void testSimpleMergeTaskReEnqueueingBySize() {
         ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
             new ShardId("index", "_na_", 1),
             IndexSettingsModule.newIndexSettings("index", mergeSchedulerSettings),
-            threadPoolMergeExecutorService
+            threadPoolMergeExecutorService,
+            merge -> 0
         );
         // sort backlogged merges by size
         PriorityQueue<MergeTask> backloggedMergeTasks = new PriorityQueue<>(16, Comparator.comparingLong(MergeTask::estimatedMergeSize));
@@ -347,7 +350,8 @@ public void testMergeSourceWithFollowUpMergesRunSequentially() throws Exception
                 ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                     new ShardId("index", "_na_", 1),
                     IndexSettingsModule.newIndexSettings("index", settings),
-                    threadPoolMergeExecutorService
+                    threadPoolMergeExecutorService,
+                    merge -> 0
                 )
             ) {
                 MergeSource mergeSource = mock(MergeSource.class);
@@ -420,7 +424,8 @@ public void testMergesRunConcurrently() throws Exception {
                 ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                     new ShardId("index", "_na_", 1),
                     IndexSettingsModule.newIndexSettings("index", settings),
-                    threadPoolMergeExecutorService
+                    threadPoolMergeExecutorService,
+                    merge -> 0
                 )
             ) {
                 // at least 1 extra merge than there are concurrently allowed
@@ -504,7 +509,8 @@ public void testSchedulerCloseWaitsForRunningMerge() throws Exception {
                 ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                     new ShardId("index", "_na_", 1),
                     IndexSettingsModule.newIndexSettings("index", settings),
-                    threadPoolMergeExecutorService
+                    threadPoolMergeExecutorService,
+                    merge -> 0
                 )
             ) {
                 CountDownLatch mergeDoneLatch = new CountDownLatch(1);
@@ -576,7 +582,8 @@ public void testAutoIOThrottleForMergeTasksWhenSchedulerDisablesIt() throws Exce
             ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                 new ShardId("index", "_na_", 1),
                 indexSettings,
-                threadPoolMergeExecutorService
+                threadPoolMergeExecutorService,
+                merge -> 0
             )
         ) {
             threadPoolMergeScheduler.merge(mergeSource, randomFrom(MergeTrigger.values()));
@@ -605,7 +612,8 @@ public void testAutoIOThrottleForMergeTasks() throws Exception {
             ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                 new ShardId("index", "_na_", 1),
                 indexSettings,
-                threadPoolMergeExecutorService
+                threadPoolMergeExecutorService,
+                merge -> 0
             )
         ) {
             threadPoolMergeScheduler.merge(mergeSource, randomFrom(MergeTrigger.values()));
@@ -621,7 +629,8 @@ public void testAutoIOThrottleForMergeTasks() throws Exception {
             ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                 new ShardId("index", "_na_", 1),
                 indexSettings,
-                threadPoolMergeExecutorService
+                threadPoolMergeExecutorService,
+                merge -> 0
             )
         ) {
             // merge submitted upon closing
@@ -637,7 +646,8 @@ public void testAutoIOThrottleForMergeTasks() throws Exception {
             ThreadPoolMergeScheduler threadPoolMergeScheduler = new ThreadPoolMergeScheduler(
                 new ShardId("index", "_na_", 1),
                 indexSettings,
-                threadPoolMergeExecutorService
+                threadPoolMergeExecutorService,
+                merge -> 0
             )
         ) {
             // merge submitted upon closing
@@ -668,7 +678,7 @@ static class TestThreadPoolMergeScheduler extends ThreadPoolMergeScheduler {
             IndexSettings indexSettings,
             ThreadPoolMergeExecutorService threadPoolMergeExecutorService
         ) {
-            super(shardId, indexSettings, threadPoolMergeExecutorService);
+            super(shardId, indexSettings, threadPoolMergeExecutorService, merge -> 0);
         }
 
         @Override

Original file line number	Diff line number	Diff line change
`@@ -2870,7 +2870,7 @@ private final class EngineThreadPoolMergeScheduler extends ThreadPoolMergeSchedu`
`2870`	`2870`	`IndexSettings indexSettings,`
`2871`	`2871`	`ThreadPoolMergeExecutorService threadPoolMergeExecutorService`
`2872`	`2872`	`) {`
`2873`		`- super(shardId, indexSettings, threadPoolMergeExecutorService);`
	`2873`	`+ super(shardId, indexSettings, threadPoolMergeExecutorService, InternalEngine.this::estimateMergeBytes);`
`2874`	`2874`	`}`
`2875`	`2875`
`2876`	`2876`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -14,11 +14,12 @@`
`14`	`14`	`import org.elasticsearch.common.util.concurrent.ConcurrentCollections;`
`15`	`15`	`import org.elasticsearch.core.Nullable;`
`16`	`16`	`import org.elasticsearch.index.engine.ThreadPoolMergeScheduler.MergeTask;`
`17`		`-import org.elasticsearch.index.merge.OnGoingMerge;`
`18`	`17`	`import org.elasticsearch.threadpool.ThreadPool;`
`19`	`18`
`20`	`19`	`import java.util.Comparator;`
	`20`	`+import java.util.List;`
`21`	`21`	`import java.util.Set;`
	`22`	`+import java.util.concurrent.CopyOnWriteArrayList;`
`22`	`23`	`import java.util.concurrent.ExecutorService;`
`23`	`24`	`import java.util.concurrent.PriorityBlockingQueue;`
`24`	`25`	`import java.util.concurrent.RejectedExecutionException;`
`@@ -74,7 +75,7 @@ public class ThreadPoolMergeExecutorService {`
`74`	`75`	`private final int concurrentMergesFloorLimitForThrottling;`
`75`	`76`	`private final int concurrentMergesCeilLimitForThrottling;`
`76`	`77`
`77`		`- private volatile MergeEventConsumer mergeEventConsumer;`
	`78`	`+ private final List<MergeEventListener> mergeEventListeners = new CopyOnWriteArrayList<>();`
`78`	`79`
`79`	`80`	`public static @Nullable ThreadPoolMergeExecutorService maybeCreateThreadPoolMergeExecutorService(`
`80`	`81`	`ThreadPool threadPool,`
`@@ -130,13 +131,18 @@ boolean submitMergeTask(MergeTask mergeTask) {`
`130`	`131`	`);`
`131`	`132`	`}`
`132`	`133`	`// then enqueue the merge task proper`
`133`		`- queuedMergeTasks.add(mergeTask);`
	`134`	`+ enqueueMergeTask(mergeTask);`
`134`	`135`	`return true;`
`135`	`136`	`}`
`136`	`137`	`}`
`137`	`138`
`138`	`139`	`void reEnqueueBackloggedMergeTask(MergeTask mergeTask) {`
	`140`	`+ enqueueMergeTask(mergeTask);`
	`141`	`+ }`
	`142`	`+`
	`143`	`+ private void enqueueMergeTask(MergeTask mergeTask) {`
`139`	`144`	`queuedMergeTasks.add(mergeTask);`
	`145`	`+ mergeEventListeners.forEach(l -> l.onMergeQueued(mergeTask.getOnGoingMerge(), mergeTask.getEstimateMergeMemoryBytes()));`
`140`	`146`	`}`
`141`	`147`
`142`	`148`	`public boolean allDone() {`
`@@ -204,6 +210,7 @@ private void runMergeTask(MergeTask mergeTask) {`
`204`	`210`	`if (mergeTask.supportsIOThrottling()) {`
`205`	`211`	`ioThrottledMergeTasksCount.decrementAndGet();`
`206`	`212`	`}`
	`213`	`+ mergeEventListeners.forEach(l -> l.onMergeCompleted(mergeTask.getOnGoingMerge()));`
`207`	`214`	`}`
`208`	`215`	`}`
`209`	`216`
`@@ -216,6 +223,7 @@ private void abortMergeTask(MergeTask mergeTask) {`
`216`	`223`	`if (mergeTask.supportsIOThrottling()) {`
`217`	`224`	`ioThrottledMergeTasksCount.decrementAndGet();`
`218`	`225`	`}`
	`226`	`+ mergeEventListeners.forEach(l -> l.onMergeAborted(mergeTask.getOnGoingMerge()));`
`219`	`227`	`}`
`220`	`228`	`}`
`221`	`229`
`@@ -281,15 +289,8 @@ public boolean usingMaxTargetIORateBytesPerSec() {`
`281`	`289`	`return MAX_IO_RATE.getBytes() == targetIORateBytesPerSec.get();`
`282`	`290`	`}`
`283`	`291`
`284`		`- public void registerMergeEventConsumer(MergeEventConsumer consumer) {`
`285`		`- assert this.mergeEventConsumer == null;`
`286`		`- this.mergeEventConsumer = consumer;`
`287`		`- }`
`288`		`-`
`289`		`- public interface MergeEventConsumer {`
`290`		`- void onMergeQueued(OnGoingMerge merge, long estimateMergeMemoryBytes);`
`291`		`-`
`292`		`- void onMergeCompleted(OnGoingMerge merge);`
	`292`	`+ public void registerMergeEventListener(MergeEventListener consumer) {`
	`293`	`+ mergeEventListeners.add(consumer);`
`293`	`294`	`}`
`294`	`295`
`295`	`296`	`// exposed for tests`
Original file line number	Diff line number	Diff line change
`@@ -660,6 +660,8 @@ public void testMergeTasksExecuteInSizeOrder() {`
`660`	`660`	`}`
`661`	`661`	`}`
`662`	`662`
	`663`	`+ // todo: testMergeEventListeners`
	`664`	`+`
`663`	`665`	`static ThreadPoolMergeExecutorService getThreadPoolMergeExecutorService(ThreadPool threadPool) {`
`664`	`666`	`ThreadPoolMergeExecutorService threadPoolMergeExecutorService = ThreadPoolMergeExecutorService`
`665`	`667`	`.maybeCreateThreadPoolMergeExecutorService(`