Threadpool merge executor is aware of available disk space

docs/changelog/127613.yaml

@@ -0,0 +1,5 @@
+pr: 127613
+summary: Threadpool merge executor is aware of available disk space
+area: Engine
+type: feature
+issues: []

server/src/main/java/org/elasticsearch/common/settings/ClusterSettings.java

@@ -88,6 +88,7 @@
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.IndexingPressure;
 import org.elasticsearch.index.MergePolicyConfig;
+import org.elasticsearch.index.engine.ThreadPoolMergeExecutorService;
 import org.elasticsearch.index.engine.ThreadPoolMergeScheduler;
 import org.elasticsearch.index.shard.IndexingStatsSettings;
 import org.elasticsearch.indices.IndexingMemoryController;
@@ -629,6 +630,9 @@ public void apply(Settings value, Settings current, Settings previous) {
         MergePolicyConfig.DEFAULT_MAX_MERGED_SEGMENT_SETTING,
         MergePolicyConfig.DEFAULT_MAX_TIME_BASED_MERGED_SEGMENT_SETTING,
         ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING,
+        ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_HIGH_WATERMARK_SETTING,
+        ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_HIGH_MAX_HEADROOM_SETTING,
+        ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_CHECK_INTERVAL_SETTING,
         TransportService.ENABLE_STACK_OVERFLOW_AVOIDANCE,
         DataStreamGlobalRetentionSettings.DATA_STREAMS_DEFAULT_RETENTION_SETTING,
         DataStreamGlobalRetentionSettings.DATA_STREAMS_MAX_RETENTION_SETTING,

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

@@ -55,7 +55,7 @@ public class ThreadPoolMergeScheduler extends MergeScheduler implements Elastics
     private final ThreadPoolMergeExecutorService threadPoolMergeExecutorService;
     private final PriorityQueue<MergeTask> backloggedMergeTasks = new PriorityQueue<>(
         16,
-        Comparator.comparingLong(MergeTask::estimatedMergeSize)
+        Comparator.comparingLong(MergeTask::estimatedRemainingMergeSize)
     );
     private final Map<MergePolicy.OneMerge, MergeTask> runningMergeTasks = new HashMap<>();
     // set when incoming merges should be throttled (i.e. restrict the indexing rate)
@@ -266,7 +266,7 @@ private void checkMergeTaskThrottling() {
     // exposed for tests
     // synchronized so that {@code #closed}, {@code #runningMergeTasks} and {@code #backloggedMergeTasks} are modified atomically
     synchronized Schedule schedule(MergeTask mergeTask) {
-        assert mergeTask.isRunning() == false;
+        assert mergeTask.hasStartedRunning() == false;
         if (closed) {
             // do not run or backlog tasks when closing the merge scheduler, instead abort them
             return Schedule.ABORT;
@@ -280,6 +280,7 @@ synchronized Schedule schedule(MergeTask mergeTask) {
             assert added : "starting merge task [" + mergeTask + "] registered as already running";
             return Schedule.RUN;
         } else {
+            assert mergeTask.hasStartedRunning() == false;
             backloggedMergeTasks.add(mergeTask);
             return Schedule.BACKLOG;
         }
@@ -403,8 +404,14 @@ public void setIORateLimit(long ioRateLimitBytesPerSec) {
             this.rateLimiter.setMBPerSec(ByteSizeValue.ofBytes(ioRateLimitBytesPerSec).getMbFrac());
         }
 
-        public boolean isRunning() {
-            return mergeStartTimeNS.get() > 0L;
+        /**
+         * Returns {@code true} if this task is currently running, or was run in the past.
+         * An aborted task (see {@link #abort()}) is considered as NOT run.
+         */
+        public boolean hasStartedRunning() {
+            boolean isRunning = mergeStartTimeNS.get() > 0L;
+            assert isRunning != false || rateLimiter.getTotalBytesWritten() == 0L;
+            return isRunning;
         }
 
         /**
@@ -415,7 +422,7 @@ public boolean isRunning() {
          */
         @Override
         public void run() {
-            assert isRunning() == false;
+            assert hasStartedRunning() == false;
             assert ThreadPoolMergeScheduler.this.runningMergeTasks.containsKey(onGoingMerge.getMerge())
                 : "runNowOrBacklog must be invoked before actually running the merge task";
             try {
@@ -480,7 +487,7 @@ public void run() {
          * (by the {@link org.apache.lucene.index.IndexWriter}) to any subsequent merges.
          */
         void abort() {
-            assert isRunning() == false;
+            assert hasStartedRunning() == false;
             assert ThreadPoolMergeScheduler.this.runningMergeTasks.containsKey(onGoingMerge.getMerge()) == false
                 : "cannot abort a merge task that's already running";
             if (verbose()) {
@@ -509,10 +516,17 @@ void abort() {
             }
         }
 
-        long estimatedMergeSize() {
+        /**
+         * Before the merge task started running, this returns the estimated required disk space for the merge to complete
+         * (i.e. the estimated disk space size of the resulting segment following the merge).
+         * While the merge is running, the returned estimation is updated to take into account the data that's already been written.
+         * After the merge completes, the estimation returned here should ideally be close to "0".
+         */
+        long estimatedRemainingMergeSize() {
             // TODO is it possible that `estimatedMergeBytes` be `0` for correctly initialize merges,
             // or is it always the case that if `estimatedMergeBytes` is `0` that means that the merge has not yet been initialized?
-            return onGoingMerge.getMerge().getStoreMergeInfo().estimatedMergeBytes();
+            long estimatedMergeSize = onGoingMerge.getMerge().getStoreMergeInfo().estimatedMergeBytes();
+            return Math.max(0L, estimatedMergeSize - rateLimiter.getTotalBytesWritten());
         }
 
         public long getMergeMemoryEstimateBytes() {

server/src/main/java/org/elasticsearch/indices/IndicesService.java

@@ -295,10 +295,6 @@ protected void doStart() {
     IndicesService(IndicesServiceBuilder builder) {
         this.settings = builder.settings;
         this.threadPool = builder.threadPool;
-        this.threadPoolMergeExecutorService = ThreadPoolMergeExecutorService.maybeCreateThreadPoolMergeExecutorService(
-            threadPool,
-            settings
-        );
         this.pluginsService = builder.pluginsService;
         this.nodeEnv = builder.nodeEnv;
         this.parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE)
@@ -321,6 +317,11 @@ protected void doStart() {
         this.bigArrays = builder.bigArrays;
         this.scriptService = builder.scriptService;
         this.clusterService = builder.clusterService;
+        this.threadPoolMergeExecutorService = ThreadPoolMergeExecutorService.maybeCreateThreadPoolMergeExecutorService(
+            threadPool,
+            clusterService.getClusterSettings(),
+            nodeEnv
+        );
         this.projectResolver = builder.projectResolver;
         this.client = builder.client;
         this.idFieldDataEnabled = INDICES_ID_FIELD_DATA_ENABLED_SETTING.get(clusterService.getSettings());
@@ -368,7 +369,8 @@ public void onRemoval(ShardId shardId, String fieldName, boolean wasEvicted, lon
                     indicesFieldDataCache,
                     cacheCleaner,
                     indicesRequestCache,
-                    indicesQueryCache
+                    indicesQueryCache,
+                    threadPoolMergeExecutorService
                 );
             } catch (IOException e) {
                 throw new UncheckedIOException(e);

server/src/test/java/org/elasticsearch/index/IndexModuleTests.java

@@ -194,13 +194,17 @@ public void setUp() throws Exception {
             emptyMap()
         );
         threadPool = new TestThreadPool("test");
-        threadPoolMergeExecutorService = ThreadPoolMergeExecutorService.maybeCreateThreadPoolMergeExecutorService(threadPool, settings);
         circuitBreakerService = new NoneCircuitBreakerService();
         PageCacheRecycler pageCacheRecycler = new PageCacheRecycler(settings);
         bigArrays = new BigArrays(pageCacheRecycler, circuitBreakerService, CircuitBreaker.REQUEST);
         scriptService = new ScriptService(settings, Collections.emptyMap(), Collections.emptyMap(), () -> 1L);
-        clusterService = ClusterServiceUtils.createClusterService(threadPool);
+        clusterService = ClusterServiceUtils.createClusterService(threadPool, ClusterSettings.createBuiltInClusterSettings(settings));
         nodeEnvironment = new NodeEnvironment(settings, environment);
+        threadPoolMergeExecutorService = ThreadPoolMergeExecutorService.maybeCreateThreadPoolMergeExecutorService(
+            threadPool,
+            clusterService.getClusterSettings(),
+            nodeEnvironment
+        );
         mapperRegistry = new IndicesModule(Collections.emptyList()).getMapperRegistry();
         indexNameExpressionResolver = TestIndexNameExpressionResolver.newInstance(threadPool.getThreadContext());
     }