forked from myzhan/boomer
/
boomer.go
80 lines (65 loc) · 1.5 KB
/
boomer.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
package boomer
import (
"flag"
"log"
"os"
"os/signal"
"runtime"
"strings"
"syscall"
)
// Run accepts a slice of Task and connects
// to a locust master.
func Run(tasks ...*Task) {
// support go version below 1.5
runtime.GOMAXPROCS(runtime.NumCPU())
if !flag.Parsed() {
flag.Parse()
}
if *runTasks != "" {
// Run tasks without connecting to the master.
taskNames := strings.Split(*runTasks, ",")
for _, task := range tasks {
if task.Name == "" {
continue
} else {
for _, name := range taskNames {
if name == task.Name {
log.Println("Running " + task.Name)
task.Fn()
}
}
}
}
return
}
if maxRPS > 0 {
log.Println("Max RPS that boomer may generate is limited to", maxRPS)
maxRPSEnabled = true
}
var r *runner
client := newClient()
r = &runner{
tasks: tasks,
client: client,
nodeID: getNodeID(),
}
Events.Subscribe("boomer:quit", r.onQuiting)
r.getReady()
c := make(chan os.Signal)
signal.Notify(c, syscall.SIGINT)
<-c
Events.Publish("boomer:quit")
// wait for quit message is sent to master
<-disconnectedFromServer
log.Println("shut down")
}
var runTasks *string
var maxRPS int64
var maxRPSThreshold int64
var maxRPSEnabled = false
var maxRPSControlChannel = make(chan bool)
func init() {
runTasks = flag.String("run-tasks", "", "Run tasks without connecting to the master, multiply tasks is seperated by comma. Usually, it's for debug purpose.")
flag.Int64Var(&maxRPS, "max-rps", 0, "Max RPS that boomer can generate.")
}