fix: Do not block scheduler if there are stages ready for scheduling (#26268)

## Description

Do not block query scheduler when there is a stage ready to be
scheduled.

## Motivation and Context

Consider a query with two stages, Stage 1 and Stage 2.

Before this change when one for the stages is blocked the scheduler will
block and wait for it to unblock preventing the second stage from
scheduling.

## Impact

Reduced scheduling delays in certain cases

## Test Plan

Tested with a simple INSERT INTO query with scaled writers enabled.

ScaledWriters scheduler has a fixed 200ms delay. Delay in ScaledWriter
scheduler was causing slow split scheduling for TableScan stage.

## Contributor checklist

- [x] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [x] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [x] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [x] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [x] Adequate tests were added if applicable.
- [x] CI passed.

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== NO RELEASE NOTE ==
```

Author: arhimondr

presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/SqlQueryScheduler.java

@@ -27,6 +27,7 @@
 import com.facebook.presto.execution.RemoteTask;
 import com.facebook.presto.execution.RemoteTaskFactory;
 import com.facebook.presto.execution.SqlStageExecution;
+import com.facebook.presto.execution.StageExecutionId;
 import com.facebook.presto.execution.StageExecutionInfo;
 import com.facebook.presto.execution.StageExecutionState;
 import com.facebook.presto.execution.StageId;
@@ -423,6 +424,7 @@ private void schedule()
             Set<StageId> completedStages = new HashSet<>();
 
             List<ExecutionSchedule> sectionExecutionSchedules = new LinkedList<>();
+            Map<StageExecutionId, ListenableFuture<?>> blockedStages = new HashMap<>();
 
             while (!Thread.currentThread().isInterrupted()) {
                 // remove finished section
@@ -445,19 +447,24 @@ private void schedule()
                         .forEach(sectionExecutionSchedules::add);
 
                 while (sectionExecutionSchedules.stream().noneMatch(ExecutionSchedule::isFinished)) {
-                    List<ListenableFuture<?>> blockedStages = new ArrayList<>();
-
                     List<StageExecutionAndScheduler> executionsToSchedule = sectionExecutionSchedules.stream()
                             .flatMap(schedule -> schedule.getStagesToSchedule().stream())
                             .collect(toImmutableList());
 
+                    boolean allBlocked = true;
                     for (StageExecutionAndScheduler stageExecutionAndScheduler : executionsToSchedule) {
                         long startCpuNanos = THREAD_MX_BEAN.getCurrentThreadCpuTime();
                         long startWallNanos = System.nanoTime();
 
                         SqlStageExecution stageExecution = stageExecutionAndScheduler.getStageExecution();
                         stageExecution.beginScheduling();
 
+                        ListenableFuture<?> stillBlocked = blockedStages.get(stageExecution.getStageExecutionId());
+                        if (stillBlocked != null && !stillBlocked.isDone()) {
+                            continue;
+                        }
+                        blockedStages.remove(stageExecution.getStageExecutionId());
+
                         // perform some scheduling work
                         ScheduleResult result = stageExecutionAndScheduler.getStageScheduler()
                                 .schedule();
@@ -475,7 +482,10 @@ private void schedule()
                             stageExecution.schedulingComplete();
                         }
                         else if (!result.getBlocked().isDone()) {
-                            blockedStages.add(result.getBlocked());
+                            blockedStages.put(stageExecution.getStageExecutionId(), result.getBlocked());
+                        }
+                        else {
+                            allBlocked = false;
                         }
                         stageExecutionAndScheduler.getStageLinkage()
                                 .processScheduleResults(stageExecution.getState(), result.getNewTasks());
@@ -535,11 +545,11 @@ else if (!result.getBlocked().isDone()) {
                     }
 
                     // wait for a state change and then schedule again
-                    if (!blockedStages.isEmpty()) {
+                    if (allBlocked && !blockedStages.isEmpty()) {
                         try (TimeStat.BlockTimer timer = schedulerStats.getSleepTime().time()) {
-                            tryGetFutureValue(whenAnyComplete(blockedStages), 1, SECONDS);
+                            tryGetFutureValue(whenAnyComplete(blockedStages.values()), 1, SECONDS);
                         }
-                        for (ListenableFuture<?> blockedStage : blockedStages) {
+                        for (ListenableFuture<?> blockedStage : blockedStages.values()) {
                             blockedStage.cancel(true);
                         }
                     }