Fix driver scheduler ready queue reservation leak (#17919)

Author: JackieTien97

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/schedule/queue/IndexedBlockingReserveQueue.java

@@ -49,6 +49,7 @@ public synchronized E poll() throws InterruptedException {
     E output = pollFirst();
     size--;
     reservedSize++;
+    markReserved(output);
     return output;
   }
 
@@ -68,15 +69,44 @@ public synchronized void repush(E element) {
       throw new NullPointerException(CalcMessages.PUSHED_ELEMENT_IS_NULL);
     }
     pushToQueue(element);
-    reservedSize--;
+    decreaseReservedSizeIfNecessary(element);
     size++;
     this.notifyAll();
   }
 
   /**
    * For task that is not in readyQueue when it's cleared, it won't be added into the queue again.
    */
-  public synchronized void decreaseReservedSize() {
-    this.reservedSize--;
+  public synchronized boolean decreaseReservedSize(E element) {
+    if (element == null) {
+      throw new NullPointerException(CalcMessages.PUSHED_ELEMENT_IS_NULL);
+    }
+    return decreaseReservedSizeIfNecessary(element);
+  }
+
+  public final synchronized int getReservedSize() {
+    return reservedSize;
+  }
+
+  @Override
+  public synchronized void clear() {
+    super.clear();
+    this.reservedSize = 0;
+  }
+
+  protected void markReserved(E element) {
+    // Do nothing by default.
+  }
+
+  protected boolean releaseReserved(E element) {
+    return true;
+  }
+
+  private boolean decreaseReservedSizeIfNecessary(E element) {
+    if (!releaseReserved(element) || reservedSize <= 0) {
+      return false;
+    }
+    reservedSize--;
+    return true;
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java

@@ -347,16 +347,18 @@ private void clearDriverTask(DriverTask task) {
             return;
           case READY:
             task.setStatus(DriverTaskStatus.ABORTED);
-            readyQueue.remove(task.getDriverTaskId());
+            if (readyQueue.remove(task.getDriverTaskId()) == null) {
+              readyQueue.decreaseReservedSize(task);
+            }
             break;
           case BLOCKED:
             task.setStatus(DriverTaskStatus.ABORTED);
             blockedTasks.remove(task);
-            readyQueue.decreaseReservedSize();
+            readyQueue.decreaseReservedSize(task);
             break;
           case RUNNING:
             task.setStatus(DriverTaskStatus.ABORTED);
-            readyQueue.decreaseReservedSize();
+            readyQueue.decreaseReservedSize(task);
             break;
           case FINISHED:
             break;
@@ -422,6 +424,10 @@ public long getReadyQueueTaskCount() {
     return readyQueue.size();
   }
 
+  public long getReadyQueueReservedTaskCount() {
+    return readyQueue.getReservedSize();
+  }
+
   public long getBlockQueueTaskCount() {
     return blockedTasks.size();
   }
@@ -497,6 +503,7 @@ public boolean readyToRunning(DriverTask task) {
       task.lock();
       try {
         if (task.getStatus() != DriverTaskStatus.READY) {
+          readyQueue.decreaseReservedSize(task);
           return false;
         }
 
@@ -553,7 +560,7 @@ public void runningToFinished(DriverTask task, ExecutionContext context) {
         }
         task.updateSchedulePriority(context);
         task.setStatus(DriverTaskStatus.FINISHED);
-        readyQueue.decreaseReservedSize();
+        readyQueue.decreaseReservedSize(task);
       } finally {
         task.unlock();
       }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/multilevelqueue/MultilevelPriorityQueue.java

@@ -191,6 +191,16 @@ protected DriverTask get(DriverTask driverTask) {
         "MultilevelPriorityQueue does not support access element by get.");
   }
 
+  @Override
+  protected void markReserved(DriverTask task) {
+    task.markReservedInReadyQueue();
+  }
+
+  @Override
+  protected boolean releaseReserved(DriverTask task) {
+    return task.releaseReservedInReadyQueue();
+  }
+
   @Override
   protected void clearAllElements() {
     highestPriorityLevelQueue.clear();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java

@@ -60,6 +60,7 @@ public class DriverTask implements IDIndexedAccessible {
   private final DriverTaskHandle driverTaskHandle;
   private long lastEnterReadyQueueTime;
   private long lastEnterBlockQueueTime;
+  private boolean reservedInReadyQueue;
 
   private long estimatedMemorySize;
 
@@ -213,6 +214,18 @@ public void setLastEnterBlockQueueTime(long lastEnterBlockQueueTime) {
     this.lastEnterBlockQueueTime = lastEnterBlockQueueTime;
   }
 
+  public void markReservedInReadyQueue() {
+    reservedInReadyQueue = true;
+  }
+
+  public boolean releaseReservedInReadyQueue() {
+    if (!reservedInReadyQueue) {
+      return false;
+    }
+    reservedInReadyQueue = false;
+    return true;
+  }
+
   /** a comparator of ddl, the less the ddl is, the low order it has. */
   public static class TimeoutComparator implements Comparator<DriverTask> {
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/metric/DriverSchedulerMetricSet.java

@@ -39,6 +39,7 @@ private DriverSchedulerMetricSet() {
   public static final String READY_QUEUED_TIME = "ready_queued_time";
   public static final String BLOCK_QUEUED_TIME = "block_queued_time";
   public static final String READY_QUEUE_TASK_COUNT = "ready_queue_task_count";
+  public static final String READY_QUEUE_RESERVED_TASK_COUNT = "ready_queue_reserved_task_count";
   public static final String BLOCK_QUEUE_TASK_COUNT = "block_queue_task_count";
   private static final String TIMEOUT_QUEUE_SIZE = "timeout_queue_task_count";
   private static final String QUERY_MAP_SIZE = "query_map_size";
@@ -67,6 +68,13 @@ public void bindTo(AbstractMetricService metricService) {
         DriverScheduler::getReadyQueueTaskCount,
         Tag.NAME.toString(),
         READY_QUEUE_TASK_COUNT);
+    metricService.createAutoGauge(
+        Metric.DRIVER_SCHEDULER.toString(),
+        MetricLevel.IMPORTANT,
+        DriverScheduler.getInstance(),
+        DriverScheduler::getReadyQueueReservedTaskCount,
+        Tag.NAME.toString(),
+        READY_QUEUE_RESERVED_TASK_COUNT);
     metricService.createAutoGauge(
         Metric.DRIVER_SCHEDULER.toString(),
         MetricLevel.IMPORTANT,
@@ -109,6 +117,11 @@ public void unbindFrom(AbstractMetricService metricService) {
         Metric.DRIVER_SCHEDULER.toString(),
         Tag.NAME.toString(),
         READY_QUEUE_TASK_COUNT);
+    metricService.remove(
+        MetricType.AUTO_GAUGE,
+        Metric.DRIVER_SCHEDULER.toString(),
+        Tag.NAME.toString(),
+        READY_QUEUE_RESERVED_TASK_COUNT);
     metricService.remove(
         MetricType.AUTO_GAUGE,
         Metric.DRIVER_SCHEDULER.toString(),

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/DefaultDriverSchedulerTest.java

@@ -197,6 +197,47 @@ public void testReadyToRunning() {
     clear();
   }
 
+  @Test
+  public void testAbortReadyTaskAfterPollReleasesReadyQueueReservation() throws Exception {
+    IMPPDataExchangeManager mockMPPDataExchangeManager =
+        Mockito.mock(IMPPDataExchangeManager.class);
+    manager.setBlockManager(mockMPPDataExchangeManager);
+    ITaskScheduler defaultScheduler = manager.getScheduler();
+    IDriver mockDriver = Mockito.mock(IDriver.class);
+    DriverTaskHandle driverTaskHandle =
+        new DriverTaskHandle(
+            1,
+            (MultilevelPriorityQueue) manager.getReadyQueue(),
+            OptionalInt.of(Integer.MAX_VALUE));
+
+    QueryId queryId = new QueryId("test");
+    FragmentInstanceId instanceId =
+        new FragmentInstanceId(new PlanFragmentId(queryId, 0), "inst-0");
+    DriverTaskId driverTaskID = new DriverTaskId(instanceId, 0);
+    Mockito.when(mockDriver.getDriverTaskId()).thenReturn(driverTaskID);
+    DriverTask testTask =
+        new DriverTask(mockDriver, 100L, DriverTaskStatus.READY, driverTaskHandle, 0, false);
+
+    try {
+      manager.registerTaskToQueryMap(queryId, testTask);
+      manager.submitTaskToReadyQueue(testTask);
+
+      DriverTask polledTask = manager.getReadyQueue().poll();
+      Assert.assertSame(testTask, polledTask);
+      Assert.assertEquals(1, manager.getReadyQueueReservedTaskCount());
+
+      manager.abortFragmentInstance(instanceId, new RuntimeException("test abort"));
+
+      Assert.assertEquals(DriverTaskStatus.ABORTED, testTask.getStatus());
+      Assert.assertEquals(0, manager.getReadyQueue().size());
+      Assert.assertEquals(0, manager.getReadyQueueReservedTaskCount());
+      Assert.assertFalse(defaultScheduler.readyToRunning(polledTask));
+      Assert.assertEquals(0, manager.getReadyQueueReservedTaskCount());
+    } finally {
+      clear();
+    }
+  }
+
   @Test
   public void testRunningToReady() {
     IMPPDataExchangeManager mockMPPDataExchangeManager =