利用时间轮实现定时器

func (tw *TimeWheel) tickHandler()  {
    //获取指针指向的槽
	l := tw.slots[tw.currentPos]
    //扫描当前槽中过期的定任务,时间复杂度为O(n)，并执行回调
	tw.scanAndRunTask(l)
    //如果当前指针已经到时间轮的末尾，则赋值=0
	if tw.currentPos == tw.slotNum - 1 {
		tw.currentPos = 0
	}else {
        //指针自增
		tw.currentPos ++
	}
}

func (tw *TimeWheel) scanAndRunTask(l *list.List)  {
	for e := l.Front(); e != nil; {
		task := e.Value.(*Task)
		if task.circle > 0 {
			task.circle --
			e = e.Next()
			continue
		}
        //执行定时任务. 每个任务都需要一个协程，需优化
		go task.job(task.data)
		next := e.Next()
        //删除任务
		l.Remove(e)
		delete(tw.taskMap, task.key)
		e = next
	}
}

package timewheel

import (
	"container/list"
	"sync/atomic"
	"time"
)
// Job 延时任务回调函数
type Job func(interface{})

// Task 延时任务
type Task struct {
	delay  time.Duration // 延迟时间
	circle int           // 时间轮需要转动几圈
	key    int64   // 定时器唯一标识, 用于删除定时器
	data   interface{}      // 回调函数参数
	job    Job
}


type TimeWheel struct {
	id int64 //auto incr
	tick time.Duration//in millisecond 时间轮上时间单位 ms
	slots    []*list.List // 时间轮槽
	currentPos        int              // 当前指针指向哪一个槽
	taskMap   map[int64]int
	slotNum           int              // 槽数量
	addTaskChannel    chan *Task        // 新增任务channel
	removeTaskChannel chan int64 // 删除任务channel
	stopChannel       chan bool        // 停止定时器channel
	ticker *time.Ticker
}

func NewTimeWheel(tick time.Duration, slotNum int) *TimeWheel{
	tw := &TimeWheel{
		tick:              tick,
		currentPos:        0,
		slots: make([]*list.List, slotNum),
		taskMap: make(map[int64]int),
		slotNum:           slotNum,
		addTaskChannel:    make(chan *Task),
		removeTaskChannel: make(chan int64),
		stopChannel:       make(chan bool),
	}
	tw.initSlots()
	return tw
}

func (tw *TimeWheel) initSlots()  {
	for i := 0; i < tw.slotNum; i++ {
		tw.slots[i] = list.New()
	}
}

// 扫描链表中过期定时器, 并执行回调函数 时间复杂度为O(n)
func (tw *TimeWheel) scanAndRunTask(l *list.List)  {
	for e := l.Front(); e != nil; {
		task := e.Value.(*Task)
		if task.circle > 0 {
			task.circle --
			e = e.Next()
			continue
		}
		go task.job(task.data)
		next := e.Next()
		l.Remove(e)
		delete(tw.taskMap, task.key)
		e = next
	}
}

func (tw *TimeWheel) tickHandler()  {
	l := tw.slots[tw.currentPos]
	tw.scanAndRunTask(l)
	if tw.currentPos == tw.slotNum - 1 {
		tw.currentPos = 0
	}else {
		tw.currentPos ++
	}
}

func (tw *TimeWheel) getRoundCount(d time.Duration) (int,int) {

	circle := int(d/tw.tick)/tw.slotNum
	pos := (tw.currentPos +  int(d/tw.tick)) % tw.slotNum
	return circle, pos
}

func (tw *TimeWheel) addTask(task *Task)  {
	circle, pos := tw.getRoundCount(task.delay)
	task.circle = circle

	tw.slots[pos].PushBack(task)
	tw.taskMap[task.key] = pos
}

func (tw *TimeWheel) removeTask(key int64)  {
	pos, ok := tw.taskMap[key]
	if !ok {
		return
	}
	l := tw.slots[pos]
	for e := l.Front(); e != nil; {
		task := e.Value.(*Task)
		if task.key == key {
			l.Remove(e)
			delete(tw.taskMap, key)
			return
		}
		e = e.Next()
	}
}

func (tw *TimeWheel) removeAllTask()  {
	for _, l := range tw.slots {
		for e := l.Front(); e != nil; {
			task := e.Value.(*Task)
			if task != nil {

				next := e.Next()
				l.Remove(e)
				if task.key != 0{
					delete(tw.taskMap, task.key)
				}
				e = next
			}else {
				return
			}
		}
	}
}

func (tw *TimeWheel) Start()  {
	tw.ticker = time.NewTicker(tw.tick)
	go func() {
		for {
			select {
			case <-tw.ticker.C:
				tw.tickHandler()
			case task := <-tw.addTaskChannel:
				tw.addTask(task)
			case key := <-tw.removeTaskChannel:
				tw.removeTask(key)
			case <-tw.stopChannel:
				tw.ticker.Stop()
				tw.removeAllTask()
				return
			}
		}
	}()
}
func (tw *TimeWheel) Stop()  {
	tw.stopChannel <- true
}
func (tw *TimeWheel) AddTimer(delay time.Duration, data interface{}, job Job) int64 {
	id := atomic.AddInt64(&tw.id, 1)
	tw.addTaskChannel <- &Task{
		key: id,
		delay:  delay,
		data:   data,
		job:    job,
	}
	return id
}

func (tw *TimeWheel) RemoveTimer(key int64)  {
	tw.removeTaskChannel <- key
}

package timewheel

import (
	"fmt"
	"testing"
	"time"
)

func TestTimeWheel_Start(t *testing.T) {
	tw := NewTimeWheel(time.Millisecond, 3600)
	tw.Start()

	id := tw.AddTimer(10*time.Millisecond, "123", func(i interface{}) {
		fmt.Println(i)
	})
	fmt.Println("timer id: ", id)

	time.Sleep(100*time.Millisecond)
	tw.Stop()
}

func (tw *TimeWheel) AfterFunc(d time.Duration, f func()) *Timer {
	t := &Timer{
		expiration: time.Now().UTC().Add(d).UnixMilli(),
		task:       f,
	}
	tw.addOrRun(t)
	return t
}

func (tw *TimeWheel) addOrRun(t *Timer) {
	if !tw.add(t) {
		workerID := t.expiration % tw.WorkPool.WorkerPoolSize
		//将请求消息发送给任务队列
		tw.WorkPool.TaskQueue[workerID] <- t.task
	}
}

func (tw *TimeWheel) add(t *Timer) bool {
	currentTime := atomic.LoadInt64(&tw.currentTime)
	if t.expiration < currentTime + tw.tickMs {
		return false
	}else if t.expiration < currentTime + tw.wheelTime {
		virtualID := t.expiration / tw.tickMs  //需要多少滴答数
		b := tw.bucket[virtualID%tw.wheelSize] //pos = 所需滴答数 % wheelSize
		b.Add(t)

		b.SetExpiration(virtualID * tw.tickMs)
	}else {
		overflowWheel := atomic.LoadPointer(&tw.overflowWheel)
		if overflowWheel == nil {
			atomic.CompareAndSwapPointer(
				&tw.overflowWheel,
				nil,
				unsafe.Pointer(newTimingWheel(tw.wheelTime, tw.wheelSize, currentTime)),
				)
			overflowWheel = atomic.LoadPointer(&tw.overflowWheel)
		}
		//递归添加到下一级定时器中
		(*TimeWheel)(overflowWheel).add(t)
	}

	return true
}

package timewheel

import (
	"pro2d/src/common"
	"pro2d/src/components/workpool"
	"sync/atomic"
	"time"
	"unsafe"
)

type TimeWheel struct {
	ticker    *time.Ticker
	tickMs    int64 //一滴答的时间 1ms 可以自定义 我们这里选择使用1ms
	wheelSize int64
	startMs   int64 //开始时间 in millisecond
	endMs     int64
	wheelTime int64 //跑完一圈所需时间
	level 	  int64 //层级

	//时间刻度 列表
	bucket []*bucket
	currentTime int64 //当前时间 in millisecond
	prevflowWheel unsafe.Pointer // type: *TimingWheel
	overflowWheel unsafe.Pointer // type: *TimingWheel
	exitC     chan struct{}

	WorkPool *workpool.WorkPool
}

func NewTimeWheel(tick time.Duration, wheelSize int64) *TimeWheel {
	//转化为毫秒
	tickMs := int64(tick / time.Millisecond)
	//如果小于零
	if tickMs <=0 {
		panic("tick must be greater than or equal to 1 ms")
	}

	startMs := time.Now().UnixMilli() //ms

	workpool := workpool.NewWorkPool(common.WorkerPoolSize, common.MaxTaskPerWorker)
	return newTimingWheel(tickMs, wheelSize, startMs, 0, nil, workpool)
}

func newTimingWheel(tick, wheelSize int64, start, level int64, prev *TimeWheel, pool *workpool.WorkPool) *TimeWheel {
	buckets := make([]*bucket, wheelSize)
	for i := range buckets {
		buckets[i] = newBucket()
	}

	return &TimeWheel{
		tickMs:      tick,
		wheelSize:   wheelSize,
		startMs:     start,
		endMs:       wheelSize * tick + start,
		wheelTime:   wheelSize * tick,
		bucket:      buckets,
		currentTime: truncate(start, tick),
		exitC:       make(chan struct{}),
		WorkPool: pool,

		prevflowWheel: unsafe.Pointer(prev),
		level: level,
	}
}

func truncate(dst, m int64) int64 {
	return dst - dst%m
}

func (tw *TimeWheel) add(t *Timer) bool {
	currentTime := atomic.LoadInt64(&tw.currentTime)
	if t.expiration < currentTime + tw.tickMs {
		return false
	}else if t.expiration < currentTime + tw.wheelTime {
		virtualID := t.expiration / tw.tickMs  //需要多少滴答数
		b := tw.bucket[virtualID%tw.wheelSize] //pos = 所需滴答数 % wheelSize
		b.Add(t)

		b.SetExpiration(virtualID * tw.tickMs)
	}else {
		overflowWheel := atomic.LoadPointer(&tw.overflowWheel)
		if overflowWheel == nil {
			level := atomic.LoadInt64(&tw.level) + 1
			atomic.CompareAndSwapPointer(
				&tw.overflowWheel,
				nil,
				unsafe.Pointer(newTimingWheel(tw.wheelTime, tw.wheelSize, currentTime, level, tw , tw.WorkPool)),
				)
			overflowWheel = atomic.LoadPointer(&tw.overflowWheel)
		}
		//递归添加到下一级定时器中
		(*TimeWheel)(overflowWheel).add(t)
	}

	return true
}

func (tw *TimeWheel) addOrRun(t *Timer) {
	if !tw.add(t) {
		workerID := t.expiration % tw.WorkPool.WorkerPoolSize
		//将请求消息发送给任务队列
		tw.WorkPool.TaskQueue[workerID] <- t.task
	}
}

//拨动时钟
func (tw *TimeWheel) advanceClock(expiration int64) {
	level := atomic.LoadInt64(&tw.level)
	currentTime := truncate(expiration, tw.tickMs)
	atomic.StoreInt64(&tw.currentTime, currentTime)

	if level == 0 {
		virtualID := expiration / tw.tickMs    //需要多少滴答数
		b := tw.bucket[virtualID%tw.wheelSize] //pos = 所需滴答数 % wheelSize
		b.Flush(tw.addOrRun)
	} else {
		prevflowWheel := atomic.LoadPointer(&tw.prevflowWheel)
		if prevflowWheel != nil {
			virtualID := expiration / tw.tickMs    //需要多少滴答数
			b := tw.bucket[virtualID%tw.wheelSize] //pos = 所需滴答数 % wheelSize
			b.Flush((*TimeWheel)(prevflowWheel).addOrRun)
		}
	}

	//如果基础的时钟指针转完了一圈，则递归拨动下一级时钟
	if currentTime >= tw.endMs {
		atomic.StoreInt64(&tw.startMs, currentTime)
		atomic.StoreInt64(&tw.endMs, currentTime + tw.wheelTime)

		overflowWheel := atomic.LoadPointer(&tw.overflowWheel)
		if overflowWheel != nil {
			(*TimeWheel)(overflowWheel).advanceClock(currentTime)
		}
	}
}


func (tw *TimeWheel) AfterFunc(d time.Duration, f func()) *Timer {
	t := &Timer{
		expiration: time.Now().UTC().Add(d).UnixMilli(),
		task:       f,
	}
	tw.addOrRun(t)
	return t
}

func (tw *TimeWheel) Start()  {
	tw.ticker = time.NewTicker(time.Duration(tw.tickMs) * time.Millisecond)
	tw.WorkPool.StartWorkerPool()

	go func() {
		for {
			select {
			case t := <- tw.ticker.C:
				tw.advanceClock(t.UnixMilli())
			case <- tw.exitC:
				return
			}
		}
	}()
}

func (tw *TimeWheel) Stop()  {
	tw.exitC <- struct{}{}
}

package timewheel

import (
	"container/list"
	"sync"
	"sync/atomic"
	"unsafe"
)
type bucket struct {
	//过期时间
	expiration int64

	mu sync.Mutex
	//相同过期时间的任务队列
	timers *list.List
}

func newBucket() *bucket  {
	return &bucket{
		expiration: -1,
		mu:      sync.Mutex{},
		timers:  list.New(),
	}
}

func (b *bucket) SetExpiration(expiration int64)  {
	atomic.AddInt64(&b.expiration, expiration)
}

func (b *bucket) Add(t *Timer) {
	b.mu.Lock()
	defer b.mu.Unlock()

	e := b.timers.PushBack(t)
	t.setBucket(b)
	t.element = e
}

func (b *bucket) Flush(reinsert func(*Timer)) {
	b.mu.Lock()
	defer b.mu.Unlock()

	for e := b.timers.Front(); e != nil; {
		next := e.Next()
		t := e.Value.(*Timer)
		b.remove(t)

		reinsert(t)
		e = next
	}
}

func (b *bucket) remove(t *Timer) bool {
	if t.getBucket() != b {
		return false
	}
	b.timers.Remove(t.element)
	t.setBucket(nil)
	t.element = nil
	return true
}

func (b *bucket) Remove(t *Timer) bool {
	b.mu.Lock()
	defer b.mu.Unlock()

	return b.remove(t)
}

type Timer struct {
	expiration int64
	//要被执行的任务
	task func()

	b unsafe.Pointer
	element *list.Element
}

func (t *Timer) setBucket(b *bucket)  {
	atomic.StorePointer(&t.b, unsafe.Pointer(b))
}

func (t *Timer) getBucket() *bucket {
	return (*bucket)(atomic.LoadPointer(&t.b))
}

package timewheel

import (
	"fmt"
	"testing"
	"time"
)
var tw *TimeWheel

func Add()  {
	fmt.Println("ADD : 123456")
	tw.AfterFunc(6*time.Second, Add)
}

func Add1()  {
	fmt.Println("GET : 78901112")
	tw.AfterFunc(9*time.Second, Add1)
}

func TestTimeWheel_AfterFunc(t *testing.T) {

	tw = NewTimeWheel(time.Second, 5)
	tw.Start()
	defer tw.Stop()


	Add()
	Add1()
	time.Sleep(time.Second * 200)
}

package workpool

type Job func()

type WorkPool struct {
	WorkerPoolSize int64
	MaxTaskPerWorker int64
	TaskQueue []chan Job
}

func NewWorkPool(poolSize, maxTaskSize int64) *WorkPool {
	return &WorkPool{
		WorkerPoolSize: poolSize,
		MaxTaskPerWorker: maxTaskSize,
		TaskQueue: make([]chan Job, poolSize),
	}
}

//StartOneWorker 启动一个Worker工作流程
func (wp *WorkPool) StartOneWorker(workerID int, taskQueue chan Job) {
	//不断的等待队列中的消息
	for {
		select {
		case job := <-taskQueue:
			_ = workerID
			job()
		}
	}
}

func (wp *WorkPool) StartWorkerPool() {
	//遍历需要启动worker的数量，依此启动
	for i := 0; i < int(wp.WorkerPoolSize); i++ {
		//一个worker被启动
		//给当前worker对应的任务队列开辟空间
		wp.TaskQueue[i] = make(chan Job,  wp.MaxTaskPerWorker)
		//启动当前Worker，阻塞的等待对应的任务队列是否有消息传递进来
		go wp.StartOneWorker(i, wp.TaskQueue[i])
	}
}