Handle aborted merge tasks at any point

albertzaharovits · albertzaharovits · commit b2e65a2e5c97 · 2025-02-23T12:42:29.000+02:00
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorService.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeExecutorService.java
@@ -62,20 +62,21 @@ private ThreadPoolMergeExecutorService(ThreadPool threadPool) {
         this.maxConcurrentMerges = threadPool.info(ThreadPool.Names.MERGE).getMax();
     }
 
-    void submitMergeTask(MergeTask mergeTask) {
+    boolean submitMergeTask(MergeTask mergeTask) {
         assert mergeTask.isRunning() == false;
-        assert mergeTask.isOnGoingMergeAborted() == false;
         // first enqueue the runnable that runs exactly one merge task (the smallest it can find)
         if (enqueueMergeTaskExecution() == false) {
-            // if the threadpool cannot run the merge, just abort it
+            // if the thread pool cannot run the merge, just abort it
             mergeTask.abortOnGoingMerge();
+            return false;
         } else {
             if (mergeTask.supportsIOThrottling()) {
                 // count enqueued merge tasks that support IO auto throttling, and maybe adjust IO rate for all
                 maybeUpdateIORateBytesPerSec(submittedIOThrottledMergeTasksCount.incrementAndGet());
             }
             // then enqueue the merge task proper
             enqueueMergeTask(mergeTask);
+            return true;
         }
     }
 
@@ -92,17 +93,19 @@ private boolean enqueueMergeTaskExecution() {
             executorService.execute(() -> {
                 boolean interrupted = false;
                 // one such runnable always executes a SINGLE merge task from the queue
-                // this is important for merge queue statistics, i.e. the executor's queue size equals the merge tasks' queue size
+                // this is important for merge queue statistics, i.e. the executor's queue size represents the current amount of merges
+                MergeTask smallestMergeTask = null;
                 try {
                     while (true) {
                         try {
                             // will block if there are backlogged merges until they're enqueued again
-                            MergeTask smallestMergeTask = queuedMergeTasks.take();
+                            smallestMergeTask = queuedMergeTasks.take();
                             // let the task's scheduler decide if it can actually run the merge task now
                             if (smallestMergeTask.runNowOrBacklog()) {
                                 runMergeTask(smallestMergeTask);
                                 break;
                             }
+                            smallestMergeTask = null;
                             // the merge task is backlogged by the merge scheduler, try to get the next smallest one
                             // it's then the duty of the said merge scheduler to re-enqueue the backlogged merge task when it can be run
                         } catch (InterruptedException e) {
@@ -111,6 +114,9 @@ private boolean enqueueMergeTaskExecution() {
                         }
                     }
                 } finally {
+                    if (smallestMergeTask != null && smallestMergeTask.supportsIOThrottling()) {
+                        submittedIOThrottledMergeTasksCount.decrementAndGet();
+                    }
                     if (interrupted) {
                         Thread.currentThread().interrupt();
                     }
@@ -126,25 +132,16 @@ private boolean enqueueMergeTaskExecution() {
 
     private void runMergeTask(MergeTask mergeTask) {
         assert mergeTask.isRunning() == false;
+        boolean added = currentlyRunningMergeTasks.add(mergeTask);
+        assert added : "starting merge task [" + mergeTask + "] registered as already running";
         try {
-            if (mergeTask.isOnGoingMergeAborted()) {
-                return;
-            }
-            boolean added = currentlyRunningMergeTasks.add(mergeTask);
-            assert added : "starting merge task [" + mergeTask + "] registered as already running";
-            try {
-                if (mergeTask.supportsIOThrottling()) {
-                    mergeTask.setIORateLimit(ByteSizeValue.ofBytes(targetIORateBytesPerSec.get()).getMbFrac());
-                }
-                mergeTask.run();
-            } finally {
-                boolean removed = currentlyRunningMergeTasks.remove(mergeTask);
-                assert removed : "completed merge task [" + mergeTask + "] not registered as running";
-            }
-        } finally {
             if (mergeTask.supportsIOThrottling()) {
-                submittedIOThrottledMergeTasksCount.decrementAndGet();
+                mergeTask.setIORateLimit(ByteSizeValue.ofBytes(targetIORateBytesPerSec.get()).getMbFrac());
             }
+            mergeTask.run();
+        } finally {
+            boolean removed = currentlyRunningMergeTasks.remove(mergeTask);
+            assert removed : "completed merge task [" + mergeTask + "] not registered as running";
         }
     }
 
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java
@@ -19,6 +19,7 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.RateLimitedIndexOutput;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.core.TimeValue;
@@ -99,12 +100,20 @@ public void refreshConfig() {
     }
 
     @Override
-    public void merge(MergeSource mergeSource, MergeTrigger trigger) throws IOException {
+    public void merge(MergeSource mergeSource, MergeTrigger trigger) {
         if (closed) {
             // avoid pulling from the merge source when closing
             return;
         }
-        MergePolicy.OneMerge merge = mergeSource.getNextMerge();
+        MergePolicy.OneMerge merge = null;
+        try {
+            merge = mergeSource.getNextMerge();
+        } catch (IllegalStateException e) {
+            if (verbose()) {
+                message("merge task poll failed, likely that index writer is failed");
+            }
+            // ignore exception, we expect the IW failure to be logged elsewhere
+        }
         if (merge != null) {
             submitNewMergeTask(mergeSource, merge, trigger);
         }
@@ -145,10 +154,13 @@ protected void handleMergeException(Throwable t) {
         throw new MergePolicy.MergeException(t);
     }
 
-    private void submitNewMergeTask(MergeSource mergeSource, MergePolicy.OneMerge merge, MergeTrigger mergeTrigger) {
-        MergeTask mergeTask = newMergeTask(mergeSource, merge, mergeTrigger);
-        threadPoolMergeExecutorService.submitMergeTask(mergeTask);
-        checkMergeTaskThrottling();
+    private boolean submitNewMergeTask(MergeSource mergeSource, MergePolicy.OneMerge merge, MergeTrigger mergeTrigger) {
+        try {
+            MergeTask mergeTask = newMergeTask(mergeSource, merge, mergeTrigger);
+            return threadPoolMergeExecutorService.submitMergeTask(mergeTask);
+        } finally {
+            checkMergeTaskThrottling();
+        }
     }
 
     private void checkMergeTaskThrottling() {
@@ -179,16 +191,15 @@ private MergeTask newMergeTask(MergeSource mergeSource, MergePolicy.OneMerge mer
         );
     }
 
-    // synchronized so that {@code #currentlyRunningMergeTasks} and {@code #backloggedMergeTasks} are modified atomically
+    // synchronized so that {@code #closed}, {@code #currentlyRunningMergeTasks} and {@code #backloggedMergeTasks} are modified atomically
     private synchronized boolean runNowOrBacklog(MergeTask mergeTask) {
         assert mergeTask.isRunning() == false;
-        assert mergeTask.isOnGoingMergeAborted() == false;
         if (closed) {
-            // Do not backlog tasks when closing the merge scheduler, instead abort them.
-            // Aborted task are not actually executed.
+            // Do not backlog or execute tasks when closing the merge scheduler, instead abort them.
             mergeTask.abortOnGoingMerge();
-            return true;
-        } else if (currentlyRunningMergeTasks.size() < config.getMaxThreadCount()) {
+            throw new ElasticsearchException("merge task aborted because scheduler is shutting down");
+        }
+        if (currentlyRunningMergeTasks.size() < config.getMaxThreadCount()) {
             boolean added = currentlyRunningMergeTasks.put(mergeTask.onGoingMerge.getMerge(), mergeTask) == null;
             assert added : "starting merge task [" + mergeTask + "] registered as already running";
             return true;
@@ -225,7 +236,7 @@ private synchronized void enqueueBackloggedTasks() {
             if (backloggedMergeTask == null) {
                 break;
             }
-            // no need to abort merge tasks now, they will be aborted when the scheduler tries to run them
+            // no need to abort merge tasks now, they will be aborted on the spot when the scheduler gets to run them
             threadPoolMergeExecutorService.enqueueMergeTask(backloggedMergeTask);
         }
     }
@@ -307,7 +318,6 @@ public boolean isRunning() {
         @Override
         public void run() {
             assert isRunning() == false;
-            assert isOnGoingMergeAborted() == false;
             assert ThreadPoolMergeScheduler.this.currentlyRunningMergeTasks.containsKey(onGoingMerge.getMerge())
                 : "runNowOrBacklog must be invoked before actually running the merge task";
             try {
@@ -360,21 +370,8 @@ public void run() {
                     message(String.format(Locale.ROOT, "merge task %s end", this));
                 }
                 mergeDone(this);
-                if (ThreadPoolMergeScheduler.this.closed == false) {
-                    // kick-off next merge, if any
-                    MergePolicy.OneMerge nextMerge = null;
-                    try {
-                        nextMerge = mergeSource.getNextMerge();
-                    } catch (IllegalStateException e) {
-                        if (verbose()) {
-                            message("merge task poll failed, likely that index writer is failed");
-                        }
-                        // ignore exception, we expect the IW failure to be logged elsewhere
-                    }
-                    if (nextMerge != null) {
-                        submitNewMergeTask(mergeSource, nextMerge, MergeTrigger.MERGE_FINISHED);
-                    }
-                }
+                // kick-off next merge, if any
+                merge(mergeSource, MergeTrigger.MERGE_FINISHED);
             }
         }
 
@@ -388,13 +385,9 @@ void abortOnGoingMerge() {
             doneMergeTaskCount.incrementAndGet();
         }
 
-        boolean isOnGoingMergeAborted() {
-            return onGoingMerge.getMerge().isAborted();
-        }
-
         @Override
         public String toString() {
-            return name + (isOnGoingMergeAborted() ? " (aborted)" : "");
+            return name + (onGoingMerge.getMerge().isAborted() ? " (aborted)" : "");
         }
     }
 
