Go实现简单时间轮

简单时间轮

该时间轮是由切片外加链表实现

// TimeWheel can execute job after waiting given duration
type TimeWheel struct {  
   interval time.Duration  
   ticker   *time.Ticker  
   slots    []*list.List  
  
   timer             map[string]int  
   currentPos        int  
   slotNum           int  
   addTaskChannel    chan task  
   removeTaskChannel chan string  
   stopChannel       chan bool  
}  
  
type task struct {  
   delay  time.Duration  
   circle int  
   key    string  
   job    func()  
}  
  
// New creates a new time wheel//  
// 初始化时间轮，时间轮的执行间隙和插槽数量作为入参传入，返回对应结构体  
func New(interval time.Duration, slotNum int) *TimeWheel {  
   if interval <= 0 || slotNum <= 0 {  
      return nil  
   }  
   tw := &TimeWheel{  
      interval:          interval,  
      slots:             make([]*list.List, slotNum),  
      timer:             make(map[string]int),  
      currentPos:        0,  
      slotNum:           slotNum,  
      addTaskChannel:    make(chan task),  
      removeTaskChannel: make(chan string),  
      stopChannel:       make(chan bool),  
   }  
   tw.initSlots()  
  
   return tw  
}  
  
// 初始化插槽，每个槽放置一个链表  
func (tw *TimeWheel) initSlots() {  
   for i := 0; i < tw.slotNum; i++ {  
      tw.slots[i] = list.New()  
   }  
}  
  
// Start starts ticker for time wheel//  
// 定义一个定时器，将传入的时间间隔作为入参，返回的定时器赋值给时间轮的ticker字段  
func (tw *TimeWheel) Start() {  
   tw.ticker = time.NewTicker(tw.interval)  
   go tw.start()  
}  
  
// Stop stops the time wheelfunc (tw *TimeWheel) Stop() {  
   tw.stopChannel <- true  
}  
  
// AddJob add new job into pending queue//  
// 添加一个任务，将任务key、x个时间单位后、以及执行的任务函数作为task结构体添加到时间轮对象中  
func (tw *TimeWheel) AddJob(delay time.Duration, key string, job func()) {  
   if delay < 0 {  
      return  
   }  
   tw.addTaskChannel <- task{delay: delay, key: key, job: job}  
}  
  
// RemoveJob add remove job from pending queue// if job is done or not found, then nothing happened  
func (tw *TimeWheel) RemoveJob(key string) {  
   if key == "" {  
      return  
   }  
   tw.removeTaskChannel <- key  
}  
  
// 监听时间轮中所有的通道字段，根据其类型执行对应的函数  
func (tw *TimeWheel) start() {  
   for {  
      select {  
      case <-tw.ticker.C:  
         // 每一个时间间隔都会执行tickHandler函数  
         tw.tickHandler()  
      case task := <-tw.addTaskChannel:  
         // 添加任务  
         tw.addTask(&task)  
      case key := <-tw.removeTaskChannel:  
         // 删除任务  
         tw.removeTask(key)  
      case <-tw.stopChannel:  
         // 监听终止信号  
         tw.ticker.Stop()  
         return  
      }  
   }  
}  
  
// 根据当前时间轮指针获取其存放任务的链表，然后执行scanAndRunTask()函数，扫描属于当前时间轮的任务并执行它  
func (tw *TimeWheel) tickHandler() {  
   l := tw.slots[tw.currentPos]  
   tw.scanAndRunTask(l)  
   // 判断当前时间轮指针是否等于插槽的大小，如果等于则赋值0重新开始，否则自增  
   if tw.currentPos == tw.slotNum-1 {  
      tw.currentPos = 0  
   } else {  
      tw.currentPos++  
   }  
}  
  
// 扫描传入的链表中，判断是否有属于当前轮次的任务  
// 判断的依据是task结构体的circle字段，如果circle等于0，则代表属于当前轮次任务，执行  
func (tw *TimeWheel) scanAndRunTask(l *list.List) {  
   for e := l.Front(); e != nil; {  
      task := e.Value.(*task)  
      // 判断circle字段，等于0则代表属于当前轮次，执行  
      // > 0 的话则代表不属于当前轮次，跳过该任务  
      if task.circle > 0 {  
         task.circle--  
         e = e.Next()  
         continue  
      }  
  
      // task.circle == 0 才会执行以下代码  
      go func() {  
         // 另外起协程执行任务  
         // 如果执行任务过程中出现panic，使用recover()函数捕获，不影响当前主协程运行  
         defer func() {  
            if err := recover(); err != nil {  
               logger.Error(err)  
            }  
         }()  
         job := task.job  
         job()  
      }()  
      next := e.Next()  
      // 从任务链表中删除该结点  
      l.Remove(e)  
      if task.key != "" {  
         delete(tw.timer, task.key)  
      }  
      e = next  
   }  
}  
  
// 添加任务  
func (tw *TimeWheel) addTask(task *task) {  
   pos, circle := tw.getPositionAndCircle(task.delay)  
   task.circle = circle  
   // 将计算得到的位置和圈数设置后，追加到所对应槽位的链表尾部  
   tw.slots[pos].PushBack(task)  
  
   if task.key != "" {  
      tw.timer[task.key] = pos  
   }  
}  
  
// 根据任务的时间间隔计算当前任务在时间轮中的圈数以及格子位置  
func (tw *TimeWheel) getPositionAndCircle(d time.Duration) (pos int, circle int) {  
   // 获取该任务在在多久个时间间隔单位后执行  
   delaySeconds := int(d.Seconds())  
   intervalSeconds := int(tw.interval.Seconds())  
   // 获取到是第几圈  
   circle = int(delaySeconds / intervalSeconds / tw.slotNum)  
   // 计算出当前任务所出于轮中的哪一个槽  
   pos = int(tw.currentPos+delaySeconds/intervalSeconds) % tw.slotNum  
  
   return  
}  
  
// 删除任务  
func (tw *TimeWheel) removeTask(key string) {  
   // 根据key去查找要删除任务所在的时间轮槽位置  
   position, ok := tw.timer[key]  
   if !ok {  
      return  
   }  
   // 获取链表  
   l := tw.slots[position]  
   // 遍历链表，如果key== 删除任务key，则直接删除  
   for e := l.Front(); e != nil; {  
      task := e.Value.(*task)  
      if task.key == key {  
         delete(tw.timer, task.key)  
         l.Remove(e)  
      }  
  
      e = e.Next()  
   }  
}