optimize timedsched

timedsched.go

@@ -1,35 +1,56 @@
 package kcp
 
 import (
+	"container/heap"
 	"runtime"
 	"sync"
 	"time"
 )
 
-// SystemTimedSched is the library level fixed-interval timer scheduler,
-var SystemTimedSched *TimedSched = NewTimedSched(runtime.NumCPU(), 10*time.Millisecond)
+// SystemTimedSched is the library level timed-scheduler
+var SystemTimedSched *TimedSched = NewTimedSched(runtime.NumCPU())
+
+type timedFunc struct {
+	execute func()
+	ts      time.Time
+}
+
+// a heap for sorted time
+type timedFuncHeap []timedFunc
+
+func (h timedFuncHeap) Len() int           { return len(h) }
+func (h timedFuncHeap) Less(i, j int) bool { return h[i].ts.Before(h[j].ts) }
+func (h timedFuncHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *timedFuncHeap) Push(x interface{}) { *h = append(*h, x.(timedFunc)) }
+
+func (h *timedFuncHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}
 
 // TimedSched represents the control struct for timed parallel scheduler
 type TimedSched struct {
 	// prepending tasks
-	prependTasks    []func()
+	prependTasks    []timedFunc
 	prependLock     sync.Mutex
 	chPrependNotify chan struct{}
 
 	// tasks will be distributed through chTask
-	chTask   chan func()
-	duration time.Duration // interval for scheduler
+	chTask chan timedFunc
 
 	dieOnce sync.Once
 	die     chan struct{}
 }
 
 // NewTimedSched creates a parallel scheduler with parameters parallel and duration
-func NewTimedSched(parallel int, duration time.Duration) *TimedSched {
+func NewTimedSched(parallel int) *TimedSched {
 	ts := new(TimedSched)
-	ts.chTask = make(chan func())
+	ts.chTask = make(chan timedFunc)
 	ts.die = make(chan struct{})
-	ts.duration = duration
 	ts.chPrependNotify = make(chan struct{}, 1)
 
 	for i := 0; i < parallel; i++ {
@@ -40,20 +61,33 @@ func NewTimedSched(parallel int, duration time.Duration) *TimedSched {
 }
 
 func (ts *TimedSched) sched() {
-	var tasks []func()
+	var tasks timedFuncHeap
 	timer := time.NewTimer(0)
 	for {
 		select {
 		case task := <-ts.chTask:
-			tasks = append(tasks, task)
-			if len(tasks) == 1 { // trigger timer
-				timer.Reset(ts.duration)
+			now := time.Now()
+			// delayed! execute immediately
+			if now.After(task.ts) {
+				task.execute()
+			} else {
+				heap.Push(&tasks, task)
+				// activate timer if timer has hibernated
+				// due to no tasks.
+				if tasks.Len() == 1 {
+					timer.Reset(task.ts.Sub(now))
+				}
 			}
 		case <-timer.C:
-			for k := range tasks {
-				tasks[k]()
+			for tasks.Len() > 0 {
+				now := time.Now()
+				if now.After(tasks[0].ts) {
+					heap.Pop(&tasks).(timedFunc).execute()
+				} else {
+					timer.Reset(tasks[0].ts.Sub(now))
+					break
+				}
 			}
-			tasks = tasks[:0]
 		case <-ts.die:
 			return
 		}
@@ -83,9 +117,9 @@ func (ts *TimedSched) prepend() {
 }
 
 // Put a function awaiting to be executed
-func (ts *TimedSched) Put(f func()) {
+func (ts *TimedSched) Put(f func(), duration time.Duration) {
 	ts.prependLock.Lock()
-	ts.prependTasks = append(ts.prependTasks, f)
+	ts.prependTasks = append(ts.prependTasks, timedFunc{f, time.Now().Add(duration)})
 	ts.prependLock.Unlock()
 	select {
 	case ts.chPrependNotify <- struct{}{}: