Run tasks in the determinsitic task queue upto a cutoff time

server/src/test/java/org/elasticsearch/gateway/GatewayServiceTests.java

@@ -300,12 +300,10 @@ public void testScheduledRecoveryNoOpWhenNewTermBegins() {
         // The 1st schedule is effectively cancelled if a new term begins
         final TimeValue elapsed = TimeValue.timeValueMinutes(1);
         final ClusterState clusterStateOfTerm2 = incrementTerm(clusterService.state());
-        deterministicTaskQueue.scheduleAt(
+        deterministicTaskQueue.scheduleAtAndRunUpTo(
             elapsed.millis(),
             () -> setClusterStateTaskQueue.submitTask(randomAlphaOfLength(5), new SetClusterStateTask(clusterStateOfTerm2), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
         assertThat(gatewayService.currentPendingStateRecovery, not(sameInstance(pendingStateRecoveryOfTerm1)));
         // The 1st scheduled recovery is now a no-op
         deterministicTaskQueue.advanceTime();
@@ -384,20 +382,16 @@ public void testScheduledRecoveryWithRecoverAfterNodes() {
 
         // The 1st schedule is cancelled when data nodes drop below recoverAfterDataNodes
         final TimeValue elapsed = TimeValue.timeValueMinutes(1);
-        deterministicTaskQueue.scheduleAt(
+        deterministicTaskQueue.scheduleAtAndRunUpTo(
             elapsed.millis(),
             () -> setDataNodeCountTaskQueue.submitTask(randomAlphaOfLength(5), new SetDataNodeCountTask(recoverAfterNodes - 1), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
 
         // The 2nd scheduled recovery when data nodes are above recoverAfterDataNodes again
-        deterministicTaskQueue.scheduleAt(
+        deterministicTaskQueue.scheduleAtAndRunUpTo(
             elapsed.millis() * 2,
             () -> setDataNodeCountTaskQueue.submitTask(randomAlphaOfLength(5), new SetDataNodeCountTask(recoverAfterNodes), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
 
         assertThat(gatewayService.currentPendingStateRecovery, sameInstance(pendingStateRecoveryOfInitialTerm));
         assertThat(deterministicTaskQueue.getCurrentTimeMillis(), equalTo(elapsed.millis() * 2));

test/framework/src/main/java/org/elasticsearch/common/util/concurrent/DeterministicTaskQueue.java

@@ -101,6 +101,19 @@ public void runAllTasksInTimeOrder() {
         }
     }
 
+    /**
+     * Run all {@code runnableTasks} and {@code deferredTasks} that are scheduled to run before or on the given {@code timeInMillis}.
+     * The current time will be set to {@code timeInMillis} once the method returns.
+     */
+    public void runTasksUpToTimeInOrder(long timeInMillis) {
+        runAllRunnableTasks();
+        while (nextDeferredTaskExecutionTimeMillis <= timeInMillis) {
+            advanceTime();
+            runAllRunnableTasks();
+        }
+        currentTimeMillis = timeInMillis;
+    }
+
     /**
      * @return whether there are any runnable tasks.
      */
@@ -174,6 +187,14 @@ public void scheduleAt(final long executionTimeMillis, final Runnable task) {
         }
     }
 
+    /**
+     * Similar to {@link #scheduleAt} but also advance time to {@code executionTimeMillis} and run all eligible tasks.
+     */
+    public void scheduleAtAndRunUpTo(final long executionTimeMillis, final Runnable task) {
+        scheduleAt(executionTimeMillis, task);
+        runTasksUpToTimeInOrder(executionTimeMillis);
+    }
+
     private void scheduleDeferredTask(DeferredTask deferredTask) {
         nextDeferredTaskExecutionTimeMillis = Math.min(nextDeferredTaskExecutionTimeMillis, deferredTask.executionTimeMillis());
         latestDeferredExecutionTime = Math.max(latestDeferredExecutionTime, deferredTask.executionTimeMillis());

test/framework/src/test/java/org/elasticsearch/common/util/concurrent/DeterministicTaskQueueTests.java

@@ -14,10 +14,14 @@
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -229,6 +233,53 @@ public void testRunInTimeOrder() {
         assertThat(strings, contains("foo", "bar"));
     }
 
+    public void testRunTasksUpToTimeInOrder() {
+        final DeterministicTaskQueue taskQueue = new DeterministicTaskQueue();
+        // The queue does _not_ have to be a clean slate before test
+        if (randomBoolean()) {
+            taskQueue.scheduleAt(randomLongBetween(1, 100), () -> {});
+            taskQueue.runAllTasksInTimeOrder();
+        }
+
+        final long cutoffTimeInMillis = randomLongBetween(taskQueue.getCurrentTimeMillis(), 1000);
+        final Set<Integer> seenNumbers = new HashSet<>();
+
+        final int nRunnableTasks = randomIntBetween(0, 10);
+        IntStream.range(0, nRunnableTasks).forEach(i -> taskQueue.scheduleNow(() -> seenNumbers.add(i)));
+
+        final int nDeferredTasksUpToCutoff = randomIntBetween(0, 10);
+        IntStream.range(0, nDeferredTasksUpToCutoff)
+            .forEach(i -> taskQueue.scheduleAt(randomLongBetween(0, cutoffTimeInMillis), () -> seenNumbers.add(i + nRunnableTasks)));
+
+        IntStream.range(0, randomIntBetween(0, 10))
+            .forEach(
+                i -> taskQueue.scheduleAt(
+                    randomLongBetween(cutoffTimeInMillis + 1, 2 * cutoffTimeInMillis),
+                    () -> seenNumbers.add(i + nRunnableTasks + nDeferredTasksUpToCutoff)
+                )
+            );
+
+        taskQueue.runTasksUpToTimeInOrder(cutoffTimeInMillis);
+        assertThat(seenNumbers, equalTo(IntStream.range(0, nRunnableTasks + nDeferredTasksUpToCutoff).boxed().collect(Collectors.toSet())));
+        assertThat(taskQueue.getCurrentTimeMillis(), equalTo(cutoffTimeInMillis));
+    }
+
+    public void testScheduleAtAndRunUpTo() {
+        final DeterministicTaskQueue taskQueue = new DeterministicTaskQueue();
+        final Set<Integer> seenNumbers = new HashSet<>();
+
+        final int nRunnableTasks = randomIntBetween(0, 10);
+        IntStream.range(0, nRunnableTasks).forEach(i -> taskQueue.scheduleNow(() -> seenNumbers.add(i)));
+        taskQueue.scheduleAt(500, () -> seenNumbers.add(500));
+        taskQueue.scheduleAt(800, () -> seenNumbers.add(800));
+
+        final long executionTimeMillis = randomLongBetween(1, 400);
+        taskQueue.scheduleAtAndRunUpTo(executionTimeMillis, () -> seenNumbers.add(nRunnableTasks));
+
+        assertThat(seenNumbers, equalTo(IntStream.rangeClosed(0, nRunnableTasks).boxed().collect(Collectors.toSet())));
+        assertThat(taskQueue.getCurrentTimeMillis(), equalTo(executionTimeMillis));
+    }
+
     public void testThreadPoolEnqueuesTasks() {
         final DeterministicTaskQueue taskQueue = new DeterministicTaskQueue();
         final List<String> strings = new ArrayList<>(2);