Closing not quite there....

albertzaharovits · albertzaharovits · commit 9d2d29da0743 · 2025-02-18T20:56:15.000+02:00
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeQueue.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeQueue.java
@@ -16,14 +16,17 @@
 import org.elasticsearch.index.engine.ThreadPoolMergeScheduler.MergeTask;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
-public class ThreadPoolMergeQueue {
+public class ThreadPoolMergeQueue implements Closeable {
     /**
      * Floor for IO write rate limit of individual merge tasks (we will never go any lower than this)
      */
@@ -48,6 +51,7 @@ public class ThreadPoolMergeQueue {
     private final AtomicLong targetIORateBytesPerSec = new AtomicLong(START_IO_RATE.getBytes());
     private final ExecutorService executorService;
     private final int maxConcurrentMerges;
+    private final AtomicBoolean closed = new AtomicBoolean();
 
     public static @Nullable ThreadPoolMergeQueue maybeCreateThreadPoolMergeQueue(ThreadPool threadPool, Settings settings) {
         if (ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.get(settings)) {
@@ -63,51 +67,66 @@ private ThreadPoolMergeQueue(ThreadPool threadPool) {
     }
 
     void submitMergeTask(MergeTask mergeTask) {
-        enqueueMergeTask(mergeTask);
-        if (mergeTask.supportsIOThrottling()) {
-            // count submitted merge tasks that support IO auto throttling, and maybe adjust IO rate for all
-            maybeUpdateIORateBytesPerSec(submittedIOThrottledMergeTasksCount.incrementAndGet());
+        assert mergeTask.isRunning() == false;
+        assert mergeTask.isOnGoingMergeAborted() == false;
+        if (closed.get() || enqueueMergeTaskExecution() == false) {
+            mergeTask.abortOnGoingMerge();
+        } else {
+            if (mergeTask.supportsIOThrottling()) {
+                // count submitted merge tasks that support IO auto throttling, and maybe adjust IO rate for all
+                maybeUpdateIORateBytesPerSec(submittedIOThrottledMergeTasksCount.incrementAndGet());
+            }
+            enqueueMergeTask(mergeTask);
         }
-        executeSmallestMergeTask();
     }
 
     void enqueueMergeTask(MergeTask mergeTask) {
         queuedMergeTasks.add(mergeTask);
     }
 
-    private void executeSmallestMergeTask() {
-        final AtomicReference<MergeTask> smallestMergeTask = new AtomicReference<>();
+    private boolean enqueueMergeTaskExecution() {
         try {
             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
-                while (true) {
-                    try {
-                        // will block if there are backlogged merges until they're enqueued again
-                        smallestMergeTask.set(queuedMergeTasks.take());
-                    } catch (InterruptedException e) {
-                        Thread.currentThread().interrupt();
-                        return;
+                try {
+                    while (true) {
+                        try {
+                            // will block if there are backlogged merges until they're enqueued again
+                            MergeTask smallestMergeTask = queuedMergeTasks.take();
+                            // the merge task's scheduler might backlog rather than execute the task
+                            // it's then the duty of the said merge scheduler to re-enqueue the backlogged merge task
+                            if (smallestMergeTask.runNowOrBacklog()) {
+                                runMergeTask(smallestMergeTask);
+                                // one runnable one merge task
+                                return;
+                            }
+                            // the merge task is backlogged by the merge scheduler, try to get the next smallest one
+                        } catch (InterruptedException e) {
+                            interrupted = true;
+                            // this runnable must always run exactly a single merge task
+                        }
                     }
-                    // the merge task's scheduler might backlog rather than execute the task
-                    // it's then the duty of the said merge scheduler to re-enqueue the backlogged merge task
-                    if (smallestMergeTask.get().runNowOrBacklog()) {
-                        runMergeTask(smallestMergeTask.get());
-                        // one runnable one merge task
-                        return;
+                } finally {
+                    if (interrupted) {
+                        Thread.currentThread().interrupt();
                     }
-                    // the merge task is backlogged by the merge scheduler,
                 }
             });
-        } catch (Exception e) {
-            if (smallestMergeTask.get() != null) {
-                smallestMergeTask.get().onRejection(e);
-            }
+            return true;
+        } catch (RejectedExecutionException e) {
+            // cannot execute merges because the executor is shutting down
+            return false;
         }
     }
 
     private void runMergeTask(MergeTask mergeTask) {
         assert mergeTask.isRunning() == false;
+        if (closed.get()) {
+            mergeTask.abortOnGoingMerge();
+            return;
+        }
         boolean added = currentlyRunningMergeTasks.add(mergeTask);
         assert added : "starting merge task [" + mergeTask + "] registered as already running";
         try {
@@ -183,4 +202,9 @@ public boolean isEmpty() {
         assert isEmpty == false || submittedIOThrottledMergeTasksCount.get() == 0L;
         return isEmpty;
     }
+
+    @Override
+    public void close() throws IOException {
+        closed.set(true);
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java b/server/src/main/java/org/elasticsearch/index/engine/ThreadPoolMergeScheduler.java
@@ -56,6 +56,7 @@ public class ThreadPoolMergeScheduler extends MergeScheduler implements Elastics
     // how many {@link MergeTask}s have kicked off (this is used to name them).
     private final AtomicLong mergeTaskCount = new AtomicLong();
     private final AtomicLong mergeTaskDoneCount = new AtomicLong();
+    private final AtomicBoolean closed = new AtomicBoolean();
 
     public ThreadPoolMergeScheduler(ShardId shardId, IndexSettings indexSettings, ThreadPoolMergeQueue threadPoolMergeQueue) {
         this.shardId = shardId;
@@ -88,17 +89,28 @@ public void refreshConfig() {
         // if maxMergeCount changed, maybe we need to toggle merge task throttling
         checkMergeTaskThrottling();
         // if maxThreadCount changed, maybe some backlogged merges are now allowed to run
-        enqueueBacklogged();
+        enqueueBackloggedTasks();
     }
 
     @Override
     public void merge(MergeSource mergeSource, MergeTrigger trigger) throws IOException {
+        if (closed.get()) {
+            // avoid pulling from the merge source when closing
+            return;
+        }
         MergePolicy.OneMerge merge = mergeSource.getNextMerge();
         if (merge != null) {
             submitNewMergeTask(mergeSource, merge, trigger);
         }
     }
 
+    @Override
+    public MergeScheduler clone() {
+        // Lucene IW makes a clone internally but since we hold on to this instance
+        // the clone will just be the identity.
+        return this;
+    }
+
     /**
      * A callback allowing for custom logic before an actual merge starts.
      */
@@ -109,17 +121,20 @@ protected void beforeMerge(OnGoingMerge merge) {}
      */
     protected void afterMerge(OnGoingMerge merge) {}
 
+    /**
+     * A callback that's invoked when indexing should throttle down indexing in order to let merging to catch up.
+     */
     protected void enableIndexingThrottling(int numRunningMerges, int numQueuedMerges, int configuredMaxMergeCount) {}
 
+    /**
+     * A callback that's invoked when indexing should un-throttle because merging caught up.
+     * This is invoked sometime after {@link #enableIndexingThrottling(int, int, int)} was invoked in the first place.
+     */
     protected void disableIndexingThrottling(int numRunningMerges, int numQueuedMerges, int configuredMaxMergeCount) {}
 
-    @Override
-    public MergeScheduler clone() {
-        // Lucene IW makes a clone internally but since we hold on to this instance
-        // the clone will just be the identity.
-        return this;
-    }
-
+    /**
+     * A callback for exceptions thrown while merging.
+     */
     protected void handleMergeException(Throwable t) {
         throw new MergePolicy.MergeException(t);
     }
@@ -162,13 +177,15 @@ private MergeTask newMergeTask(MergeSource mergeSource, MergePolicy.OneMerge mer
     // synchronized so that {@code #currentlyRunningMergeTasks} and {@code #backloggedMergeTasks} are modified atomically
     private synchronized boolean runNowOrBacklog(MergeTask mergeTask) {
         assert mergeTask.isRunning() == false;
-        if (currentlyRunningMergeTasks.size() >= config.getMaxThreadCount()) {
-            backloggedMergeTasks.add(mergeTask);
-            return false;
-        } else {
+        assert mergeTask.isOnGoingMergeAborted() == false;
+        // if the merge scheduler is closed it will abort all merges before they start running
+        if (closed.get() || 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;
+        } else {
+            backloggedMergeTasks.add(mergeTask);
+            return false;
         }
     }
 
@@ -178,15 +195,20 @@ private void mergeDone(MergeTask mergeTask) {
             boolean removed = currentlyRunningMergeTasks.remove(mergeTask.onGoingMerge.getMerge()) != null;
             assert removed : "completed merge task [" + mergeTask + "] not registered as running";
             // when one merge is done, maybe a backlogged one can now execute
-            enqueueBacklogged();
+            enqueueBackloggedTasks();
+            // when closing, we wait for all running merges to finish
+            if (currentlyRunningMergeTasks.isEmpty()) {
+                notifyAll();
+            }
         }
         mergeTaskDoneCount.incrementAndGet();
         checkMergeTaskThrottling();
     }
 
-    private synchronized void enqueueBacklogged() {
+    private synchronized void enqueueBackloggedTasks() {
         int maxBackloggedTasksToEnqueue = config.getMaxThreadCount() - currentlyRunningMergeTasks.size();
-        while (maxBackloggedTasksToEnqueue-- > 0) {
+        // enqueue all backlogged tasks when closing, let the queue deal with them
+        while (closed.get() || maxBackloggedTasksToEnqueue-- > 0) {
             MergeTask backloggedMergeTask = backloggedMergeTasks.poll();
             if (backloggedMergeTask == null) {
                 break;
@@ -271,6 +293,11 @@ public boolean isRunning() {
 
         @Override
         public void run() {
+            assert isOnGoingMergeAborted() == false;
+            if (ThreadPoolMergeScheduler.this.closed.get()) {
+                abortOnGoingMerge();
+                return;
+            }
             if (mergeStartTimeNS.compareAndSet(0L, System.nanoTime()) == false) {
                 throw new IllegalStateException("Cannot run the same merge task multiple times");
             }
@@ -302,6 +329,7 @@ public void run() {
                 if (t instanceof MergePolicy.MergeAbortedException) {
                     // OK to ignore. This is what Lucene's ConcurrentMergeScheduler does
                 } else {
+                    // sometimes this might double-abort a merge, but that's OK
                     abortOnGoingMerge();
                     handleMergeException(t);
                 }
@@ -317,32 +345,27 @@ public void run() {
                         mergeTracking.mergeFinished(onGoingMerge.getMerge(), onGoingMerge, tookMS);
                     } finally {
                         mergeDone(this);
-                        // 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");
+                        if (ThreadPoolMergeScheduler.this.closed.get() == 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);
                             }
-                            // ignore exception, we expect the IW failure to be logged elsewhere
-                        }
-                        if (nextMerge != null) {
-                            submitNewMergeTask(mergeSource, nextMerge, MergeTrigger.MERGE_FINISHED);
                         }
                     }
                 }
             }
         }
 
-        public void onRejection(Exception e) {
-            if (verbose()) {
-                message(String.format(Locale.ROOT, "merge task [%s] rejected by thread pool, aborting", onGoingMerge.getId()));
-            }
-            abortOnGoingMerge();
-        }
-
-        private void abortOnGoingMerge() {
+        void abortOnGoingMerge() {
             // This would interrupt an IndexWriter if it were actually performing the merge. We just set this here because it seems
             // appropriate as we are not going to move forward with the merge.
             onGoingMerge.getMerge().setAborted();
@@ -351,6 +374,10 @@ private void abortOnGoingMerge() {
             mergeSource.onMergeFinished(onGoingMerge.getMerge());
         }
 
+        boolean isOnGoingMergeAborted() {
+            return onGoingMerge.getMerge().isAborted();
+        }
+
         @Override
         public String toString() {
             return name;
@@ -375,6 +402,35 @@ protected void message(String message) {
         super.message(message);
     }
 
+    @Override
+    public void close() throws IOException {
+        super.close();
+        closed.set(true);
+        boolean interrupted = false;
+        try {
+            synchronized (this) {
+                // enqueue all backlogged merge tasks, the merge queue assumes that backlogged tasks are always re-enqueued
+                enqueueBackloggedTasks();
+                // supercharge running merges
+                currentlyRunningMergeTasks.values().forEach(runningTask -> runningTask.setIORateLimit(Double.POSITIVE_INFINITY));
+                // wait until all running merges are done
+                while (currentlyRunningMergeTasks.isEmpty() == false) {
+                    try {
+                        // wait with a timeout, just to cover for something that failed to notify
+                        wait(1000);
+                    } catch (InterruptedException e) {
+                        // ignore interruption, we will retry until all currently running merge tasks are done
+                        interrupted = true;
+                    }
+                }
+            }
+        } finally {
+            if (interrupted) {
+                Thread.currentThread().interrupt();
+            }
+        }
+    }
+
     private static double nsToSec(long ns) {
         return ns / (double) TimeUnit.SECONDS.toNanos(1);
     }
diff --git a/server/src/main/java/org/elasticsearch/threadpool/DefaultBuiltInExecutorBuilders.java b/server/src/main/java/org/elasticsearch/threadpool/DefaultBuiltInExecutorBuilders.java
@@ -149,7 +149,7 @@ public Map<String, ExecutorBuilder> getBuilders(Settings settings, int allocated
         if (ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.get(settings)) {
             result.put(
                 ThreadPool.Names.MERGE,
-                new ScalingExecutorBuilder(ThreadPool.Names.MERGE, 1, allocatedProcessors, TimeValue.timeValueMinutes(5), false)
+                new ScalingExecutorBuilder(ThreadPool.Names.MERGE, 1, allocatedProcessors, TimeValue.timeValueMinutes(5), true)
             );
         }
         result.put(

Original file line number	Diff line number	Diff line change
`@@ -149,7 +149,7 @@ public Map<String, ExecutorBuilder> getBuilders(Settings settings, int allocated`
`149`	`149`	`if (ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.get(settings)) {`
`150`	`150`	`result.put(`
`151`	`151`	`ThreadPool.Names.MERGE,`
`152`		`- new ScalingExecutorBuilder(ThreadPool.Names.MERGE, 1, allocatedProcessors, TimeValue.timeValueMinutes(5), false)`
	`152`	`+ new ScalingExecutorBuilder(ThreadPool.Names.MERGE, 1, allocatedProcessors, TimeValue.timeValueMinutes(5), true)`
`153`	`153`	`);`
`154`	`154`	`}`
`155`	`155`	`result.put(`