1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package main
4
5import (
6 "bufio"
7 "context"
8 "errors"
9 "fmt"
10 "io"
11 "math/rand/v2"
12 "net"
13 "net/http"
14 "net/netip"
15 "net/url"
16 "strings"
17 "sync"
18 "sync/atomic"
19 "time"
20)
21
22type config struct {
23 base string
24 scenario string
25 vars varsFlag
26 clients int
27 rate float64
28 duration time.Duration
29 warmup time.Duration
30 streams float64
31 token string
32 gzip bool
33 jsonOut string
34}
35
36// sample is one finished request.
37type sample struct {
38 req int
39 latency time.Duration
40 bytes int64
41 err string // "" when the request got the expected status
42}
43
44// job is one request to send; at is when it was due, so that open-loop
45// latencies include the time spent waiting for a client (no coordinated
46// omission).
47type job struct {
48 req int
49 at time.Time
50}
51
52func run(ctx context.Context, cfg config, sc *Scenario) (*Report, error) {
53 base := strings.TrimRight(cfg.base, "/")
54 tr := &http.Transport{
55 MaxIdleConns: cfg.clients * 2,
56 MaxIdleConnsPerHost: cfg.clients * 2,
57 DisableCompression: true, // the Accept-Encoding header is set by hand
58 IdleConnTimeout: 90 * time.Second,
59 }
60 client := &http.Client{Transport: tr, Timeout: 30 * time.Second}
61 defer tr.CloseIdleConnections()
62
63 if err := preflight(ctx, client, cfg, base, sc); err != nil {
64 return nil, err
65 }
66
67 // Streams are opened first and stay open for the whole run.
68 sctx, scancel := context.WithCancel(ctx)
69 defer scancel()
70 streamStats := openStreams(sctx, cfg, base, sc)
71
72 if cfg.warmup > 0 {
73 fire(ctx, client, cfg, base, sc, cfg.warmup, nil)
74 }
75 var mu sync.Mutex
76 var samples []sample
77 collect := func(s sample) {
78 mu.Lock()
79 samples = append(samples, s)
80 mu.Unlock()
81 }
82 start := time.Now()
83 missed := fire(ctx, client, cfg, base, sc, cfg.duration, collect)
84 elapsed := time.Since(start)
85 scancel()
86 streams := streamStats.wait()
87 if err := ctx.Err(); err != nil {
88 return nil, fmt.Errorf("interrupted after %v: %w", elapsed.Round(time.Millisecond), err)
89 }
90
91 return newReport(cfg, sc, samples, missed, elapsed, streams), nil
92}
93
94// preflight fetches every request once and fails fast on a wrong status, so
95// that a typo in a scenario does not turn into a run of errors.
96func preflight(ctx context.Context, client *http.Client, cfg config, base string, sc *Scenario) error {
97 var errs []string
98 for _, r := range sc.Requests {
99 s := do(ctx, client, cfg, base, r)
100 if s.err != "" {
101 errs = append(errs, fmt.Sprintf("%s (%s): %s", r.Name, r.Path, s.err))
102 }
103 }
104 if len(errs) > 0 {
105 return errors.New("preflight failed:\n " + strings.Join(errs, "\n "))
106 }
107 return nil
108}
109
110// fire sends requests for d, closed-loop or open-loop, and hands each result
111// to collect (nil during the warm-up). It returns the open-loop requests that
112// found no free client.
113//
114// The deadline only stops new requests: those in flight or queued when it
115// passes still run to the end (bounded by the client timeout) and are
116// counted, so that slow responses are not dropped from the figures. Only an
117// interruption of ctx cancels them.
118func fire(ctx context.Context, client *http.Client, cfg config, base string, sc *Scenario, d time.Duration, collect func(sample)) int64 {
119 sched, cancel := context.WithTimeout(ctx, d)
120 defer cancel()
121 pk := newPicker(sc.Requests)
122 var wg sync.WaitGroup
123
124 send := func(j job) {
125 s := do(ctx, client, cfg, base, sc.Requests[j.req])
126 if ctx.Err() != nil {
127 return // interrupted: the run reports no figures
128 }
129 s.req = j.req
130 if !j.at.IsZero() {
131 s.latency = time.Since(j.at)
132 }
133 if collect != nil {
134 collect(s)
135 }
136 }
137
138 if cfg.rate <= 0 {
139 for i := range cfg.clients {
140 rng := rand.New(rand.NewPCG(uint64(i), uint64(time.Now().UnixNano())))
141 wg.Go(func() {
142 for sched.Err() == nil {
143 send(job{req: pk.pick(rng)})
144 }
145 })
146 }
147 wg.Wait()
148 return 0
149 }
150
151 // The queue absorbs scheduling hiccups of the load generator; a request
152 // that finds it full is missed.
153 jobs := make(chan job, cfg.clients)
154 for range cfg.clients {
155 wg.Go(func() {
156 for j := range jobs {
157 send(j)
158 }
159 })
160 }
161 var missed int64
162 rng := rand.New(rand.NewPCG(1, uint64(time.Now().UnixNano())))
163 interval := time.Duration(float64(time.Second) / cfg.rate)
164 next := time.Now()
165loop:
166 for {
167 if wait := time.Until(next); wait > 0 {
168 select {
169 case <-sched.Done():
170 break loop
171 case <-time.After(wait):
172 }
173 }
174 if sched.Err() != nil {
175 break
176 }
177 select {
178 case jobs <- job{req: pk.pick(rng), at: next}:
179 default:
180 missed++
181 }
182 next = next.Add(interval)
183 }
184 close(jobs)
185 wg.Wait()
186 if collect == nil {
187 return 0
188 }
189 return missed
190}
191
192// do sends one request and reads the whole body.
193func do(ctx context.Context, client *http.Client, cfg config, base string, r Request) sample {
194 req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+r.Path, nil)
195 if err != nil {
196 return sample{err: err.Error()}
197 }
198 header(req, cfg)
199 t := time.Now()
200 resp, err := client.Do(req)
201 if err != nil {
202 return sample{latency: time.Since(t), err: shortErr(err)}
203 }
204 n, err := io.Copy(io.Discard, resp.Body)
205 _ = resp.Body.Close() // fully read: closing frees the connection
206 s := sample{latency: time.Since(t), bytes: n}
207 switch {
208 case err != nil:
209 s.err = shortErr(err)
210 case resp.StatusCode != r.Status:
211 s.err = fmt.Sprintf("status %d, want %d", resp.StatusCode, r.Status)
212 }
213 return s
214}
215
216func header(req *http.Request, cfg config) {
217 req.Header.Set("User-Agent", "ermite-loadtest")
218 if cfg.gzip {
219 req.Header.Set("Accept-Encoding", "gzip, br, zstd")
220 }
221 if cfg.token != "" {
222 req.Header.Set("Authorization", "Bearer "+cfg.token)
223 }
224}
225
226func shortErr(err error) string {
227 s := err.Error()
228 if i := strings.LastIndex(s, ": "); i >= 0 {
229 s = s[i+2:]
230 }
231 return s
232}
233
234// streamConnectTimeout bounds the wait for the headers of a stream.
235const streamConnectTimeout = 5 * time.Second
236
237// streamSet follows the open streams of a run.
238type streamSet struct {
239 wg sync.WaitGroup
240 stats []*streamStat
241}
242
243type streamStat struct {
244 name string
245 want int
246 opened atomic.Int64
247 refused atomic.Int64 // a status other than 200, or a failed dial
248 dropped atomic.Int64 // closed by the server before the end of the run
249 events atomic.Int64
250 ttfb struct {
251 sync.Mutex
252 all []time.Duration
253 }
254}
255
256func openStreams(ctx context.Context, cfg config, base string, sc *Scenario) *streamSet {
257 set := &streamSet{}
258 if cfg.streams <= 0 {
259 return set
260 }
261 // Each stream is one visitor: its own connection and, against a loopback
262 // server, its own source address (127.1.x.y), since the server bounds the
263 // streams per client address.
264 loopback := isLoopback(base)
265 n := 0
266 for _, s := range sc.Streams {
267 st := &streamStat{name: s.Name, want: int(float64(s.Count)*cfg.streams + .5)}
268 set.stats = append(set.stats, st)
269 for range st.want {
270 client := streamClient(loopback, n)
271 n++
272 set.wg.Go(func() { follow(ctx, client, cfg, base+s.Path, st) })
273 }
274 }
275 // Let them connect before the measured run starts; a stream still not
276 // open after streamConnectTimeout has failed (Report.Failed).
277 deadline := time.Now().Add(streamConnectTimeout + time.Second)
278 for time.Now().Before(deadline) {
279 done := true
280 for _, st := range set.stats {
281 if int(st.opened.Load()+st.refused.Load()) < st.want {
282 done = false
283 }
284 }
285 if done {
286 break
287 }
288 time.Sleep(20 * time.Millisecond)
289 }
290 return set
291}
292
293func isLoopback(base string) bool {
294 u, err := url.Parse(base)
295 if err != nil {
296 return false
297 }
298 ip, err := netip.ParseAddr(u.Hostname())
299 return err == nil && ip.Is4() && ip.IsLoopback()
300}
301
302// streamClient is the client of the n-th stream.
303func streamClient(loopback bool, n int) *http.Client {
304 d := &net.Dialer{Timeout: 5 * time.Second}
305 if loopback {
306 src := netip.AddrFrom4([4]byte{127, 1, byte(n / 250), byte(n%250 + 1)})
307 d.LocalAddr = &net.TCPAddr{IP: src.AsSlice()}
308 }
309 // A stream is a single request: no keep-alive, so that a stream the server
310 // closes frees its connection at once. A stream whose headers do not come
311 // in time counts as refused.
312 return &http.Client{Transport: &http.Transport{
313 DialContext: d.DialContext,
314 DisableCompression: true,
315 DisableKeepAlives: true,
316 ResponseHeaderTimeout: streamConnectTimeout,
317 }}
318}
319
320func follow(ctx context.Context, client *http.Client, cfg config, target string, st *streamStat) {
321 req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
322 if err != nil {
323 st.refused.Add(1)
324 return
325 }
326 header(req, cfg)
327 req.Header.Del("Accept-Encoding")
328 req.Header.Set("Accept", "text/event-stream")
329 t := time.Now()
330 resp, err := client.Do(req)
331 if err != nil {
332 if ctx.Err() == nil {
333 st.refused.Add(1)
334 }
335 return
336 }
337 defer func() { _ = resp.Body.Close() }()
338 if resp.StatusCode != http.StatusOK {
339 st.refused.Add(1)
340 return
341 }
342 st.opened.Add(1)
343 st.ttfb.Lock()
344 st.ttfb.all = append(st.ttfb.all, time.Since(t))
345 st.ttfb.Unlock()
346 sc := bufio.NewScanner(resp.Body)
347 sc.Buffer(make([]byte, 64<<10), 4<<20)
348 for sc.Scan() {
349 if strings.HasPrefix(sc.Text(), "event:") {
350 st.events.Add(1)
351 }
352 }
353 if ctx.Err() == nil {
354 st.dropped.Add(1)
355 }
356}
357
358func (s *streamSet) wait() []StreamReport {
359 s.wg.Wait()
360 out := make([]StreamReport, 0, len(s.stats))
361 for _, st := range s.stats {
362 out = append(out, StreamReport{
363 Name: st.name, Wanted: st.want,
364 Opened: int(st.opened.Load()), Refused: int(st.refused.Load()),
365 Dropped: int(st.dropped.Load()), Events: int(st.events.Load()),
366 Connect: summarize(st.ttfb.all),
367 })
368 }
369 return out
370}