-
Notifications
You must be signed in to change notification settings - Fork 138
/
attacker.go
109 lines (98 loc) · 2.51 KB
/
attacker.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
package attacker
import (
"context"
"math"
"net/http"
"time"
vegeta "github.com/tsenart/vegeta/v12/lib"
)
const (
DefaultRate = 50
DefaultDuration = 10 * time.Second
DefaultTimeout = 30 * time.Second
DefaultMethod = http.MethodGet
DefaultWorkers = 10
DefaultMaxWorkers = math.MaxUint64
DefaultMaxBody = int64(-1)
DefaultConnections = 10000
)
type Attacker interface {
Attack(vegeta.Targeter, vegeta.Pacer, time.Duration, string) <-chan *vegeta.Result
Stop()
}
// Options provides optional settings to attack.
type Options struct {
Rate int
Duration time.Duration
Timeout time.Duration
Method string
Body []byte
MaxBody int64
Header http.Header
Workers uint64
MaxWorkers uint64
KeepAlive bool
Connections int
Attacker Attacker
}
// Result contains the results of a single HTTP request.
type Result struct {
Latency time.Duration
// Indicates if the last result in the entire attack.
End bool
}
// Attack keeps the request running for the specified period of time.
// Results are sent to the given channel as soon as they arrive.
// When the attack is over, it gives back final statistics.
func Attack(ctx context.Context, target string, resCh chan *Result, metricsCh chan *Metrics, opts Options) {
if target == "" {
return
}
if opts.Method == "" {
opts.Method = DefaultMethod
}
if opts.Workers == 0 {
opts.Workers = DefaultWorkers
}
if opts.MaxWorkers == 0 {
opts.MaxWorkers = DefaultMaxWorkers
}
if opts.MaxBody == 0 {
opts.MaxBody = DefaultMaxBody
}
if opts.Connections == 0 {
opts.Connections = DefaultConnections
}
if opts.Attacker == nil {
opts.Attacker = vegeta.NewAttacker(
vegeta.Timeout(opts.Timeout),
vegeta.Workers(opts.Workers),
vegeta.MaxWorkers(opts.MaxWorkers),
vegeta.MaxBody(opts.MaxBody),
vegeta.Connections(opts.Connections),
vegeta.KeepAlive(opts.KeepAlive),
)
}
rate := vegeta.Rate{Freq: opts.Rate, Per: time.Second}
targeter := vegeta.NewStaticTargeter(vegeta.Target{
Method: opts.Method,
URL: target,
Body: opts.Body,
Header: opts.Header,
})
var metrics vegeta.Metrics
for res := range opts.Attacker.Attack(targeter, rate, opts.Duration, "main") {
select {
case <-ctx.Done():
opts.Attacker.Stop()
// metricsCh is already closed (as context is done) so we shouldn't send any metric
return
default:
resCh <- &Result{Latency: res.Latency}
metrics.Add(res)
metricsCh <- newMetrics(&metrics)
}
}
metrics.Close()
metricsCh <- newMetrics(&metrics)
}