go实现时间轮（timewheel)算法（简易版）

概念

时间轮算法（Time Wheel Algorithm）是一种常用于处理定时任务的数据结构和算法。它可以高效地管理和执行一系列定时任务，特别是在需要频繁添加、删除和执行任务时表现良好。时间轮算法主要用于实现定时器功能，比如在操作系统、网络通信、任务调度等领域。

基本思想

时间轮算法的基本思想是将时间划分为若干个固定大小的时间片（也称为槽位），并且维护一个环形的时间轮数据结构，每个槽位对应一个时间片。时间轮按照时间递增的方向不断前进，当前时间指针指向当前时间所在的槽位。当时间轮的指针指向某个槽位时，执行该槽位上的定时任务。

代码实现

package timewheel

import (
    "sync"
    "time"
)

// 默认一秒作为时间轮间隔
type timewheel struct {
    slot map[int64]*[]func()
    lock sync.RWMutex
}

// 利用时间戳作为key，相同时间戳进入同一个桶中，过期的桶则进行删除
// todo 想要优化可以把map换成定长list(按长度取模放入list)，后跟基于[]func列表实现的优先级队列
/*
结构体如下所示：优先级队列可以搜索一下go怎么实现的，很简单
type funcPrio []func()

type timewheel struct {
    wheel []funcPrio
    size int32
    lock sync.RWMutex
}
*/

func NewTimewheel() *timewheel {
    return &timewheel{slot: map[int64]*[]func(){}, lock: sync.RWMutex{}}
}

func (t *timewheel) AddTask(runTm int64, fn func()) bool {
    t.lock.Lock()
    if runTm < time.Now().Unix() {
       t.lock.Unlock()
       return false
    }
    if list, ok := t.slot[runTm]; ok {
       *list = append(*list, fn)
    } else {
       list := []func(){fn}
       t.slot[runTm] = &list
    }
    t.lock.Unlock()
    return true
}

// 套一层，减少go协程执行造成时间轮之间间隔可能操作一秒,不确定是否可以控制住时间间隔
func (t *timewheel) Run() {
    lasttm := time.Now().Unix()
    for {
       now := time.Now().Unix()
       if now-lasttm == 1 {
          lasttm = now
          go t.run(now)
       }
    }
}

func (t *timewheel) run(tm int64) {
    //for {
    //执行时间间隔可能有微差，但大致应该控制在1s之内，暂时我也不清楚怎么控制
    //time.Sleep(time.Second)
    //tm := time.Now().Unix()
    t.lock.Lock()
    if list, ok := t.slot[tm]; ok {
       for _, fn := range *list {
          go fn()
       }
       //自动清理list
       delete(t.slot, tm)
    }
    t.lock.Unlock()
    //}
}

测试代码

package timewheel

import (
    "math/rand"
    "testing"
    "time"
)

func TestRunTimewheel(t *testing.T) {
    tw := NewTimewheel()
    go tw.Run()
    ls := []int64{}
    for i := 0; i < 10; i++ {
       tm := rand.Int63n(int64(5)) + time.Now().Unix() + 1
       //println("add task:", tm)
       ls = append(ls, tm)
       fg := tw.AddTask(tm, func() {
          println("exec hello ", rand.Intn(20)+i)
       })
       if !fg {
          println("addtask false")
       }
    }
    time.Sleep(time.Second * 10)
}

执行结果

=== RUN   TestRunTimewheel
exec hello  6
exec hello  10
exec hello  13
exec hello  13
exec hello  9
exec hello  6
exec hello  9
exec hello  20
exec hello  11
exec hello  4
--- PASS: TestRunTimewheel (10.00s)
PASS

结论

总体实现比较简单，当然也可以进行抽象化处理，写一个优先级队列，但感觉基本思路都差不，还是简单点来吧，能理解思路就可以了