Run tasks in the determinsitic task queue upto a cutoff time

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

@@ -304,8 +304,7 @@ public void testScheduledRecoveryNoOpWhenNewTermBegins() {
             elapsed.millis(),
             () -> setClusterStateTaskQueue.submitTask(randomAlphaOfLength(5), new SetClusterStateTask(clusterStateOfTerm2), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
+        deterministicTaskQueue.runTasksUpToTimeInOrder(elapsed.millis());
         assertThat(gatewayService.currentPendingStateRecovery, not(sameInstance(pendingStateRecoveryOfTerm1)));
         // The 1st scheduled recovery is now a no-op
         deterministicTaskQueue.advanceTime();
@@ -388,16 +387,14 @@ public void testScheduledRecoveryWithRecoverAfterNodes() {
             elapsed.millis(),
             () -> setDataNodeCountTaskQueue.submitTask(randomAlphaOfLength(5), new SetDataNodeCountTask(recoverAfterNodes - 1), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
+        deterministicTaskQueue.runTasksUpToTimeInOrder(elapsed.millis());
 
         // The 2nd scheduled recovery when data nodes are above recoverAfterDataNodes again
         deterministicTaskQueue.scheduleAt(
             elapsed.millis() * 2,
             () -> setDataNodeCountTaskQueue.submitTask(randomAlphaOfLength(5), new SetDataNodeCountTask(recoverAfterNodes), null)
         );
-        deterministicTaskQueue.advanceTime();
-        deterministicTaskQueue.runAllRunnableTasks();
+        deterministicTaskQueue.runTasksUpToTimeInOrder(elapsed.millis() * 2);
 
         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,17 @@ public void runAllTasksInTimeOrder() {
         }
     }
 
+    /**
+     * Run all {@code runnableTasks} and {@code deferredTasks} that are scheduled to run before or on the given {@code timeInMillis}.
+     */
+    public void runTasksUpToTimeInOrder(long timeInMillis) {
+        runAllRunnableTasks();
+        while (nextDeferredTaskExecutionTimeMillis <= timeInMillis) {
+            advanceTime();
+            runAllRunnableTasks();
+        }
+    }
+
     /**
      * @return whether there are any runnable tasks.
      */

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,37 @@ 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(), lessThanOrEqualTo(cutoffTimeInMillis));
+    }
+
     public void testThreadPoolEnqueuesTasks() {
         final DeterministicTaskQueue taskQueue = new DeterministicTaskQueue();
         final List<String> strings = new ArrayList<>(2);