字节面试题：实现任务调度器(Scheduler)

突然想起去年面试字节时的一道题，题面大概是这样：

实现一个任务调度器 scheduler
只能同时执行 n 个任务
当有 n 个任务正在执行时，后安排的任务要排队

当时很久没面试了，紧张到冒汗😓，只写出一半就被叫停了；

今天辗转反侧睡不着，突然想起这道题，就给他写了吧 :)

实现

大致思路是：

定义 Scheduler 类：指定同时执行的任务数 n
有 schedule 方法：安排一个任务(task)，如果执行中的达到 n 则入队，否则直接执行
包装器 wrapper：包装 task，使执行 task 完毕后检查是否要出队，即执行新任务

伪码如下：

typescript 复制代码

class Scheduler {
  // 计数当前执行中的数量
  private runningTask: number
  // 最大并发
  private capacity: number
  // 等待队列
  private pendingQueue: (() => Promise<any>)[]

  constructor(capacity: number) {
    this.capacity = capacity
    this.runningTask = 0
    this.pendingQueue = []
  }

  private wrapTask<T>(task: () => Promise<T>): () => Promise<void> {
    return async () => {
      // 执行 task
      // 判断 pending 队列是否空
      // 否，出队下一个 task
      // 是，更新计数
    }
  }

  public schedule<T>(task: () => Promise<T>): void {
    // 包装 task，得到 wrappedTask
    // 判断计数是否达到 n
    // 否，则直接执行 wrappedTask
    // 是，则入队
  }
}

然后，让我们来实现它：

typescript 复制代码

class Scheduler {
  // 计数当前执行中的数量
  private runningTask: number
  // 最大并发
  private capacity: number
  // 等待队列
  private pendingQueue: (() => Promise<any>)[]

  constructor(capacity: number) {
    this.capacity = capacity
    this.runningTask = 0
    this.pendingQueue = []
  }

  private wrapTask<T>(task: () => Promise<T>): () => Promise<void> {
    return async () => {
      // 执行 task
      await task()
      // 判断 pending 队列是否空
      if (this.pendingQueue.length > 0) {
        // 否，出队下一个 task
        const nextWrappedTask = this.pendingQueue.shift()!
        nextWrappedTask()
      } else {
        // 是，更新计数
        this.runningTask--
      }
    }
  }

  public schedule<T>(task: () => Promise<T>): void {
    // 包装 task，得到 wrappedTask
    const wrappedTask = this.wrapTask(task)
    // 判断计数是否达到 n
    if (this.runningTask < this.capacity) {
      // 否，则直接执行 wrappedTask
      wrappedTask()
    } else {
      // 是，则入队
      this.pendingQueue.push(wrappedTask)
    }
  }
}

进阶

这里还可以进一步：要求 schedule 返回 Task 的执行结果，怎么写？

EZ，我们只要：让 schedule 返回一个 Promise，并在 wrappedTask 中 resolve 结果 就行了，代码如下：

PS: 这里总结一个小技巧"传递 resolve/reject 来实现异步返回值"

typescript 复制代码

class Scheduler {
  private runningTask: number
  private capacity: number
  private pendingQueue: (() => Promise<any>)[]

  constructor(capacity: number) {
    this.capacity = capacity
    this.runningTask = 0
    this.pendingQueue = []
  }

  private wrapTask<T>(task: () => Promise<T>, resolve: (value: T) => void, reject: (error: any) => void): () => Promise<void> {
    return async () => {
      try {
        const res = await task()
        resolve(res)
        const nextTask = this.pendingQueue.shift()
        if (nextTask) {
          nextTask()
        } else {
          this.runningTask--
        }
      } catch (e) {
        reject(e)
      }
    }
  }

  public async schedule<T>(task: () => Promise<T>): Promise<T> {
    return (new Promise<T>((resolve, reject) => {
      const wrappedTask = this.wrapTask(task, resolve, reject)
      if (this.runningTask < this.capacity) {
        this.runningTask++
        wrappedTask()
      } else {
        this.pendingQueue.push(wrappedTask)
      }
    }))
  }
}