147 lines
2.3 KiB
Go
147 lines
2.3 KiB
Go
package timingwheel
|
|
|
|
import (
|
|
"github.com/siddontang/golib/log"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type TaskFunc func()
|
|
|
|
type bucket struct {
|
|
c chan struct{}
|
|
tasks []TaskFunc
|
|
}
|
|
|
|
const defaultTasksSize = 16
|
|
const defaultTaskPool = 4
|
|
|
|
type TimingWheel struct {
|
|
sync.Mutex
|
|
|
|
interval time.Duration
|
|
|
|
ticker *time.Ticker
|
|
quit chan struct{}
|
|
|
|
maxTimeout time.Duration
|
|
|
|
buckets []bucket
|
|
|
|
pos int
|
|
|
|
tasks chan []TaskFunc
|
|
}
|
|
|
|
func NewTimingWheel(interval time.Duration, buckets int) *TimingWheel {
|
|
w := new(TimingWheel)
|
|
|
|
w.interval = interval
|
|
|
|
w.quit = make(chan struct{})
|
|
w.pos = 0
|
|
|
|
w.maxTimeout = time.Duration(interval * (time.Duration(buckets)))
|
|
|
|
w.buckets = make([]bucket, buckets)
|
|
|
|
for i := range w.buckets {
|
|
w.buckets[i].c = make(chan struct{})
|
|
w.buckets[i].tasks = make([]TaskFunc, 0, defaultTasksSize)
|
|
}
|
|
|
|
w.ticker = time.NewTicker(interval)
|
|
go w.run()
|
|
|
|
w.tasks = make(chan []TaskFunc, 1024)
|
|
|
|
for i := 0; i < defaultTaskPool; i++ {
|
|
go w.taskPool()
|
|
}
|
|
|
|
return w
|
|
}
|
|
|
|
func (w *TimingWheel) Stop() {
|
|
close(w.quit)
|
|
}
|
|
|
|
func (w *TimingWheel) After(timeout time.Duration) <-chan struct{} {
|
|
if timeout >= w.maxTimeout {
|
|
panic("timeout too much, over maxtimeout")
|
|
}
|
|
|
|
w.Lock()
|
|
|
|
index := (w.pos + int(timeout/w.interval)) % len(w.buckets)
|
|
|
|
b := w.buckets[index].c
|
|
|
|
w.Unlock()
|
|
|
|
return b
|
|
}
|
|
|
|
func (w *TimingWheel) AddTask(timeout time.Duration, f TaskFunc) {
|
|
if timeout >= w.maxTimeout {
|
|
panic("timeout too much, over maxtimeout")
|
|
}
|
|
|
|
w.Lock()
|
|
|
|
index := (w.pos + int(timeout/w.interval)) % len(w.buckets)
|
|
|
|
w.buckets[index].tasks = append(w.buckets[index].tasks, f)
|
|
|
|
w.Unlock()
|
|
}
|
|
|
|
func (w *TimingWheel) run() {
|
|
for {
|
|
select {
|
|
case <-w.ticker.C:
|
|
w.onTicker()
|
|
case <-w.quit:
|
|
w.ticker.Stop()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *TimingWheel) onTicker() {
|
|
w.Lock()
|
|
|
|
lastC := w.buckets[w.pos].c
|
|
tasks := w.buckets[w.pos].tasks
|
|
|
|
w.buckets[w.pos].c = make(chan struct{})
|
|
w.buckets[w.pos].tasks = w.buckets[w.pos].tasks[0:0:defaultTasksSize]
|
|
|
|
w.pos = (w.pos + 1) % len(w.buckets)
|
|
|
|
w.Unlock()
|
|
|
|
close(lastC)
|
|
|
|
w.tasks <- tasks
|
|
}
|
|
|
|
func (w *TimingWheel) taskPool() {
|
|
defer func() {
|
|
if e := recover(); e != nil {
|
|
log.Fatal("task pool fatal %v", e)
|
|
}
|
|
}()
|
|
|
|
for {
|
|
select {
|
|
case tasks := <-w.tasks:
|
|
for _, task := range tasks {
|
|
task()
|
|
}
|
|
case <-w.quit:
|
|
return
|
|
}
|
|
}
|
|
}
|