forked from tsliwowicz/go-wrk
/
loader.go
196 lines (177 loc) · 5 KB
/
loader.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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
package loader
import (
"bytes"
"fmt"
"io"
"io/ioutil"
"log"
"net/http"
"net/url"
"strings"
"sync/atomic"
"time"
"github.com/tsliwowicz/go-wrk/util"
)
const (
USER_AGENT = "go-wrk"
)
type LoadCfg struct {
duration int //seconds
goroutines int
testUrl string
reqBody string
method string
host string
header map[string]string
statsAggregator chan *RequesterStats
timeoutms int
allowRedirects bool
disableCompression bool
disableKeepAlive bool
interrupted int32
clientCert string
clientKey string
caCert string
http2 bool
}
// RequesterStats used for colelcting aggregate statistics
type RequesterStats struct {
TotRespSize int64
TotDuration time.Duration
MinRequestTime time.Duration
MaxRequestTime time.Duration
NumRequests int
NumErrs int
}
func NewLoadCfg(duration int, //seconds
goroutines int,
testUrl string,
reqBody string,
method string,
host string,
header map[string]string,
statsAggregator chan *RequesterStats,
timeoutms int,
allowRedirects bool,
disableCompression bool,
disableKeepAlive bool,
clientCert string,
clientKey string,
caCert string,
http2 bool) (rt *LoadCfg) {
rt = &LoadCfg{duration, goroutines, testUrl, reqBody, method, host, header, statsAggregator, timeoutms,
allowRedirects, disableCompression, disableKeepAlive, 0, clientCert, clientKey, caCert, http2}
return
}
func escapeUrlStr(in string) string {
qm := strings.Index(in, "?")
if qm != -1 {
qry := in[qm+1:]
qrys := strings.Split(qry, "&")
var query string = ""
var qEscaped string = ""
var first bool = true
for _, q := range qrys {
qSplit := strings.Split(q, "=")
if len(qSplit) == 2 {
qEscaped = qSplit[0] + "=" + url.QueryEscape(qSplit[1])
} else {
qEscaped = qSplit[0]
}
if first {
first = false
} else {
query += "&"
}
query += qEscaped
}
return in[:qm] + "?" + query
} else {
return in
}
}
//DoRequest single request implementation. Returns the size of the response and its duration
//On error - returns -1 on both
func DoRequest(httpClient *http.Client, header map[string]string, method, host, loadUrl, reqBody string) (respSize int, duration time.Duration) {
respSize = -1
duration = -1
loadUrl = escapeUrlStr(loadUrl)
var buf io.Reader
if len(reqBody) > 0 {
buf = bytes.NewBufferString(reqBody)
}
req, err := http.NewRequest(method, loadUrl, buf)
if err != nil {
fmt.Println("An error occured doing request", err)
return
}
for hk, hv := range header {
req.Header.Add(hk, hv)
}
req.Header.Add("User-Agent", USER_AGENT)
if host != "" {
req.Host = host
}
start := time.Now()
resp, err := httpClient.Do(req)
if err != nil {
fmt.Println("redirect?")
//this is a bit weird. When redirection is prevented, a url.Error is retuned. This creates an issue to distinguish
//between an invalid URL that was provided and and redirection error.
rr, ok := err.(*url.Error)
if !ok {
fmt.Println("An error occured doing request", err, rr)
return
}
fmt.Println("An error occured doing request", err)
}
if resp == nil {
fmt.Println("empty response")
return
}
defer func() {
if resp != nil && resp.Body != nil {
resp.Body.Close()
}
}()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
fmt.Println("An error occured reading body", err)
}
if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
duration = time.Since(start)
respSize = len(body) + int(util.EstimateHttpHeadersSize(resp.Header))
} else if resp.StatusCode == http.StatusMovedPermanently || resp.StatusCode == http.StatusTemporaryRedirect {
duration = time.Since(start)
respSize = int(resp.ContentLength) + int(util.EstimateHttpHeadersSize(resp.Header))
} else {
fmt.Println("received status code", resp.StatusCode, "from", resp.Header, "content", string(body), req)
}
return
}
//Requester a go function for repeatedly making requests and aggregating statistics as long as required
//When it is done, it sends the results using the statsAggregator channel
func (cfg *LoadCfg) RunSingleLoadSession() {
stats := &RequesterStats{MinRequestTime: time.Minute}
start := time.Now()
httpClient, err := client(cfg.disableCompression, cfg.disableKeepAlive, cfg.timeoutms, cfg.allowRedirects, cfg.clientCert, cfg.clientKey, cfg.caCert, cfg.http2)
if err != nil {
log.Fatal(err)
}
for time.Since(start).Seconds() <= float64(cfg.duration) && atomic.LoadInt32(&cfg.interrupted) == 0 {
respSize, reqDur := DoRequest(httpClient, cfg.header, cfg.method, cfg.host, cfg.testUrl, cfg.reqBody)
if respSize > 0 {
stats.TotRespSize += int64(respSize)
stats.TotDuration += reqDur
stats.MaxRequestTime = util.MaxDuration(reqDur, stats.MaxRequestTime)
stats.MinRequestTime = util.MinDuration(reqDur, stats.MinRequestTime)
stats.NumRequests++
} else {
stats.NumErrs++
}
}
cfg.statsAggregator <- stats
}
func (cfg *LoadCfg) Stop() {
atomic.StoreInt32(&cfg.interrupted, 1)
}