Skip to content

Commit 8b752ef

Browse files
1. Added comment section of absolute and current block.
2. Removed Err and absoluteTime from the configState object. 3. Added flag package for debugging purpose and removed hard coded file path. Signed-off-by: Kushal Shukla <[email protected]>
1 parent d1560eb commit 8b752ef

File tree

1 file changed

+32
-28
lines changed

1 file changed

+32
-28
lines changed

tools/load-generator/main.go

Lines changed: 32 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
package main
1515

1616
import (
17+
"flag"
1718
"fmt"
1819
"io"
1920
"log"
@@ -98,10 +99,8 @@ type BucketConfig struct {
9899
MaxTime int64 `yaml:"maxTime"`
99100
}
100101

101-
type configState struct {
102+
type bucketState struct {
102103
bucketConfig *BucketConfig
103-
Err error
104-
absoluteTime int64
105104
}
106105

107106
func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
@@ -133,14 +132,16 @@ func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
133132
}
134133

135134
// Function to load `minTime` and `maxTime` from bucket-config.yml
136-
func loadKeyConfig() (*BucketConfig, error) {
137-
filePath := "/config/bucket-config.yml"
138-
_, err := os.Stat(filePath)
135+
func loadBucketConfig() (*BucketConfig, error) {
136+
filePath := flag.String("bucketconfig-file", "/config/bucket-config.yml", "Path to the bucket configuration file")
137+
flag.Parse()
138+
139+
_, err := os.Stat(*filePath)
139140
if os.IsNotExist(err) {
140-
return nil, fmt.Errorf("file not found: %s", filePath)
141+
return nil, fmt.Errorf("file not found: %s", *filePath)
141142
}
142143

143-
data, err := os.ReadFile(filePath)
144+
data, err := os.ReadFile(*filePath)
144145
if err != nil {
145146
return nil, fmt.Errorf("error reading file: %w", err)
146147
}
@@ -154,35 +155,37 @@ func loadKeyConfig() (*BucketConfig, error) {
154155
return &bucketConfig, nil
155156
}
156157

157-
func configstate(v *BucketConfig, err error) *configState {
158-
var absolutetime int64
159-
if v != nil {
160-
absolutetime = v.MaxTime
158+
func setconfig(v *BucketConfig, err error) *bucketState {
159+
// If there is an error in reading bucket-config.yml file then just return nil.
160+
if err != nil {
161+
return nil
161162
}
162-
return &configState{
163+
return &bucketState{
163164
bucketConfig: v,
164-
Err: err,
165-
absoluteTime: absolutetime,
166165
}
167166
}
168167

169-
func (q *Querier) run(wg *sync.WaitGroup, timeBound *configState) {
168+
func (q *Querier) run(wg *sync.WaitGroup, timeBound *bucketState) {
170169
defer wg.Done()
171170
fmt.Printf("Running querier %s %s for %s\n", q.target, q.name, q.url)
172171
time.Sleep(20 * time.Second)
173172

174173
for {
175174
start := time.Now()
175+
// If timeBound is not nil, both the "absolute" and "current" blocks will run;
176+
// otherwise, only the "current" block will execute. This execution pattern is used
177+
// because if Downloaded block data is present, both the head block and downloaded block
178+
// need to be processed.
176179
runBlockMode := "current"
177180
for _, query := range q.queries {
178181
if runBlockMode == "current" {
179182
q.query(query.Expr, "current", nil)
180-
} else if timeBound.Err == nil {
183+
} else if runBlockMode == "absolute" {
181184
q.query(query.Expr, "absolute", timeBound)
182185
}
183-
if runBlockMode == "current" && timeBound.Err == nil {
186+
if runBlockMode == "current" && timeBound != nil {
184187
runBlockMode = "absolute"
185-
} else if timeBound.Err == nil {
188+
} else if timeBound != nil {
186189
runBlockMode = "current"
187190
}
188191
}
@@ -194,7 +197,7 @@ func (q *Querier) run(wg *sync.WaitGroup, timeBound *configState) {
194197
}
195198
}
196199

197-
func (q *Querier) query(expr string, timeMode string, timeBound *configState) {
200+
func (q *Querier) query(expr string, timeMode string, timeBound *bucketState) {
198201
queryCount.WithLabelValues(q.target, q.name, expr, q.qtype).Inc()
199202
start := time.Now()
200203

@@ -219,7 +222,7 @@ func (q *Querier) query(expr string, timeMode string, timeBound *configState) {
219222
qParams.Set("step", q.step)
220223
}
221224
} else if timeMode == "absolute" {
222-
blockinstime := time.Unix(0, timeBound.absoluteTime*int64(time.Millisecond))
225+
blockinstime := time.Unix(0, timeBound.bucketConfig.MaxTime*int64(time.Millisecond))
223226
qParams.Set("time", fmt.Sprintf("%d", int64(blockinstime.Unix())))
224227
}
225228
req.URL.RawQuery = qParams.Encode()
@@ -268,9 +271,13 @@ func main() {
268271
}
269272
prNumber := os.Args[2]
270273

271-
configFile, err := os.ReadFile("/etc/loadgen/config.yaml")
274+
configPath := flag.String("config-file", "/etc/loadgen/config.yaml", "Path to the configuration file")
275+
flag.Parse()
276+
277+
configFile, err := os.ReadFile(*configPath)
272278
if err != nil {
273-
log.Fatalf("Failed to load config: %v", err)
279+
fmt.Printf("Error reading config file: %v\n", err)
280+
return
274281
}
275282

276283
var config struct {
@@ -286,11 +293,8 @@ func main() {
286293

287294
var wg sync.WaitGroup
288295

289-
bucketConfig, err := loadKeyConfig()
290-
if err != nil {
291-
fmt.Printf("bucket-config.yml file is not present: %v\n", err)
292-
}
293-
timeBound := configstate(bucketConfig, err)
296+
bucketConfig, err := loadBucketConfig()
297+
timeBound := setconfig(bucketConfig, err)
294298

295299
for i, group := range config.Querier.Groups {
296300
wg.Add(1)

0 commit comments

Comments
 (0)