testBackloggedMergeTasksAreAllExecutedExactlyOnce

Author: albertzaharovits

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

@@ -36,17 +36,30 @@ public class ThreadPoolMergeExecutorService {
      * Initial value for IO write rate limit of individual merge tasks when doAutoIOThrottle is true
      */
     private static final ByteSizeValue START_IO_RATE = ByteSizeValue.ofMb(20L);
+    /**
+     * Total number of submitted merge tasks that support IO auto throttling and that have not yet been executed.
+     * This includes merge tasks that are currently running and that are backlogged (by their respective merge schedulers).
+     */
     private final AtomicInteger currentlySubmittedIOThrottledMergeTasksCount = new AtomicInteger();
+    /**
+     * The merge tasks that are waiting execution. This does NOT include backlogged or currently executing merge tasks.
+     * For instance, this can be empty while there are backlogged merge tasks awaiting re-enqueuing.
+     */
     private final PriorityBlockingQueue<MergeTask> queuedMergeTasks = new PriorityBlockingQueue<>();
-    // the set of all merge tasks currently being executed by merge threads from the pool,
-    // in order to be able to update the IO throttle rate of merge tasks also after they have started (while executing)
+    /**
+     * The set of all merge tasks currently being executed by merge threads from the pool.
+     * These are tracked notably in order to be able to update their disk IO throttle rate, after they have started, while executing.
+     */
     private final Set<MergeTask> currentlyRunningMergeTasks = ConcurrentCollections.newConcurrentSet();
     /**
      * Current IO write throttle rate, in bytes per sec, that's in effect for all currently running merge tasks,
      * across all {@link ThreadPoolMergeScheduler}s that use this instance of the queue.
      */
     private final AtomicLong targetIORateBytesPerSec = new AtomicLong(START_IO_RATE.getBytes());
     private final ExecutorService executorService;
+    /**
+     * The maximum number of concurrently running merges, given the number of threads in the pool.
+     */
     private final int maxConcurrentMerges;
 
     public static @Nullable ThreadPoolMergeExecutorService maybeCreateThreadPoolMergeExecutorService(
@@ -202,6 +215,11 @@ private static long newTargetIORateBytesPerSec(long currentTargetIORateBytesPerS
         return ByteSizeValue.ofBytes(targetIORateBytesPerSec.get()).getMbFrac();
     }
 
+    // exposed for tests
+    int getMaxConcurrentMerges() {
+        return maxConcurrentMerges;
+    }
+
     public boolean allDone() {
         return queuedMergeTasks.isEmpty()
             && currentlyRunningMergeTasks.isEmpty()

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

@@ -311,7 +311,7 @@ boolean runNowOrBacklog() {
         @Override
         public int compareTo(MergeTask other) {
             // sort smaller merges first, so they are executed before larger ones
-            return Long.compare(onGoingMerge.getMerge().estimatedMergeBytes, other.onGoingMerge.getMerge().estimatedMergeBytes);
+            return Long.compare(estimatedMergeSize(), other.estimatedMergeSize());
         }
 
         public boolean supportsIOThrottling() {
@@ -396,6 +396,10 @@ void abortOnGoingMerge() {
             doMerge(mergeSource, onGoingMerge.getMerge());
         }
 
+        long estimatedMergeSize() {
+            return onGoingMerge.getMerge().estimatedMergeBytes;
+        }
+
         @Override
         public String toString() {
             return name + (onGoingMerge.getMerge().isAborted() ? " (aborted)" : "");

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

@@ -10,14 +10,22 @@
 package org.elasticsearch.index.engine;
 
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.common.util.concurrent.DeterministicTaskQueue;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.IndexSettingsModule;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.junit.Before;
 
+import java.util.Collection;
+import java.util.concurrent.CountDownLatch;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.mockito.ArgumentMatchers.anyDouble;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -37,7 +45,7 @@ public void setUpThreadPool() {
     }
 
     public void testMergeTasksAreAbortedWhenThreadPoolIsShutdown() {
-        final TestThreadPool testThreadPool = new TestThreadPool("test");
+        TestThreadPool testThreadPool = new TestThreadPool("test");
         ThreadPoolMergeExecutorService threadPoolMergeExecutorService = ThreadPoolMergeExecutorService
             .maybeCreateThreadPoolMergeExecutorService(
                 testThreadPool,
@@ -57,4 +65,66 @@ public void testMergeTasksAreAbortedWhenThreadPoolIsShutdown() {
         verify(mergeTask, times(0)).run();
         assertTrue(threadPoolMergeExecutorService.allDone());
     }
+
+    public void testBackloggedMergeTasksAreAllExecutedExactlyOnce() throws Exception {
+        int mergeExecutorThreadCount = randomIntBetween(1, 2);
+        Settings settings = Settings.builder()
+                .put(ThreadPoolMergeScheduler.USE_THREAD_POOL_MERGE_SCHEDULER_SETTING.getKey(), true)
+                // results in few merge threads, in order to increase contention
+                .put(EsExecutors.NODE_PROCESSORS_SETTING.getKey(), mergeExecutorThreadCount)
+                .build();
+        try (TestThreadPool testThreadPool = new TestThreadPool("test", settings)) {
+            ThreadPoolMergeExecutorService threadPoolMergeExecutorService = ThreadPoolMergeExecutorService
+                    .maybeCreateThreadPoolMergeExecutorService(testThreadPool, settings);
+            assertNotNull(threadPoolMergeExecutorService);
+            assertThat(threadPoolMergeExecutorService.getMaxConcurrentMerges(), equalTo(mergeExecutorThreadCount));
+            int mergeTaskCount = randomIntBetween(3, 30);
+            CountDownLatch mergeTasksDoneLatch = new CountDownLatch(mergeTaskCount);
+            CountDownLatch mergeTasksReadyLatch = new CountDownLatch(mergeTaskCount);
+            CountDownLatch submitTaskLatch = new CountDownLatch(1);
+            Collection<ThreadPoolMergeScheduler.MergeTask> generatedMergeTasks = ConcurrentCollections.newConcurrentSet();
+            for (int i = 0; i < mergeTaskCount; i++) {
+                new Thread(()-> {
+                    ThreadPoolMergeScheduler.MergeTask mergeTask = mock(ThreadPoolMergeScheduler.MergeTask.class);
+                    when(mergeTask.isRunning()).thenReturn(false);
+                    boolean supportsIOThrottling = randomBoolean();
+                    when(mergeTask.supportsIOThrottling()).thenReturn(supportsIOThrottling);
+                    long mergeSize = randomNonNegativeLong();
+                    when(mergeTask.estimatedMergeSize()).thenReturn(mergeSize);
+                    doAnswer(mock -> {
+                        // each individual merge task can either "run" or be "backlogged"
+                        boolean runNowOrBacklog = randomBoolean();
+                        if (runNowOrBacklog) {
+                            mergeTasksDoneLatch.countDown();
+                        } else {
+                            testThreadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
+                                // reenqueue backlogged merge task
+                                threadPoolMergeExecutorService.reEnqueueBackloggedMergeTask(mergeTask);
+                            });
+                        }
+                        return runNowOrBacklog;
+                    }).when(mergeTask).runNowOrBacklog();
+                    generatedMergeTasks.add(mergeTask);
+                    mergeTasksReadyLatch.countDown();
+                    // make all threads submit merge tasks at once
+                    safeAwait(submitTaskLatch);
+                    threadPoolMergeExecutorService.submitMergeTask(mergeTask);
+                }).start();
+            }
+            safeAwait(mergeTasksReadyLatch);
+            submitTaskLatch.countDown();
+            safeAwait(mergeTasksDoneLatch);
+            assertBusy(() -> {
+                for (ThreadPoolMergeScheduler.MergeTask mergeTask : generatedMergeTasks) {
+                    verify(mergeTask, times(1)).run();
+                    if (mergeTask.supportsIOThrottling()) {
+                        verify(mergeTask).setIORateLimit(anyDouble());
+                    } else {
+                        verify(mergeTask, times(0)).setIORateLimit(anyDouble());
+                    }
+                }
+                threadPoolMergeExecutorService.allDone();
+            });
+        }
+    }
 }