/
scheduler.go
126 lines (106 loc) · 3.7 KB
/
scheduler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
// Package scheduler provides a periodic task and implements a task scheduler. Its
// primary purpose is to manage the volume of scraping requests to a target.
package scheduler
import (
"github.com/opinionated/utils/log"
"time"
)
// Scheduler runs Schedulable tasks asynchronously at a specified time. Allows for dynamic
// task rescheduling and multithreaded task adding.
type Scheduler struct {
queue []Schedulable // sorted array of tasks to run
addTask chan Schedulable // tasks are put on here when they are able to run
quit chan bool // signal the scheduler to stop once no more tasks are ready
ready chan Schedulable // the next task to run. This needs to be buffered or deadlock will occur
isRunning bool // if the scheduler is running
cycleTime int
}
// MakeScheduler creates a new Scheduler.
// QueueSize is how many tasks can be held.
// BufferSize is how many tasks can be added to the queue each cycle. If bufferSize < num tasks
// being added at once then some of the adds will block unless run as goroutines choice of buffer
// size shouldn't make a huge difference
func MakeScheduler(queueSize, bufferSize int) *Scheduler {
// NOTE: go figures out that you are returning a pointer to the local variable and puts it on the heap for you
return &Scheduler{make([]Schedulable, 0, queueSize),
make(chan Schedulable, bufferSize),
make(chan bool),
make(chan Schedulable, 1),
false,
15}
}
// Add adds a new schedulable to the run queue. May block if the waiting queue is full.
func (scheduler *Scheduler) Add(schedulable Schedulable) {
scheduler.addTask <- schedulable
}
// SetCycleTime of the main loop.
func (scheduler *Scheduler) SetCycleTime(time int) {
scheduler.cycleTime = time
}
// Start the Scheduler asynchronously. Does not block.
func (scheduler *Scheduler) Start() {
go scheduler.Run()
}
// Stop the scheduler from running. May block until Scheduler ends.
func (scheduler *Scheduler) Stop() {
scheduler.quit <- true
}
// IsRunning returns true if the scheduler is running, otherwise returns flase.
func (scheduler *Scheduler) IsRunning() bool {
return scheduler.isRunning
}
// Run the main scheduler loop.
func (scheduler *Scheduler) Run() {
scheduler.isRunning = true
// signals the loop to run every (cycleTime) seconds
ticker := time.NewTicker(time.Duration(scheduler.cycleTime) * time.Second)
for {
didAdd := false // keep track of adds so we only sort when we need to
AddNewTasksLoop:
for {
// add tasks from buffered channel to queue until all waiting tasks are added
select {
case s := <-scheduler.addTask:
scheduler.queue = append(scheduler.queue, s)
didAdd = true
default:
break AddNewTasksLoop
}
}
// only sort if we added a new task
// TODO: reschedule tasks properly
if didAdd {
By(SortLowToHigh).Sort(scheduler.queue)
}
i := 0 // declare outside so we can use later
didRemove := false
for ; i < len(scheduler.queue); i++ {
if scheduler.queue[i].TimeRemaining() < scheduler.cycleTime {
// run any tasks that are ready
go scheduler.queue[i].Run(scheduler)
didRemove = true
} else { // no tasks to run this cycle
break
}
}
if didRemove {
// remove unused elements
// do it this way to make sure we avoid mem leaks
// (something could be sitting in an unused part of the queue and not get cleared)
copy(scheduler.queue[0:], scheduler.queue[i:])
for j := 1; j <= i; j++ {
scheduler.queue[len(scheduler.queue)-j] = nil
}
scheduler.queue = scheduler.queue[:len(scheduler.queue)-i]
}
select {
case <-ticker.C:
// wait until next time step
case <-scheduler.quit:
scheduler.isRunning = false
log.Warn("Done with scheduler")
ticker.Stop()
return
}
}
}