/
workload.go
88 lines (73 loc) · 1.53 KB
/
workload.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
package keygen
import (
"context"
"fmt"
"os"
"time"
"github.com/dustin/go-humanize"
)
type workload struct {
id int
c Config
generated int
timeout <-chan time.Time
endTime time.Time
ctx context.Context
out func(string)
}
func (w *workload) startWork() {
w.logf("(ID %d) Workload starting", w.id)
if w.c.TimeToRun > 0 {
dur := time.Duration(w.c.TimeToRun) * time.Second
w.endTime = time.Now().Add(dur)
w.timeout = time.After(dur)
}
for !w.isDone() {
w.itemDone(generateKey())
}
}
func (w *workload) isDone() bool {
select {
case <-w.ctx.Done():
w.done(true)
return true
case <-w.timeout:
w.done(false)
return true
default:
}
if w.c.NumToGen > 0 && w.generated >= w.c.NumToGen {
w.done(false)
return true
}
return false
}
func (w *workload) done(canceled bool) {
w.logf("(ID %d) Workload exiting", w.id)
if !canceled && w.c.ExitOnComplete {
os.Exit(w.c.ExitCode)
}
}
func (w *workload) itemDone(desc string) {
w.generated = w.generated + 1
var count string
if w.c.NumToGen > 0 {
count = fmt.Sprintf(" %d/%d", w.generated, w.c.NumToGen)
} else {
count = fmt.Sprintf(" %d/Inf", w.generated)
}
timeleft := ""
if w.timeout != nil {
timeleft = " " + humanize.RelTime(time.Now(), w.endTime, "left", "overdue")
}
if len(desc) > 0 {
desc = ": " + desc
}
w.logf("(ID %d%s%s) Item done%s", w.id, count, timeleft, desc)
}
func (w *workload) log(s string) {
w.out(s)
}
func (w *workload) logf(format string, v ...interface{}) {
w.out(fmt.Sprintf(format, v...))
}