overflow_probe/scheduler/scheduler.go

package scheduler

import (
	"errors"
	"reflect"
	"runtime"
	"sort"
	"time"
)

const (
	MAX_TASKS = 10000
)

type Task struct {
	id          string
	intervalSec uint64
	TaskFunc    string
	period      time.Duration
	taskFunc    map[string]interface{}
	funcParams  map[string]([]interface{})
	lastAt      time.Time
	nextAt      time.Time
	running     bool
}

func NewTask(id string, intervel uint64) *Task {
	return &Task{
		id,
		intervel,
		"",
		0,
		make(map[string]interface{}),
		make(map[string]([]interface{})),
		time.Unix(0, 0),
		time.Unix(0, 0),
		true,
	}
}

func (t *Task) runnable() bool {
	return time.Now().After(t.nextAt)
}

func (t *Task) run() {
	if !t.running {
		return
	}
	taskFunc := reflect.ValueOf(t.taskFunc[t.TaskFunc])
	taskParams := t.funcParams[t.TaskFunc]
	in := make([]reflect.Value, len(taskParams))
	for i, param := range taskParams {
		in[i] = reflect.ValueOf(param)
	}
	taskFunc.Call(in)
	t.lastAt = time.Now()
	t.addNextAt()
}

func (t *Task) invoke(TaskFunc interface{}, params ...interface{}) error {
	if reflect.TypeOf(TaskFunc).Kind() != reflect.Func {
		return errors.New("Not a function type.")
	}

	funcName := runtime.FuncForPC(reflect.ValueOf((TaskFunc)).Pointer()).Name()
	t.taskFunc[funcName] = TaskFunc
	t.funcParams[funcName] = params
	t.TaskFunc = funcName
	t.addNextAt()

	return nil
}

func (t *Task) addNextAt() {
	if t.lastAt == time.Unix(0, 0) {
		t.lastAt = time.Now()
	}

	if t.period == 0 {
		t.period = time.Duration(t.intervalSec)
		t.nextAt = t.lastAt.Add(1 * time.Second)
	} else {
		t.nextAt = t.lastAt.Add((t.period * time.Second) - time.Second)
	}

}

type Scheduler struct {
	Tasks [MAX_TASKS]*Task
	size  int
	stop  chan bool
}

func (s *Scheduler) Len() int {
	return s.size
}

func (s *Scheduler) Swap(i, j int) {
	s.Tasks[i], s.Tasks[j] = s.Tasks[j], s.Tasks[i]
}

func (s *Scheduler) Less(i, j int) bool {
	return s.Tasks[j].nextAt.After(s.Tasks[i].nextAt)
}

func (s *Scheduler) RunnableTasks() ([MAX_TASKS]*Task, int) {
	runnableTasks := [MAX_TASKS]*Task{}
	n := 0
	sort.Sort(s)
	for i := 0; i < s.size; i++ {
		if s.Tasks[i].runnable() {
			runnableTasks[n] = s.Tasks[i]
			n++
		} else {
			break
		}
	}
	return runnableTasks, n
}

func (s *Scheduler) nextRun() (*Task, time.Time) {
	if s.size <= 0 {
		return nil, time.Now()
	}
	sort.Sort(s)
	return s.Tasks[0], s.Tasks[0].nextAt
}

func (s *Scheduler) addTask(id string, intervalSec uint64) *Task {
	s.remove(id)

	Task := NewTask(id, intervalSec)
	s.Tasks[s.size] = Task
	s.size++
	return Task
}

func (s *Scheduler) runAll() {
	runnableTasks, n := s.RunnableTasks()

	if n != 0 {
		for i := 0; i < n; i++ {
			go runnableTasks[i].run()
		}
	}
}

func (s *Scheduler) remove(id string) error {
	i := 0
	var ex bool
	for ; i < s.size; i++ {
		if s.Tasks[i].id == id {
			ex = true
			break
		}
	}
	if !ex {
		return nil
	}
	for j := (i + 1); j < s.size; j++ {
		s.Tasks[i] = s.Tasks[j]
		i++
	}
	s.size = s.size - 1

	return nil
}

func (s *Scheduler) Start() {
	s.stop = make(chan bool, 1)
	ticker := time.NewTicker(time.Second * 1)

	go func() {
		for {
			select {
			case <-ticker.C:
				s.runAll()
			case <-s.stop:
				return
			}
		}
	}()
}

func (s *Scheduler) Stop() {
	s.stop <- true
}

func (s *Scheduler) NewSchedule(id string, interval uint64, fn interface{}) error {
	return s.addTask(id, interval).invoke(fn, id)
}

func (s *Scheduler) RemoveSchedule(id string) error {
	i := 0
	var ex bool
	for ; i < s.size; i++ {
		if s.Tasks[i].id == id {
			ex = true
			break
		}
	}
	if !ex {
		return nil
	}
	for j := (i + 1); j < s.size; j++ {
		s.Tasks[i] = s.Tasks[j]
		i++
	}
	s.size = s.size - 1

	return nil
}

func (s *Scheduler) RemoveAllSchedule() error {
	for i := 0; i < s.size; i++ {
		s.Tasks[i] = nil
	}
	s.size = 0

	return nil
}

func (s *Scheduler) UpdateInterval(id string, interval uint64) {
	for i := 0; i < s.Len(); i++ {
		if s.Tasks[i].id == id {
			s.Tasks[i].intervalSec = interval
			s.Tasks[i].period = 0
		}
	}
}

func (s *Scheduler) StartSchedule(id string) error {
	i := 0
	var ex bool
	for ; i < s.size; i++ {
		if s.Tasks[i].id == id {
			ex = true
			break
		}
	}
	if !ex {
		return errors.New("Cannot find the task. ID : " + id)
	}
	if s.Tasks[i].running {
		return errors.New("Already started task. ID : " + id)
	}
	s.Tasks[i].running = true
	return nil
}

func (s *Scheduler) StopSchedule(id string) error {
	i := 0
	var ex bool
	for ; i < s.size; i++ {
		if s.Tasks[i].id == id {
			ex = true
			break
		}
	}
	if s.Tasks[i].running {
		s.Tasks[i].running = false
	}

	if !ex {
		return errors.New("Cannot find the task. ID : " + id)
	}
	if !s.Tasks[i].running {
		return errors.New("Already stopped task. ID : " + id)
	}
	s.Tasks[i].running = false
	return nil
}