React Scheduler - 优先级调度

开篇

Scheduler 调度器的核心除了「时间切片」requestHostCallback 外，还有一个功能就是「优先级调度」。

它会根据任务的优先级（任务设定的过期时间）决定先执行哪个任务。为了更快捷地查找到最高优先级任务，采用了数据结构「小顶堆」来存储任务。

另一篇「React Scheduler - 时间切片」可以看这里。

一、小顶堆

堆是一个完全二叉树，小顶堆的特点是：任意节点都小于等于其左右节点值。

由于堆是一个完全二叉树，因此适用于数组存储法，给定一个节点的下标 i（i 从 1 开始），它的父节点一定为 arr[i / 2]，左子节点为 arr[2 * i]，右子节点为 arr[2 * i + 1]。

与常规小顶堆不同，Scheduler 下使用的小顶堆是从下标 0 开始。

在 Scheduler 中，小顶堆采用了「插入式」创建方式，即：每次插入一个节点，重新堆化小顶堆。

小顶堆的常用方法如下：

function push() {} // 向小顶堆添加一个节点（先加入到堆尾，依次向上去重新堆化最小堆）
function peek() {} // 取出小顶堆的堆顶元素
function pop() {} // 删除小顶堆的堆顶元素，（删除堆顶，将堆尾元素放到堆顶，然后从上往下重新堆化最小堆）
function siftUp() {} // 从下往上堆化最小堆
function siftDown() {} // 从上往下堆化最小堆

小顶堆的具体实现如下：

type HeapNode = {
  id: number,
  sortIndex: number,
};
type Heap = Array<HeapNode>;

// 1、向小顶堆添加一个节点（加入到堆尾，重新堆化小顶堆）
function push(heap: Heap, node: HeapNode): void {
  const index = heap.length;
  heap.push(node);
  siftUp(heap, node, index); // 从 index 尾部位置开始，向上堆化小顶堆
}

// 2、查找堆顶元素（仅查找，不删除）
function peek(heap: Heap): HeapNode | null {
  const first = heap[0];
  return first === undefined ? null : first;
}

// 3、删除堆顶元素
function pop(heap: Heap): HeapNode | null {
  const first = heap[0];
  if (first !== undefined) {
    const last = heap.pop() as HeapNode;
    if (last !== first) {
      heap[0] = last;
      siftDown(heap, last, 0); // 从 0 头部位置开始，向下进行堆化最小堆
    }
    return first;
  } else {
    return null;
  }
}

// 优先级比较函数
function compare(a, b) {
  // 首先比较排序索引，然后比较任务id
  const diff = a.sortIndex - b.sortIndex;
  return diff !== 0 ? diff : a.id - b.id;
}

// 4、从堆尾节点开始到堆顶节点（从后往前）
function siftUp(heap: Heap, node: HeapNode, i: number) {
  let index = i;
  while (true) {
    // 堆数据结构：父节点 = index / 2，这里的 >>> 等价于 Math.floor(index / 2)
    // 由于堆是以 0 为堆顶，而非 1，在获取 parentIndex 时需要减 1
    const parentIndex = (index - 1) >>> 1;
    const parent = heap[parentIndex];
    // 父任务比新任务数值大，交换位置
    if (parent !== undefined && compare(parent, node) > 0) {
      heap[parentIndex] = node;
      heap[index] = parent;
      index = parentIndex;
    } else {
      return;
    }
  }
}

// 5、从堆顶节点开始到堆尾
function siftDown(heap: Heap, node: HeapNode, i: number) {
  let index = i;
  const length = heap.length;
  while (index < length) {
    // 0 是最小堆根节点，这里的 left 和 right 和常见的堆数据结构设计不太一致
    const leftIndex = (index + 1) * 2 - 1;
    const rightIndex = leftIndex + 1;
    const left = heap[leftIndex];
    const right = heap[rightIndex];
    // 如果 left 或 right 小于当前节点，交换位置
    if (left !== undefined && compare(left, node) < 0) {
      // 拿 left 和 right 进行比较
      if (right !== undefined && compare(right, left) < 0) {
        heap[index] = right;
        heap[rightIndex] = node;
        index = rightIndex;
      } else {
        heap[index] = left;
        heap[leftIndex] = node;
        index = leftIndex;
      }
    } else if (right !== undefined && compare(right, node) < 0) {
      heap[index] = right;
      heap[rightIndex] = node;
      index = rightIndex;
    } else {
      return;
    }
  }
}

二、Scheduler 优先级

Scheduler 优先级是独立于 React 的优先级（lane），它提供给任务执行的优先级定义有六种：

export const NoPriority = 0; // 没有优先级
export const ImmediatePriority = 1; // 直接优先级
export const UserBlockingPriority = 2; // 用户阻塞优先级
export const NormalPriority = 3; // 普通优先级
export const LowPriority = 4; // 低优先级
export const IdlePriority = 5; // 空闲优先级

在调度一个任务时，会根据传入的任务优先级，来计算 timeout 超时时间，不同优先级对应的超时时间定义如下：

var maxSigned31BitInt = 1073741823;
// 立即超时(执行)。如果不定义 delay，expirationTime 过期时间的计算就等于 currentTime - 1
var IMMEDIATE_PRIORITY_TIMEOUT = -1;
var USER_BLOCKING_PRIORITY_TIMEOUT = 250; // 用户操作优先级超时时间
var NORMAL_PRIORITY_TIMEOUT = 5000; // 正常优先级超时时间
var LOW_PRIORITY_TIMEOUT = 10000; // 低优先级超时时间
// 从不超时 Never times out
var IDLE_PRIORITY_TIMEOUT = maxSigned31BitInt;

将两者进行关联：根据任务优先级计算任务超时时间，是通过 timeoutForPriorityLevel 来实现：

function timeoutForPriorityLevel(priorityLevel) {
  switch (priorityLevel) {
    case ImmediatePriority:
      return IMMEDIATE_PRIORITY_TIMEOUT;
    case UserBlockingPriority:
      return USER_BLOCKING_PRIORITY_TIMEOUT;
    case IdlePriority:
      return IDLE_PRIORITY_TIMEOUT;
    case LowPriority:
      return LOW_PRIORITY_TIMEOUT;
    case NormalPriority:
    default:
      return NORMAL_PRIORITY_TIMEOUT;
  }
}

三、Scheduler 任务队列

Scheduler 是以小顶堆数据结构存储多个执行任务，小顶堆采用数组表示如下：

var taskQueue = []; // 已就绪任务队列
var timerQueue = []; // 延迟任务队列
var taskIdCounter = 1; // 维护 Scheduler Task 插入顺序。

其中根据任务的不同又会分为 taskQueue 已就绪任务队列和 timerQueue 延迟任务队列。

当任务标识了 delay 信息，会被加入到 timerQueue 延迟任务队列，否则加入到 taskQueue 就绪任务队列调用 requestHostCallback 时间切片执行任务。

这里的单个任务，你可以理解为是触发一次 React 更新。

Scheduler 工作原理图如下：

下面我们在 Scheduler 核心入口函数 unstable_scheduleCallback 看看任务进来后的「优先级调度」逻辑。

四、Scheduler 核心入口函数

Scheduler 暴露了 unstable_scheduleCallback(priorityLevel, callback, options) 来调度任务。其中：

1. priorityLevel 是调度任务的优先级；
2. callback 是需要执行的更新任务；
3. optoins 里面可以通过指定 delay 延迟执行任务 或 timeout 定义任务的超时时间。

4.1、计算任务开始时间和过期时间

• 任务的开始时间 startTime，决定了任务是要进入 taskQueue 还是 timerQueue 任务队列；
• 任务的过期时间 expirationTime，决定了任务在执行队列中的执行优先级顺序。

function unstable_scheduleCallback(priorityLevel, callback, options) {
  var currentTime = getCurrentTime(); // requestHostCallback 中暴露的方法
  var startTime，timeout; // 超时时间
  // 要延迟的任务，将 delay 加入到 startTime
  if (typeof options === 'object' && options !== null) {
    var delay = options.delay;
    if (typeof delay === 'number' && delay > 0) {
      startTime = currentTime + delay; // 延迟
    } else {
      startTime = currentTime;
    }
    // 根据任务优先级计算超时时间
    timeout = typeof options.timeout === 'number'
      ? options.timeout
      : timeoutForPriorityLevel(priorityLevel);
  } else {
    // 不需要延迟的任务，只需计算任务超时时间
    timeout = timeoutForPriorityLevel(priorityLevel);
    startTime = currentTime;
  }

  // 任务过期时间 = 当前时间(currentTime) + 延时(delay) + 优先级超时(timeout)
  var expirationTime = startTime + timeout;

  ...
}

4.2、创建 Task 任务

一个调度任务上包含了：执行回调 callback、开始时间 startTime、过期时间 expirationTime、排序 sortIndex 等信息。

function unstable_scheduleCallback(priorityLevel, callback, options) {
  ...

  var newTask = {
    id: taskIdCounter++,
    callback,
    priorityLevel,
    startTime,
    expirationTime,
    sortIndex: -1, // 用于小顶堆的重建排序
  };

  ...
}

4.3、加入任务队列

1. 如果任务还不需要执行（startTime > currentTime），将任务加入到 timerQueue 队列，延迟任务队列的排序 sortIndex 基于 startTime 进行；
2. 如果任务满足执行，将任务加入到 taskQueue，就绪任务队列的排序 sortIndex 则根据 expirationTime 过期时间绝定谁先执行。

function unstable_scheduleCallback(priorityLevel, callback, options) {
  ...

  // 延迟任务，加入到 timerQueue
  if (startTime > currentTime) {
    newTask.sortIndex = startTime; // 小顶堆的重建基于 startTime
    push(timerQueue, newTask); // 加入队列，重新堆化
    // 如果 taskQueue 任务队列中没有任务，并且当前 newTask 任务是 timerQueue 延迟任务队列中最先被执行的任务（最小堆的堆顶）
    if (peek(taskQueue) === null && newTask === peek(timerQueue)) {
      if (isHostTimeoutScheduled) {
        cancelHostTimeout(); // 取消现有的超时任务，开启新的超时任务
      } else {
        isHostTimeoutScheduled = true;
      }
      // Schedule a timeout.（注意：因为是延迟任务，要等 startTime - currentTime 后才会执行）
      requestHostTimeout(handleTimeout, startTime - currentTime);
    }
  }
  // 正常任务，加入到 taskQueue
  else {
    // 排序索引，根据时间将任务优先级进行排序，值越小越先被执行
    newTask.sortIndex = expirationTime;
    push(taskQueue, newTask);
    // 开启调度
    if (!isHostCallbackScheduled && !isPerformingWork) {
      isHostCallbackScheduled = true;
      requestHostCallback(flushWork); // flushWork 用于执行当前所有任务
    }
  }
}

4.4、flushWork 执行 taskQueue

对于 taskQueue 就绪任务，调用 requestHostCallback 时间切片，通过异步宏任务的方式调度 flushWork。

// 这是在执行工作时设置的，以防止再次进入
var isPerformingWork = false;
var currentTask = null;
var currentPriorityLevel = NormalPriority;

function flushWork(hasTimeRemaining, initialTime) {
  isPerformingWork = true;
  const previousPriorityLevel = currentPriorityLevel;
  try {
    return workLoop(hasTimeRemaining, initialTime);
  } finally {
    currentTask = null;
    currentPriorityLevel = previousPriorityLevel;
    isPerformingWork = false;
  }
}

flushWork 的核心是执行 workLoop 同步循环执行 taskQueue 中的 task.callback。

注意，这里的 workLoop 仅仅是 Scheduler 执行任务队列使用，与执行 React Fiber 执行单元没有直接关系。

function workLoop(hasTimeRemaining, initialTime) {
  let currentTime = initialTime; // requestHostCallback 调度任务的开始时间
  advanceTimers(currentTime); // 梳理延迟队列中是否有过期任务，让其加入到 taskQueue 任务队列中
  currentTask = peek(taskQueue); // 取出任务队列的第一个要执行的任务
  
  while (currentTask !== null) {
    // 如果当前任务未过期，但时间切片用尽了，中断执行，让出执行权给浏览器
    if (currentTask.expirationTime > currentTime && (!hasTimeRemaining || shouldYieldToHost())) {
      break;
    }

    const callback = currentTask.callback;
    if (callback !== null) {
      currentTask.callback = null;
      currentPriorityLevel = currentTask.priorityLevel;
      const didUserCallbackTimeout = currentTask.expirationTime <= currentTime; // 已经超时
      const continuationCallback = callback(didUserCallbackTimeout); // 执行 callback
      currentTime = getCurrentTime();
      // 如果 callback 任务执行完成会返回 null，若需要让出主线程但任务还未执行完成，则返回的是函数，详见下文 performConcurrentWorkOnRoot
      if (typeof continuationCallback === 'function') {
        currentTask.callback = continuationCallback; // 保存 callback，等待下次宏任务继续从中断的位置执行
      } else {
        if (currentTask === peek(taskQueue)) {
          pop(taskQueue); // 执行完任务，从堆中移除
        }
      }
      advanceTimers(currentTime); // 重新梳理 timerQueue
    } else {
      // 当前任务不存在 callback 要执行，从堆中移除该任务
      pop(taskQueue);
    }
    
    currentTask = peek(taskQueue);
  }

  // 返回任务队列的执行状态：是否有剩余任务工作
  if (currentTask !== null) {
    return true;
  } else {
    const firstTimer = peek(timerQueue);
    if (firstTimer !== null) {
      requestHostTimeout(handleTimeout, firstTimer.startTime - currentTime);
    }
    return false;
  }
}

1. 通过 while 不断从 taskQueue 小顶堆中取出第一个要执行的任务；
2. 当 taskQueue 中的任务执行完毕，或者时间切片用尽，中断循环，将执行权交给浏览器，开启下一次异步任务请求；
3. 执行 task.callback，梳理 timerQueue；
4. 如果说时间切片用完了，需要让出执行权，但 callback 中的 tasks 还没有执行完，保存当前任务进度 currentTask.callback = continuationCallback；
5. 若此时任务已过期，需要立即执行完成，不会等待下一个宏任务，而是继续循环执行此任务，直至完成；
6. 若此时任务未过期，跳出 workloop，等待 requestHostCallback 开启下一次宏任务请求。

这里如果不是很理解，可以往下看「六、如何使用？」部分，其中的 Example 会解释这一点。

workLoop 每次执行完一个 taskQueue task 后，会调用 advanceTimers 方法来看 timerQueue 中是否有任务需要移交到 taskQueue 队列中执行，具体实现如下：

function advanceTimers(currentTime) {
  let timer = peek(timerQueue);
  while (timer !== null) {
    if (timer.callback === null) {
      pop(timerQueue);
    } else if (timer.startTime <= currentTime) {
      // 任务已过期，需要被执行，转移到任务队列。
      pop(timerQueue);
      timer.sortIndex = timer.expirationTime; // 在加入 tastQueue 前，将任务的过期时间赋值给 sortIndex，用于任务时间的排序
      push(taskQueue, timer);
    } else { // 不需要转移
      return;
    }
    timer = peek(timerQueue);
  }
}

4.5、handleTimeout 处理 timerQueue

timerQueue 延迟任务，会不断地去使用 requestHostTimeout 执行 handleTimeout 梳理延迟任务，当任务移交到 taskQueue 队列后，调用 requestHostCallback(flushWork) 执行任务。

function handleTimeout(currentTime) {
  isHostTimeoutScheduled = false;
  advanceTimers(currentTime); // 重新梳理 task

  if (!isHostCallbackScheduled) {
    // advanceTimers 梳理后如果有新任务
    if (peek(taskQueue) !== null) {
      isHostCallbackScheduled = true;
      requestHostCallback(flushWork);
    } else {
      // 如果 taskQueue 仍然为空，就开始递归的调用该方法
      const firstTimer = peek(timerQueue);
      if (firstTimer !== null) {
        requestHostTimeout(handleTimeout, firstTimer.startTime - currentTime);
      }
    }
  }
}

五、更新视角下的 Scheduler 工作流程

下面我们从更新视角串联一下上面介绍的 Scheduler 流程。

当我们调用 setState 触发更新，会进入到 scheduleUpdateOnFiber --> ensureRootIsScheduled 中，如果是异步更新，调用 ScheduleCallback（Scheduler 核心入口函数）并传入 callback -> performConcurrentWorkOnRoot。

export function scheduleUpdateOnFiber(fiber, lane, eventTime) {
  ensureRootIsScheduled(root, eventTime);
}
function ensureRootIscheduled(root: FiberRoot, currentTime: number) {
  // ...
  const newCallbackPriority = returnNextLanesPriority();
  const schedulerPriorityLevel = lanePriorityToSchedulerPriority(newCallbackPriority);
  let newCallbackNode = scheduleCallback( // unstable_scheduleCallback
    schedulerPriorityLevel, // callback 执行优先级
    performConcurrentWorkOnRoot.bind(null, root),
  );
}

在 scheduleCallback 中会创建一个 Task 并调用 requestHostCallback 时间切片来开启宏任务执行 performConcurrentWorkOnRoot。

在这里会开启 renderRootConcurrent 渲染阶段进入 workLoopConcurrent 依次执行每个 Fiber 执行单元。

function performConcurrentWorkOnRoot(root, didTimeout) {
  const originalCallbackNode = root.callbackNode;
  // ...
  let exitStatus = renderRootConcurrent(root, lanes);
  // ...
  // 1、任务若是完成，在 commitRootImpl 会清除掉 root.callbackNode，若没有完成，说明当然还有剩余任务，返回当前函数，先将执行权交还给浏览器。
  if (root.callbackNode === originalCallbackNode) {
    return performConcurrentWorkOnRoot.bind(null, root);
  }
  // 2、若任务执行完成，返回 null
  return null;
}

function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
  // ...
  do {
    try { // 这里使用 do while + try catch，防止执行过程中出错，能够重新执行
      workLoopConcurrent();
      break;
    } catch (thrownValue) {
      handleError(root, thrownValue);
    }
  } while (true);
  // ...
}

function workLoopConcurrent() {
  while (workInProgress !== null && !shouldYield()) { // unstable_shouldYield
    performUnitOfWork(workInProgress);
  }
}

shouldYield 方法则是 requestHostCallback 时间切片中提供的 shouldYieldToHost 方法，返回是否要将执行权交给浏览器：

function unstable_shouldYield() {
  const currentTime = getCurrentTime();
  advanceTimers(currentTime);
  return shouldYieldToHost(); // requestHostCallback 中的是否让出执行权方法
}

六、如何使用？

从上面我们得知，每一个调度任务会采用发起宏任务高频 5ms 方式去执行，如果 5ms 执行过后任务还未执行完毕，则会根据任务的过期时间，决定：

1. 任务过期，不会让出主线程，同步将任务执行完毕；
2. 任务未过期，先将主线程交还给浏览器，开启下一个宏任务再继续执行。

这段逻辑在 workLoop 中：

function workLoop(hasTimeRemaining, initialTime) {
  let currentTime = initialTime; // requestHostCallback 调度任务的开始时间
  currentTask = peek(taskQueue); // 取出任务队列的第一个要执行的任务
  while (currentTask !== null) {
    // 1、如果当前任务未过期，但时间切片用尽了，需要让出执行权给浏览器
    if (currentTask.expirationTime > currentTime && (!hasTimeRemaining || shouldYieldToHost())) {
      break;
    }

    // 2、否则，时间切片未用尽，或任务过期，同步执行完毕
    const callback = currentTask.callback;
    const didUserCallbackTimeout = currentTask.expirationTime <= currentTime; // 已经超时
    const continuationCallback = callback(didUserCallbackTimeout); // 执行 callback

    currentTime = getCurrentTime();
    // 如果 callback 任务执行完成会返回 null，若需要让出主线程但任务还未执行完成，则返回的是函数，详见下文 performConcurrentWorkOnRoot
    if (typeof continuationCallback === 'function') {
      currentTask.callback = continuationCallback; // 保存 callback，等待下次宏任务继续执行
    } else {
      pop(taskQueue); // 执行完任务，从堆中移除
    }
    currentTask = peek(taskQueue);
  }
}

下面，我们通过 Scheduler 调度一次任务、10个工作单元，来理解任务中断/续上场景。

6.1、任务已过期，同步执行中断任务

假设我们给任务安排的是 ImmediatePriority（-1） 立即过期优先级，在 Scheduler workLoop 中将以同步方式去循环执行，而源码中为 performWork 提供的 didUserCallbackTimeout 参数变得尤为重要，它可以防止任务执行出现死循环。

import { unstable_scheduleCallback, unstable_ImmediatePriority, unstable_UserBlockingPriority, unstable_shouldYield } from 'scheduler';

const clickBtn = () => {
  const work = {
    count: 10,
    priority: unstable_ImmediatePriority,
    timeout: false, // 是否已过期超时
  }

  // 4、执行任务单元
  const performUnitOfWork = () => {
    const start = Date.now();
    while (Date.now() - start < 2) {}
    work.count --;
  }

  // 3、循环执行任务单元
  const workLoop = () => {
    // 这里的 work.timeout 就是 didUserCallbackTimeout，避免出现死循环
    while (work.count > 0 && (!unstable_shouldYield() || work.timeout)) {
      performUnitOfWork();
    }
  }

  // 2、执行任务
  const performWork = (timeout: boolean) => {
    console.log('timeout: ', timeout);
    work.timeout = timeout;
    workLoop();
    if (work.count > 0) {
      console.log('需要先让出主线程给浏览器，本次更新任务剩余：', work.count);
      return performWork;
    }
    console.log('任务执行完成！');
  }

  console.log('开始执行任务！');
  // 1、调度
  unstable_scheduleCallback(
    work.priority, 
    performWork,
  );
}

return (
  <div>
    <button onClick={clickBtn}>点击按钮</button>
    <div id="container"></div>
  </div>
)

点击按钮，执行打印输出如下：

开始执行任务！
timeout:  true
任务执行完成！

可见，对于过期的任务，并未让出主线程，而是以同步方式一次性执行完成。

6.2、任务未过期，异步宏任务执行中断任务

如果我们给任务设置优先级为 UserBlockingPriority，它的过期时间为 250ms，对于这个场景来说直到 work 执行完毕任务都不会过期，因此会多次开启宏任务来继续执行中断任务。

我们把优先级改成 unstable_UserBlockingPriority，点击按钮打印输出如下：

最后

感谢阅读，如有不足之处，欢迎指出。

借鉴:
1. Scheduler
2. postMessage & Scheduler
3. 不用一行代码，搞懂React调度器原理