Kafka 延迟操作模块（一）：TimingWheel时间轮

kafka使用的是层级时间时间轮处理实现延迟功能。例如我们的手表，手表由时针、分针和秒针组成，它们各自有独立的刻度，这就是典型的分层时间轮。和手表不一样的是，Kafka 自己有专门的术语。在 Kafka 中，手表中的“一格”叫“一个桶（Bucket）”，而“推进”对应于 Kafka 中的“滴答”，也就是 tick。

Kafka 的分层时间轮算法在实现上主要涉及 TimingWheel、TimerTaskList、TimerTaskEntry，以及 TimerTask 这 4 个类，各个类的作用说明如下：

• TimerTask ：特质类型，用于描述延时任务。它是一个 Runnable 类，Kafka 使用一个单独线程异步添加延时请求到时间轮。
• TimerTaskList ：时间格，采用环形双向链表实现，记录位于同一个时间格中的延时任务。提供 O(1) 时间复杂度的请求插入和删除。
• TimerTaskEntry ：时间格链表中的一个结点，是对延时任务 TimerTask 的封装。
• TimingWheel ：建模时间轮类型，采用定长数组记录放置时间格。统一管理下辖的所有 Bucket 以及定时任务。

TimerTask 类
       

 TimerTask 建模的是 Kafka 中的定时任务，它的定义如下：

trait TimerTask extends Runnable {
  // 表示这个定时任务的超时时间
  // 通常是request.timeout.ms参数值
  // 每个TimerTask实例关联一个TimerTaskEntry
  // 就是说每个定时任务需要知道它在哪个Bucket链表下的哪个链表元素上
  val delayMs: Long // timestamp in millisecond
  // timerTaskEntry 用于封装当前延时任务并记录到时间格中，属于延时任务与时间格之间建立关系的桥梁。
  private[this] var timerTaskEntry: TimerTaskEntry = null
  // 取消定时任务，原理就是将关联的timerTaskEntry置空
  def cancel(): Unit = {
    synchronized {
      if (timerTaskEntry != null) timerTaskEntry.remove()
      timerTaskEntry = null
    }
  }
  // 关联timerTaskEntry，原理是给timerTaskEntry字段赋值
  private[timer] def setTimerTaskEntry(entry: TimerTaskEntry): Unit = {
    synchronized {
      // if this timerTask is already held by an existing timer task entry,
      // we will remove such an entry first.
      if (timerTaskEntry != null && timerTaskEntry != entry)
        timerTaskEntry.remove()

      timerTaskEntry = entry
    }
  }
  // 获取关联的timerTaskEntry实例
  private[timer] def getTimerTaskEntry(): TimerTaskEntry = {
    timerTaskEntry
  }
}

TimerTaskEntry

        TimerTaskEntry 表征的是 Bucket 链表下的一个元素。本质上是一个链表结点的定义，其字段定义了结点的前置和后置指针，以及所属的时间格.TimerTaskEntry 类对象在被构造时会建立延时任务timerTask与结点之间的映射关系，并提供了获取延时任务取消状态，以及移除对应延时任务的操作。
类字段定义如下：

private[timer] class TimerTaskEntry(val timerTask: TimerTask,
                                    val expirationMs: Long // 定义了定时任务的过期时间
                                    // 假设有个 PRODUCE 请求在当前时间 1 点钟被发送到 Broker，超时时间是 30 秒，
                                    // 那么，该请求必须在 1 点 30 秒之前完成，否则将被视为超时。
                                    // 这里的 1 点 30 秒，就是 expirationMs 值。
                                   ) extends Ordered[TimerTaskEntry] {

  @volatile
  var list: TimerTaskList = null   // 绑定的Bucket链表实例
  var next: TimerTaskEntry = null  // next指针
  var prev: TimerTaskEntry = null  // prev指针

  // if this timerTask is already held by an existing timer task entry,
  // setTimerTaskEntry will remove it.
  // 关联给定的定时任务
  if (timerTask != null) timerTask.setTimerTaskEntry(this)

  // 关联定时任务是否已经被取消了
  def cancelled: Boolean = {
    timerTask.getTimerTaskEntry != this
  }
  // 从Bucket链表中移除自己
  def remove(): Unit = {
    var currentList = list
    // If remove is called when another thread is moving the entry from a task entry list to another,
    // this may fail to remove the entry due to the change of value of list. Thus, we retry until the list becomes null.
    // In a rare case, this thread sees null and exits the loop, but the other thread insert the entry to another list later.
    while (currentList != null) {
      currentList.remove(this)
      currentList = list
    }
  }

  override def compare(that: TimerTaskEntry): Int = {
    this.expirationMs compare that.expirationMs
  }
}

TimerTaskList 类
       

        TimerTaskList 描述了时间轮的一格（即时间格），在实现上采用双向链表实现，用于封装位于特定时间区间范围内的所有的延时任务

private[timer] class TimerTaskList(taskCounter: AtomicInteger // 用于标识当前这个链表，即整个分层时间轮中的总定时任务数
                                  ) extends Delayed {

  // TimerTaskList forms a doubly linked cyclic list using a dummy root entry
  // root.next points to the head
  // root.prev points to the tail
  /** 根结点 */
  private[this] val root = new TimerTaskEntry(null, -1)
  root.next = root
  root.prev = root
  // 表示这个链表所在 Bucket 的过期时间戳。
  private[this] val expiration = new AtomicLong(-1L)
}

Getter 和 Setter 方法

  // 代码使用了 AtomicLong 的 CAS 方法 getAndSet 原子性地设置了过期时间戳，之后将新过期时间戳和旧值进行比较，看看是否不同，然后返回结果。
  // 因为，目前 Kafka 使用一个 DelayQueue 统一管理所有的 Bucket，也就是 TimerTaskList 对象。
  // 随着时钟不断向前推进，原有 Bucket 会不断地过期，然后失效。当这些 Bucket 失效后，源码会重用这些 Bucket。
  // 重用的方式就是重新设置 Bucket 的过期时间，并把它们加回到 DelayQueue 中。
  // 这里进行比较的目的，就是用来判断这个 Bucket 是否要被插入到 DelayQueue。
  def setExpiration(expirationMs: Long): Boolean = {
    expiration.getAndSet(expirationMs) != expirationMs
  }

  // Get the bucket's expiration time
  def getExpiration(): Long = {
    expiration.get()
  }

此外，TimerTaskList 类还提供了 add 和 remove 方法，分别实现将给定定时任务插入到链表、从链表中移除定时任务的逻辑。

add方法

  // 将给定定时任务插入到链表
  def add(timerTaskEntry: TimerTaskEntry): Unit = {
    var done = false
    while (!done) {
      // Remove the timer task entry if it is already in any other list
      // We do this outside of the sync block below to avoid deadlocking.
      // We may retry until timerTaskEntry.list becomes null.
      // 在添加之前尝试移除该定时任务，保证该任务没有在其他链表中
      timerTaskEntry.remove()

      synchronized {
        timerTaskEntry.synchronized {
          if (timerTaskEntry.list == null) {
            // put the timer task entry to the end of the list. (root.prev points to the tail entry)
            val tail = root.prev
            timerTaskEntry.next = root
            timerTaskEntry.prev = tail
            timerTaskEntry.list = this
            // 把timerTaskEntry添加到链表末尾
            tail.next = timerTaskEntry
            root.prev = timerTaskEntry
            taskCounter.incrementAndGet()
            done = true
          }
        }
      }
    }
  }

  // Remove the specified timer task entry from this list
  // 从链表中移除定时任务
  def remove(timerTaskEntry: TimerTaskEntry): Unit = {
    synchronized {
      timerTaskEntry.synchronized {
        if (timerTaskEntry.list eq this) {
          timerTaskEntry.next.prev = timerTaskEntry.prev
          timerTaskEntry.prev.next = timerTaskEntry.next
          timerTaskEntry.next = null
          timerTaskEntry.prev = null
          timerTaskEntry.list = null
          taskCounter.decrementAndGet()
        }
      }
    }
  }

flush 方法
        flush 方法是清空链表中的所有元素，并对每个元素执行指定的逻辑。该方法用于将高层次时间轮 Bucket 上的定时任务重新插入回低层次的 Bucket 中

  // flush 方法是清空链表中的所有元素，并对每个元素执行指定的逻辑。
  // 该方法用于将高层次时间轮 Bucket 上的定时任务重新插入回低层次的 Bucket 中。
  def flush(f: (TimerTaskEntry)=>Unit): Unit = {
    synchronized {
      // 找到链表第一个元素
      var head = root.next
      // 开始遍历链表
      while (head ne root) {
        // 移除遍历到的链表元素
        remove(head)
        // 执行传入参数f的逻辑
        f(head)
        head = root.next
      }
      // 清空过期时间设置
      expiration.set(-1L)
    }
  }

TimingWheel 类
       

分层时间轮 TimingWheel 的定义如下：

private[timer] class TimingWheel(tickMs: Long, // 滴答一次的时长，类似于手表的例子中向前推进一格的时间。对于秒针而言，tickMs 就是 1 秒。同理，分针是 1 分，时针是 1 小时。
                                 // 在 Kafka 中，第 1 层时间轮的 tickMs 被固定为 1 毫秒，也就是说，向前推进一格 Bucket 的时长是 1 毫秒。
                                 wheelSize: Int, // 每一层时间轮上的 Bucket 数量。第 1 层的 Bucket 数量是 20。
                                 startMs: Long, // 时间轮对象被创建时的起始时间戳。
                                 taskCounter: AtomicInteger, // 这一层时间轮上的总定时任务数。
                                 queue: DelayQueue[TimerTaskList] // 将所有 Bucket 按照过期时间排序的延迟队列。
                                 // 随着时间不断向前推进，Kafka 需要依靠这个队列获取那些已过期的 Bucket，并清除它们。
                                ) {
  // 这层时间轮总时长，等于滴答时长乘以 wheelSize。以第 1 层为例，interval 就是 20 毫秒。
  // 由于下一层时间轮的滴答时长就是上一层的总时长，因此，第 2 层的滴答时长就是 20 毫秒，总时长是 400 毫秒，以此类推。
  private[this] val interval = tickMs * wheelSize
  // 时间轮下的所有 Bucket 对象，也就是所有 TimerTaskList 对象。
  private[this] val buckets = Array.tabulate[TimerTaskList](wheelSize) { _ => new TimerTaskList(taskCounter) }
  // 当前时间戳，只是源码对它进行了一些微调整，将它设置成小于当前时间的最大滴答时长的整数倍。
  // 假设滴答时长是 20 毫秒，当前时间戳是 123 毫秒，那么，currentTime 会被调整为 120 毫秒。
  private[this] var currentTime = startMs - (startMs % tickMs) // rounding down to multiple of tickMs

  // overflowWheel can potentially be updated and read by two concurrent threads through add().
  // Therefore, it needs to be volatile due to the issue of Double-Checked Locking pattern with JVM
  // Kafka 是按需创建上层时间轮的。这也就是说，当有新的定时任务到达时，会尝试将其放入第 1 层时间轮。
  // 如果第 1 层的 interval 无法容纳定时任务的超时时间，就现场创建并配置好第 2 层时间轮，并再次尝试放入，
  // 如果依然无法容纳，那么，就再创建和配置第 3 层时间轮，
  // 以此类推，直到找到适合容纳该定时任务的第 N 层时间轮。
  @volatile private[this] var overflowWheel: TimingWheel = null
}

addOverflowWheel 方法

        每当需要一个新的上层时间轮时，代码就会调用 addOverflowWheel 方法。
        这个方法就是创建一个新的 TimingWheel 实例，也就是创建上层时间轮。所用的滴答时长等于下层时间轮总时长，而每层的轮子数都是相同的。创建完成之后，代码将新创建的实例赋值给 overflowWheel 字段。

  private[this] def addOverflowWheel(): Unit = {
    // 只有之前没有创建上层时间轮方法才会继续
    synchronized {
      if (overflowWheel == null) {
        // 创建新的TimingWheel实例
        // 滴答时长tickMs等于下层时间轮总时长
        // 每层的轮子数都是相同的
        overflowWheel = new TimingWheel(
          tickMs = interval,
          wheelSize = wheelSize,
          startMs = currentTime,
          taskCounter = taskCounter,
          queue
        )
      }
    }
  }

add方法

           用于往时间轮中添加延时任务，该方法接收一个 TimerTaskEntry 类型对象，即对延时任务 TimerTask 的封装。方法实现如下：

  def add(timerTaskEntry: TimerTaskEntry): Boolean = {
    // 获取定时任务的过期时间戳，就是这个定时任务过期时的时点。
    val expiration = timerTaskEntry.expirationMs
    // 如果该任务已然被取消了，则无需添加，直接返回
    if (timerTaskEntry.cancelled) {
      // Cancelled
      false
      // 如果该任务超时时间已过期
    } else if (expiration < currentTime + tickMs) {
      // Already expired
      false
    } else if (expiration < currentTime + interval) {
      // 如果该任务超时时间在本层时间轮覆盖时间范围内
      // Put in its own bucket
      val virtualId = expiration / tickMs
      // 计算要被放入到哪个Bucket中
      // 第 1 层的时间轮有 20 个 Bucket，每个滴答时长是 1 毫秒。那么，第 2 层时间轮的滴答时长应该就是 20 毫秒，总时长是 400 毫秒。
      // 第 2 层第 1 个 Bucket 的时间范围应该是[20，40)，第 2 个 Bucket 的时间范围是[40，60），
      // 依次类推。假设现在有个延时请求的超时时间戳是 237，那么，它就应该被插入到第 11 个 Bucket 中。
      val bucket = buckets((virtualId % wheelSize.toLong).toInt)
      // 添加到Bucket中
      // 在确定了目标 Bucket 序号之后，代码会将该定时任务添加到这个 Bucket 下，同时更新这个 Bucket 的过期时间戳。
      // 在刚刚的那个例子中，第 11 号 Bucket 的起始时间就应该是小于 237 的最大的 20 的倍数，即 220。
      bucket.add(timerTaskEntry)

      // Set the bucket expiration time
      // 如果这个 Bucket 是首次插入定时任务，那么，还同时要将这个 Bucket 加入到 DelayQueue 中，方便 Kafka 轻松地获取那些已过期 Bucket，并删除它们。
      // 设置Bucket过期时间
      // 如果该时间变更过，说明Bucket是新建或被重用，将其加回到DelayQueue
      if (bucket.setExpiration(virtualId * tickMs)) {
        // The bucket needs to be enqueued because it was an expired bucket
        // We only need to enqueue the bucket when its expiration time has changed, i.e. the wheel has advanced
        // and the previous buckets gets reused; further calls to set the expiration within the same wheel cycle
        // will pass in the same value and hence return false, thus the bucket with the same expiration will not
        // be enqueued multiple times.
        queue.offer(bucket)
      }
      true
    } else {
      // Out of the interval. Put it into the parent timer
      // 按需创建上层时间轮
      if (overflowWheel == null) addOverflowWheel()
      // 加入到上层时间轮中
      overflowWheel.add(timerTaskEntry)
    }
  }

advanceClock 方法

        用于推动当前时间轮指针

def advanceClock(timeMs: Long): Unit = {
  // 向前驱动到的时点要超过Bucket的时间范围，才是有意义的推进，否则什么都不做
  // 更新当前时间currentTime到下一个Bucket的起始时点
  if (timeMs >= currentTime + tickMs) {
    currentTime = timeMs - (timeMs % tickMs)
    // 同时尝试为上一层时间轮做向前推进动作
    if (overflowWheel != null) overflowWheel.advanceClock(currentTime)
  }
}