背景

Apache IoTDB 查询引擎目前采用 MPP 架构，一条查询 SQL 大致会经历下图几个阶段：

image

FragmentInstance 是分布式计划被拆分后实际分发到各个节点进行执行的实例。由于每个节点会同时接收来自于多个并发 Query 的多个 FragmentInstance，这些 FragmentInstance 在执行时可能由于等待上游数据而处于阻塞状态、或者数据就绪可以执行、或者超时需要被取消。因此，需要一个较为合理的调度策略，保证在分配给 FragmentInstance 的有限资源内，能够满足高并发的查询需求，同时尽可能避免出现饿死或者死锁情况。

在具体实现中，查询引擎里真正执行查询计算的算子树 Operator Tree 是类 Driver 的一个成员变量，Driver 负责控制Operator 的运行。DriverTask 是 Driver 的一层封装，也是调度模块真正调度的对象。一个 FragmentInstance 可能对应多个 Driver，而 Driver 与 DriverTask 是一一对应的。

本文主要介绍 Apache IoTDB 查询引擎在 DataNode 上如何调度和执行 DriverTask。相关代码位于包 org.apache.iotdb.db.mpp.execution.schedule

DriverTask 调度与执行

调度模块维护了两个队列：

• ReadyQueue：处于 Ready 状态的 DriverTask 队列
• TimeoutQueue：所有当前节点未结束的 DriverTask 按照超时时间排序的队列

同时处于 Blocked 状态的 DriverTask 会被放入集合 blockedTasks 进行记录。

总体而言，DriverTask 的调度执行参考了协程思想和操作系统任务调度机制。分配给查询引擎调度模块的线程数是固定的，可以通过配置项更改。来自于不同节点的 FragmentInstance 的 DriverTask 在 init 时会被加入 ReadyQueue。执行线程会不断拉取 ReadyQueue 队头的任务进行执行，每次只执行一个时间片，然后根据 DriverTask 的状态决定是否要将 DriverTask 重新放回 ReadyQueue。可以结合下图帮助理解：

DriverTask 完整的生命周期与状态

image

如上图所示，目前 DriverTask 的状态包括：

• Ready：就绪状态。在以下场景，DriverTask 会处于 Ready 状态：
  • 新建 DriverTask 时，状态会被设置为 Ready，然后加入到 ReadyQueue 中。
  • 当 DriverTask 依赖的上游数据就绪时，DataBlockManager 会调用回调接口，将其状态从 Blocked 改为 Ready，并从 blockedTasks 中移除。
  • 分配的时间片用完，会进入 ReadyQueue，并且状态从 Running 转换成 Ready。
• Blocked：阻塞状态。在以下场景，DriverTask 会处于阻塞状态：
  • 依赖的上游数据为空，需要等待上游数据时，会从 Running 状态改为 Blocked，并放入 blockedTasks。
  • 向下游输出数据的 buffer 已满暂时无法发送数据，会从 Running 状态改为 Blocked。
• Running：执行中状态。当处于 Ready 的 DriverTask 被线程调度时，从 ReadyQueue 中移除，状态变为 Running。
• Finished：完成状态。DriverTask 变为完成状态后，调度模块会清理此 DriverTask 的信息。
• Aborted：终止状态。在任何情况下，出现以下情况，DriverTask 会立即进入终止状态，并释放所有申请的资源。
  • 执行超时
  • 用户中断了查询
  • 不可恢复的 Exception

调度模块重要组件

image

上图包含了调度模块的一些重要组件，下面对调度模块重要组件进行介绍，理解这些组件的作用可以帮助您更好地阅读源码。

Worker Thread

真正负责执行 DriverTask 的物理线程，具体实现类为 DriverSchedulerThread。数量可通过配置参数进行配置，实例启动后不可改变。

DriverSchedulerThread 的实现：

• 生命周期和查询引擎一致。
• 内部主体为无限循环，只通过 InterruptedException 中断（当服务停止时会发送 InterruptedException）。
• 循环会尝试去 ReadyQueue 拉取队头的 DriverTask。若队列为空，则 Worker Thread 进入阻塞状态。
• Worker Thread 在执行 DriverTask 时，会调用 DriverTask.processFor()，然后返回 ListenableFuture。为了保证 Worker Thread 不会因为某个执行时间较长的 DriverTask 导致其他 DriverTask 饿死，引入了时间片机制。当 Driver#processFor 方法会接收一个时间片长度作为参数，processFor 会运行时间片长度的时间，执行时间超过时间片长度时，processFor() 方法会结束运行，然后返回一个 Future。(目前的时间片为代码内置的常量: 100ms。后续可能会考虑变成用户可配置的项。但是需要有范围值保护。过大的时间片会使得此机制失效，过小的则会频繁触发 DriverTask 的状态切换，影响执行效率。)
• 根据返回的 Future，会有如下操作：
  • 若 Future 被 cancel，则终止当前 DriverTask 的执行，将其设置为 Aborted 状态。
  • 若执行完成，则将 DriverTask 置为 Finished 状态。
  • 若时间片用完，则将 DriverTask 置为 Ready 状态，计算并更新调度权重，将 DriverTask 加入到 ReadyQueue。
  • 若是因阻塞导致执行权让出，则将 DriverTask 置为 Blocked 状态，并注册 Blocked → Ready 的回调逻辑。

具体流程可以结合下图进行理解：

image

代码实现为：

public void execute(DriverTask task) throws InterruptedException {
    long startNanos = ticker.read();
    // try to switch it to RUNNING
    if (!scheduler.readyToRunning(task)) {
      return;
    }
    IDriver driver = task.getDriver();
    CpuTimer timer = new CpuTimer();
    ListenableFuture<?> future = driver.processFor(EXECUTION_TIME_SLICE);
    CpuTimer.CpuDuration duration = timer.elapsedTime();
    // If the future is cancelled, the task is in an error and should be thrown.
    if (future.isCancelled()) {
      task.setAbortCause(DriverTaskAbortedException.BY_ALREADY_BEING_CANCELLED);
      scheduler.toAborted(task);
      return;
    }
    long quantaScheduledNanos = ticker.read() - startNanos;
    ExecutionContext context = new ExecutionContext();
    context.setCpuDuration(duration);
    context.setScheduledTimeInNanos(quantaScheduledNanos);
    context.setTimeSlice(EXECUTION_TIME_SLICE);
    if (driver.isFinished()) {
      scheduler.runningToFinished(task, context);
      return;
    }

    if (future.isDone()) {
      scheduler.runningToReady(task, context);
    } else {
      scheduler.runningToBlocked(task, context);
      future.addListener(
          () -> {
            try (SetThreadName driverTaskName2 =
                new SetThreadName(task.getDriver().getDriverTaskId().getFullId())) {
              scheduler.blockedToReady(task);
            }
          },
          listeningExecutor);
    }
  }
}

Sentinel Thread

负责监控 DriverTask 超时的物理线程，全局唯一，具体实现类为 DriverTaskTimeoutSentinelThread。

DriverTaskTimeoutSentinelThread 的实现：

• 内部主体为无限循环，只通过 InterruptedException 中断（当服务停止时会发送 InterruptedException）。
• 尝试去 timeoutQueue 拉取队头的 DriverTask。若队列为空，则 Sentinel Thread 进入睡眠状态。
• Sentinel 在拉取 DriverTask 时，会判断当前系统时间是否超过了超时时间：
  • 若超时，则将状态置为 Aborted 状态，走超时处理逻辑。
  • 若未超时，则睡眠至超时时间，将状态置为 Aborted 状态，走超时处理逻辑。

可以结合下图进行理解：

image

优先调度队列 ReadyQueue

目前实现参考了 Trino 的 MultilevelSplitQueue，在 IoTDB 里的实现类为 MultilevelPriorityQueue，设计思路可以参考博客 Trino 源码阅读 —— MultiLevelSplitQueue 调度机制。

该队列特点：

• 线程安全。
• 是一个阻塞队列，有最大长度限制。
• 存在任务降级机制，设计初衷是避免任务出现饿死，提升 CPU 利用率。

超时队列 TimeoutQueue

根据 DriverTask 的超时 deadline 排序的最大堆，超时时间越早的 DriverTask 就会被先做超时检查。

该队列长度应该有最大限制。

该队列特点：

• 线程安全。
• 按照 DriverTask 的调度权重排序，在 O(lgn) 的时间复杂度内完成队列元素的 pull 和 push。
• 有根据 DriverTask 的 id 做索引查询的能力，能够在 O(lgn) 的时间复杂度内完成随机元素的删除。

阻塞任务集合 BlockedTasks

处于 Blocked 状态的 DriverTask 的集合，线程安全，在 O(1) 的时间复杂度内完成元素的读取。

调度器 DriverScheduler

调度模块的核心，持有线程资源，即之前提到的 WorkerThread 和 SentinelThread。维护了 ReadyQueue 和 TimeoutQueue，FragmentInstance 可以通过 DriverScheduler 提交 Driver，DriverScheduler 负责将 Driver 封装成 DriverTask 并进一步执行。

DriverScheduler 负责切换 DriverTask 的状态，主要通过内部类 Scheduler 完成。ITaskScheduler 定义了切换 DriverTask 状态的接口，Scheduler 实现了这些接口。接口定义如下：

/** the scheduler interface of {@link DriverTask} */
public interface ITaskScheduler {

  /**
   * Switch a task from {@link DriverTaskStatus#BLOCKED} to {@link DriverTaskStatus#READY}.
   *
   * @param task the task to be switched.
   */
  void blockedToReady(DriverTask task);

  /**
   * Switch a task from {@link DriverTaskStatus#READY} to {@link DriverTaskStatus#RUNNING}.
   *
   * @param task the task to be switched.
   * @return true if it's switched to the target status successfully, otherwise false.
   */
  boolean readyToRunning(DriverTask task);

  /**
   * Switch a task from {@link DriverTaskStatus#RUNNING} to {@link DriverTaskStatus#READY}.
   *
   * @param task the task to be switched.
   * @param context the execution context of last running.
   */
  void runningToReady(DriverTask task, ExecutionContext context);

  /**
   * Switch a task from {@link DriverTaskStatus#RUNNING} to {@link DriverTaskStatus#BLOCKED}.
   *
   * @param task the task to be switched.
   * @param context the execution context of last running.
   */
  void runningToBlocked(DriverTask task, ExecutionContext context);

  /**
   * Switch a task from {@link DriverTaskStatus#RUNNING} to {@link DriverTaskStatus#FINISHED}.
   *
   * @param task the task to be switched.
   * @param context the execution context of last running.
   */
  void runningToFinished(DriverTask task, ExecutionContext context);

  /**
   * Switch a task to {@link DriverTaskStatus#ABORTED}.
   *
   * @param task the task to be switched.
   */
  void toAborted(DriverTask task);

Blocked→ Ready

总体流程可以参考下图：

image

红色三角处表示，当获取到锁之后，还需要再次确认 DriverTask 状态是否符合预期（在排队等锁时可能被 SentinelThread 改为 Aborted 状态）。若为 Aborted 状态，则后续流程全部跳过。

代码实现为：

@Override
public void blockedToReady(DriverTask task) {
  task.lock();
  try {
    if (task.getStatus() != DriverTaskStatus.BLOCKED) {
      return;
    }

    task.setStatus(DriverTaskStatus.READY);
    QUERY_METRICS.recordTaskQueueTime(
        BLOCK_QUEUED_TIME, System.nanoTime() - task.getLastEnterBlockQueueTime());
    task.setLastEnterReadyQueueTime(System.nanoTime());
    task.resetLevelScheduledTime();
    readyQueue.push(task);
    blockedTasks.remove(task);
  } finally {
    task.unlock();
  }
}

Running -> Ready

计算并更新调度权重，将 DriverTask 加入到 ReadyQueue。

@Override
public void runningToReady(DriverTask task, ExecutionContext context) {
  task.lock();
  try {
    if (task.getStatus() != DriverTaskStatus.RUNNING) {
      return;
    }
    task.updateSchedulePriority(context);
    task.setStatus(DriverTaskStatus.READY);
    task.setLastEnterReadyQueueTime(System.nanoTime());
    readyQueue.push(task);
  } finally {
    task.unlock();
  }
}

Running -> Blocked

更新调度权重，然后将 DriverTask 加入 blockedTasks。

@Override
public void runningToBlocked(DriverTask task, ExecutionContext context) {
  task.lock();
  try {
    if (task.getStatus() != DriverTaskStatus.RUNNING) {
      return;
    }
    task.updateSchedulePriority(context);
    task.setStatus(DriverTaskStatus.BLOCKED);
    task.setLastEnterBlockQueueTime(System.nanoTime());
    blockedTasks.add(task);
  } finally {
    task.unlock();
  }
}

Running -> Finished

更新调度权重，清理 DriverTask 相关信息。

@Override
public void runningToFinished(DriverTask task, ExecutionContext context) {
  task.lock();
  try {
    if (task.getStatus() != DriverTaskStatus.RUNNING) {
      return;
    }
    task.updateSchedulePriority(context);
    task.setStatus(DriverTaskStatus.FINISHED);
    clearDriverTask(task);
  } finally {
    task.unlock();
  }
}

toAborted

由于同一个 FragmentInstance 的 DriverTask 之间有依赖性，一个 DriverTask 被置为 Aborted，其余相关的 DriverTask 也应该被置为 Aborted。

@Override
public void toAborted(DriverTask task) {
  try (SetThreadName driverTaskName =
      new SetThreadName(task.getDriver().getDriverTaskId().getFullId())) {
    task.lock();
    try {
      // If a task is already in an end state, it indicates that the task is finalized in other
      // threads.
      if (task.isEndState()) {
        return;
      }
      logger.warn(
          "The task {} is aborted. All other tasks in the same query will be cancelled",
          task.getDriverTaskId());
      clearDriverTask(task);
    } finally {
      task.unlock();
    }
    QueryId queryId = task.getDriverTaskId().getQueryId();
    Map<FragmentInstanceId, Set<DriverTask>> queryRelatedTasks = queryMap.get(queryId);
    if (queryRelatedTasks != null) {
      for (Set<DriverTask> fragmentRelatedTasks : queryRelatedTasks.values()) {
        if (fragmentRelatedTasks != null) {
          for (DriverTask otherTask : fragmentRelatedTasks) {
            if (task.equals(otherTask)) {
              continue;
            }
            otherTask.lock();
            try {
              otherTask.setAbortCause(DriverTaskAbortedException.BY_QUERY_CASCADING_ABORTED);
              clearDriverTask(otherTask);
            } finally {
              otherTask.unlock();
            }
          }
        }
      }
    }
  }
}