Add timing mode

timewheel.go

@@ -7,10 +7,12 @@ import (
 )
 
 type taskSlice struct {
-	id    string
-	pos   int
-	cycle int
-	task  func()
+	id       string
+	pos      int
+	cycle    int
+	rawCycle int
+	mode     TimeMode
+	task     func()
 }
 
 type TimeWheel struct {
@@ -25,6 +27,13 @@ type TimeWheel struct {
 	removeTaskC chan string
 }
 
+type TimeMode int
+
+const (
+	TimerTypeOnce TimeMode = iota
+	TimeTypeLoop
+)
+
 func New(slots int, interval time.Duration) *TimeWheel {
 	if slots <= 0 {
 		slots = 10
@@ -38,7 +47,7 @@ func New(slots int, interval time.Duration) *TimeWheel {
 		interval:    interval,
 		ticker:      time.NewTicker(interval),
 		slots:       make([]*list.List, 0, slots),
-		taskMap:     map[string]*list.Element{},
+		taskMap:     make(map[string]*list.Element),
 		stopC:       make(chan struct{}),
 		addTaskC:    make(chan *taskSlice),
 		removeTaskC: make(chan string),
@@ -53,14 +62,16 @@ func New(slots int, interval time.Duration) *TimeWheel {
 	return t
 }
 
-func (t *TimeWheel) AddTask(id string, task func(), flower time.Time) {
+func (t *TimeWheel) AddTask(id string, mode TimeMode, task func(), flower time.Time) {
 	cycle, pos := t.getSlotPosAndCycle(flower)
 
 	t.addTaskC <- &taskSlice{
-		id:    id,
-		task:  task,
-		cycle: cycle,
-		pos:   pos,
+		id:       id,
+		task:     task,
+		cycle:    cycle,
+		mode:     mode,
+		rawCycle: cycle,
+		pos:      pos,
 	}
 }
 
@@ -110,7 +121,7 @@ func (t *TimeWheel) updateCurtSlot() {
 }
 
 func (t *TimeWheel) execute(l *list.List) {
-	for e := l.Back(); e != nil; {
+	for e := l.Front(); e != nil; {
 		event := e.Value.(*taskSlice)
 		if event.cycle > 0 {
 			event.cycle--
@@ -128,10 +139,15 @@ func (t *TimeWheel) execute(l *list.List) {
 			event.task()
 		}()
 
-		next := e.Next()
-		l.Remove(e)
-		delete(t.taskMap, event.id)
-		e = next
+		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
+		}
 	}
 }
 

timewheel_test.go

@@ -6,17 +6,19 @@ import (
 )
 
 func Test_timeWheel(t *testing.T) {
-	timeWheel := New(10, 500*time.Millisecond)
+	timeWheel := New(10, 100*time.Millisecond)
 	defer timeWheel.Stop()
 
+	<-time.After(10 * time.Millisecond)
+
 	t.Errorf("test2, %v", time.Now())
-	timeWheel.AddTask("test1", func() {
+	timeWheel.AddTask("test1", TimeTypeLoop, func() {
 		t.Errorf("test1, %v", time.Now())
 	}, time.Now().Add(time.Second))
-	timeWheel.AddTask("test2", func() {
+	timeWheel.AddTask("test2", TimeTypeLoop, func() {
 		t.Errorf("test2, %v", time.Now())
 	}, time.Now().Add(5*time.Second))
-	timeWheel.AddTask("test2", func() {
+	timeWheel.AddTask("test2", TimerTypeOnce, func() {
 		t.Errorf("test2, %v", time.Now())
 	}, time.Now().Add(3*time.Second))