testEnqueuedMergeTasksAreUnblockedWhenEstimatedMergeSizeChanges

Author: albertzaharovits

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

@@ -896,6 +896,106 @@ public void testUnavailableBudgetBlocksNewMergeTasksFromStartingExecution() thro
         }
     }
 
+    public void testEnqueuedMergeTasksAreUnblockedWhenEstimatedMergeSizeChanges() throws Exception {
+        long diskSpaceLimitBytes = randomLongBetween(10L, 100L);
+        aFileStore.usableSpace = diskSpaceLimitBytes + randomLongBetween(1L, 100L);
+        aFileStore.totalSpace = aFileStore.usableSpace + randomLongBetween(1L, 10L);
+        bFileStore.usableSpace = diskSpaceLimitBytes + randomLongBetween(1L, 100L);
+        bFileStore.totalSpace = bFileStore.usableSpace + randomLongBetween(1L, 10L);
+        boolean aHasMoreSpace = aFileStore.usableSpace > bFileStore.usableSpace;
+        Settings.Builder settingsBuilder = Settings.builder().put(settings);
+        settingsBuilder.put(ThreadPoolMergeExecutorService.INDICES_MERGE_DISK_HIGH_WATERMARK_SETTING.getKey(), diskSpaceLimitBytes + "b");
+        try (
+            ThreadPoolMergeExecutorService threadPoolMergeExecutorService = ThreadPoolMergeExecutorService
+                .maybeCreateThreadPoolMergeExecutorService(
+                    testThreadPool,
+                    ClusterSettings.createBuiltInClusterSettings(settingsBuilder.build()),
+                    nodeEnvironment
+                )
+        ) {
+            assert threadPoolMergeExecutorService != null;
+            assertThat(threadPoolMergeExecutorService.getMaxConcurrentMerges(), greaterThanOrEqualTo(1));
+            final long availableBudget = aHasMoreSpace
+                ? aFileStore.usableSpace - diskSpaceLimitBytes
+                : bFileStore.usableSpace - diskSpaceLimitBytes;
+            final AtomicLong expectedAvailableBudget = new AtomicLong(availableBudget);
+            assertBusy(
+                () -> assertThat(threadPoolMergeExecutorService.getDiskSpaceAvailableForNewMergeTasks(), is(expectedAvailableBudget.get()))
+            );
+            List<ThreadPoolMergeScheduler.MergeTask> tasksRunList = new ArrayList<>();
+            List<ThreadPoolMergeScheduler.MergeTask> tasksAbortList = new ArrayList<>();
+            int submittedMergesCount = randomIntBetween(1, 5);
+            long[] mergeSizeEstimates = new long[submittedMergesCount];
+            for (int i = 0; i < submittedMergesCount; i++) {
+                // all these merge estimates are over-budget
+                mergeSizeEstimates[i] = availableBudget + randomLongBetween(1L, 10L);
+            }
+            for (int i = 0; i < submittedMergesCount;) {
+                ThreadPoolMergeScheduler.MergeTask mergeTask = mock(ThreadPoolMergeScheduler.MergeTask.class);
+                when(mergeTask.supportsIOThrottling()).thenReturn(randomBoolean());
+                doAnswer(mock -> {
+                    Schedule schedule = randomFrom(Schedule.values());
+                    if (schedule == BACKLOG) {
+                        testThreadPool.executor(ThreadPool.Names.GENERIC).execute(() -> {
+                            // re-enqueue backlogged merge task
+                            threadPoolMergeExecutorService.reEnqueueBackloggedMergeTask(mergeTask);
+                        });
+                    } else if (schedule == RUN) {
+                        tasksRunList.add(mergeTask);
+                    } else if (schedule == ABORT) {
+                        tasksAbortList.add(mergeTask);
+                    }
+                    return schedule;
+                }).when(mergeTask).schedule();
+                // randomly let some task complete
+                if (randomBoolean()) {
+                    // this task is not blocked
+                    when(mergeTask.estimatedRemainingMergeSize()).thenReturn(randomLongBetween(0L, availableBudget));
+                } else {
+                    // this task will initially be blocked because over-budget
+                    int finalI = i;
+                    doAnswer(mock -> mergeSizeEstimates[finalI]).when(mergeTask).estimatedRemainingMergeSize();
+                    i++;
+                }
+                assertTrue(threadPoolMergeExecutorService.submitMergeTask(mergeTask));
+            }
+            // assert tasks are blocked because their estimated merge size is over the available budget
+            assertBusy(() -> {
+                assertTrue(threadPoolMergeExecutorService.isMergingBlockedDueToInsufficientDiskSpace());
+                assertThat(threadPoolMergeExecutorService.getMergeTasksQueueLength(), is(submittedMergesCount));
+                assertThat(threadPoolMergeExecutorService.getDiskSpaceAvailableForNewMergeTasks(), is(availableBudget));
+                assertThat(threadPoolMergeExecutorService.getRunningMergeTasks().size(), is(0));
+            });
+            // change estimates to be under the available budget
+            for (int i = 0; i < submittedMergesCount; i++) {
+                mergeSizeEstimates[i] = randomLongBetween(0L, availableBudget);
+            }
+            // assert tasks are all unblocked because their estimated merge size is now under the available budget
+            assertBusy(() -> {
+                assertFalse(threadPoolMergeExecutorService.isMergingBlockedDueToInsufficientDiskSpace());
+                assertThat(threadPoolMergeExecutorService.getMergeTasksQueueLength(), is(0));
+                assertThat(threadPoolMergeExecutorService.getDiskSpaceAvailableForNewMergeTasks(), is(availableBudget));
+            });
+            // assert all merge tasks are either run or aborted
+            assertBusy(() -> {
+                for (ThreadPoolMergeScheduler.MergeTask mergeTask : tasksRunList) {
+                    verify(mergeTask, times(1)).run();
+                    verify(mergeTask, times(0)).abort();
+                }
+                for (ThreadPoolMergeScheduler.MergeTask mergeTask : tasksAbortList) {
+                    verify(mergeTask, times(0)).run();
+                    verify(mergeTask, times(1)).abort();
+                }
+            });
+        }
+        if (setThreadPoolMergeSchedulerSetting) {
+            assertWarnings(
+                "[indices.merge.scheduler.use_thread_pool] setting was deprecated in Elasticsearch "
+                    + "and will be removed in a future release. See the breaking changes documentation for the next major version."
+            );
+        }
+    }
+
     public void testMergeTasksAreUnblockedWhenMoreDiskSpaceBecomesAvailable() throws Exception {
         aFileStore.totalSpace = randomLongBetween(300L, 1_000L);
         bFileStore.totalSpace = randomLongBetween(300L, 1_000L);