dashboard/pkg/timewheel/timewheel.go

package timewheel

import (
	"container/list"
	"fmt"
	"math"
	"sync"
	"time"
)

type taskSlice struct {
	id       string
	pos      int
	cycle    int
	rawCycle int
	flower   time.Duration
	mode     TimeMode
	task     func()
}

type TimeWheel struct {
	sync.Once
	curtSlot    int
	interval    time.Duration
	ticker      *time.Ticker
	slots       []*list.List
	taskMap     map[string]*list.Element
	tasksMap    map[string][]int
	stopC       chan struct{}
	addTaskC    chan *taskSlice
	removeTaskC chan string
}

type TimeMode int

const (
	TimerTypeOnce TimeMode = iota
	TimeTypeLoop
)

func New(slots int, interval time.Duration) *TimeWheel {
	if slots <= 0 {
		slots = 10
	}

	if interval <= 0 {
		interval = time.Second
	}

	t := &TimeWheel{
		interval:    interval,
		ticker:      time.NewTicker(interval),
		slots:       make([]*list.List, 0, slots),
		taskMap:     make(map[string]*list.Element),
		stopC:       make(chan struct{}),
		addTaskC:    make(chan *taskSlice),
		removeTaskC: make(chan string),
		tasksMap:    make(map[string][]int),
	}

	for i := 0; i < slots; i++ {
		t.slots = append(t.slots, list.New())
	}

	go t.run()

	return t
}

func (t *TimeWheel) AddTask(id string, mode TimeMode, task func(), flower time.Duration) {
	cycle, pos := t.getSlotPosAndCycle(flower)

	t.addTaskC <- &taskSlice{
		id:       id,
		task:     task,
		cycle:    cycle,
		mode:     mode,
		rawCycle: cycle,
		pos:      pos,
		flower:   flower,
	}
}

func (t *TimeWheel) RemoveTask(key string) {
	t.removeTaskC <- key
}

func (t *TimeWheel) Stop() {
	t.Do(func() {
		t.ticker.Stop()
		close(t.stopC)
	})
}

func (t *TimeWheel) getSlotPosAndCycle(flower time.Duration) (int, int) {
	delay := time.Until(time.Now().Add(flower))
	cycle := delay / (t.interval * time.Duration(len(t.slots)))
	pos := (t.curtSlot + int(delay/t.interval)) % len(t.slots)

	return int(cycle), pos
}

func (t *TimeWheel) getCyclePos(task *taskSlice) []int {
	if task.cycle == 0 && task.mode == TimeTypeLoop {
		var listc []int
		dur := int(math.Ceil(float64(float64(task.flower) / float64(t.interval))))
		start := task.pos
		for i := 0; i < int(math.Ceil(float64(len(t.slots))/float64(dur))); i++ {
			listc = append(listc, start%len(t.slots))
			start += dur
		}

		return listc
	}

	return nil
}

func (t *TimeWheel) run() {
	for {
		select {
		case <-t.stopC:
			return
		case <-t.ticker.C:
			t.tick()
		case task := <-t.addTaskC:
			t.addCycelTask(task)
		case key := <-t.removeTaskC:
			t.removeCycelTask(key)
		}
	}
}

func (t *TimeWheel) tick() {
	list := t.slots[t.curtSlot]
	defer t.updateCurtSlot()

	t.execute(list)
}

func (t *TimeWheel) updateCurtSlot() {
	t.curtSlot = (t.curtSlot + 1) % len(t.slots)
}

func (t *TimeWheel) execute(l *list.List) {
	for e := l.Front(); e != nil; {
		event := e.Value.(*taskSlice)
		if event.cycle > 0 {
			event.cycle--
			e = e.Next()
			continue
		}

		go func() {
			defer func() {
				if err := recover(); err != nil {

				}
			}()

			event.task()
		}()

		if event.mode == TimerTypeOnce {
			next := e.Next()
			l.Remove(e)
			delete(t.taskMap, event.id)
			e = next
		} else {
			e = e.Next()
			event.cycle = event.rawCycle
		}
	}
}

func (t *TimeWheel) addCycelTask(task *taskSlice) {
	t.removeCycelTask(task.id)

	poss := t.getCyclePos(task)
	fmt.Println(poss)
	if len(poss) > 0 {
		t.tasksMap[task.id] = poss
		for _, index := range poss {
			internalTask := &taskSlice{
				id:       fmt.Sprintf("%s_%d", task.id, index),
				pos:      index,
				cycle:    task.cycle,
				rawCycle: task.rawCycle,
				mode:     task.mode,
				task:     task.task,
			}

			t.addTask(internalTask)
		}

		return
	}

	t.addTask(task)
}

func (t *TimeWheel) removeCycelTask(key string) {
	poss := t.tasksMap[key]
	if len(poss) > 0 {
		for _, index := range poss {
			id := fmt.Sprintf("%s_%d", key, index)
			t.removeTask(id)
		}

		return
	}

	t.removeTask(key)
}

func (t *TimeWheel) addTask(task *taskSlice) {
	if _, ok := t.taskMap[task.id]; ok {
		t.removeTask(task.id)
	}

	list := t.slots[task.pos]
	etask := list.PushBack(task)
	t.taskMap[task.id] = etask
}

func (t *TimeWheel) removeTask(key string) {
	etask, ok := t.taskMap[key]
	if !ok {
		return
	}

	delete(t.taskMap, key)

	task := etask.Value.(*taskSlice)
	_ = t.slots[task.pos].Remove(etask)
}