1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package main
4
5import (
6 "embed"
7 "encoding/json"
8 "errors"
9 "fmt"
10 "math/rand/v2"
11 "os"
12 "path"
13 "regexp"
14 "strings"
15)
16
17// A Scenario is the traffic of one or more pages: weighted page requests and
18// Server-Sent Events streams kept open while the requests run, as the pages
19// of real visitors do.
20type Scenario struct {
21 Name string `json:"name"`
22 Description string `json:"description"`
23 Requests []Request `json:"requests"`
24 Streams []Stream `json:"streams"`
25}
26
27// A Request is one page or asset fetched over and over. Weight is its share
28// of the requests of the scenario.
29type Request struct {
30 Name string `json:"name"`
31 Path string `json:"path"`
32 Weight int `json:"weight"`
33 // Status is the expected status code, 200 when absent.
34 Status int `json:"status"`
35}
36
37// A Stream is a live updates feed (GET /-/events) held open during the run.
38// Count is how many are opened, scaled by the -streams flag.
39type Stream struct {
40 Name string `json:"name"`
41 Path string `json:"path"`
42 Count int `json:"count"`
43}
44
45//go:embed scenarios/*.json
46var builtin embed.FS
47
48var varRef = regexp.MustCompile(`\{\{([a-z_]+)\}\}`)
49
50// loadScenario reads a built-in scenario by name, or a JSON file by path,
51// and replaces the {{name}} references of its paths with vars.
52func loadScenario(ref string, vars map[string]string) (*Scenario, error) {
53 var data []byte
54 var err error
55 if strings.HasSuffix(ref, ".json") {
56 data, err = os.ReadFile(ref)
57 } else {
58 data, err = builtin.ReadFile(path.Join("scenarios", ref+".json"))
59 }
60 if err != nil {
61 return nil, fmt.Errorf("scenario %q: %w (built-in: %s)", ref, err, strings.Join(builtinNames(), ", "))
62 }
63 var sc Scenario
64 if err := json.Unmarshal(data, &sc); err != nil {
65 return nil, fmt.Errorf("scenario %q: %w", ref, err)
66 }
67 return &sc, sc.resolve(vars)
68}
69
70func builtinNames() []string {
71 entries, _ := builtin.ReadDir("scenarios")
72 names := make([]string, 0, len(entries))
73 for _, e := range entries {
74 names = append(names, strings.TrimSuffix(e.Name(), ".json"))
75 }
76 return names
77}
78
79// resolve checks the scenario and expands its variables.
80func (sc *Scenario) resolve(vars map[string]string) error {
81 if len(sc.Requests) == 0 {
82 return errors.New("scenario has no requests")
83 }
84 var missing []string
85 expand := func(p string) string {
86 return varRef.ReplaceAllStringFunc(p, func(m string) string {
87 name := varRef.FindStringSubmatch(m)[1]
88 v, ok := vars[name]
89 if !ok {
90 missing = append(missing, name)
91 }
92 return v
93 })
94 }
95 for i := range sc.Requests {
96 r := &sc.Requests[i]
97 if r.Weight <= 0 {
98 r.Weight = 1
99 }
100 if r.Status == 0 {
101 r.Status = 200
102 }
103 if r.Name == "" {
104 r.Name = r.Path
105 }
106 r.Path = expand(r.Path)
107 }
108 for i := range sc.Streams {
109 s := &sc.Streams[i]
110 if s.Name == "" {
111 s.Name = s.Path
112 }
113 s.Path = expand(s.Path)
114 }
115 if len(missing) > 0 {
116 return fmt.Errorf("missing -var for %s", strings.Join(missing, ", "))
117 }
118 return nil
119}
120
121// picker draws requests in proportion to their weights.
122type picker struct {
123 reqs []Request
124 total int
125}
126
127func newPicker(reqs []Request) picker {
128 p := picker{reqs: reqs}
129 for _, r := range reqs {
130 p.total += r.Weight
131 }
132 return p
133}
134
135func (p picker) pick(rng *rand.Rand) int {
136 n := rng.IntN(p.total)
137 for i, r := range p.reqs {
138 if n < r.Weight {
139 return i
140 }
141 n -= r.Weight
142 }
143 return len(p.reqs) - 1
144}