From bb039680fb5172562e135c576cbb327a45bb9acd Mon Sep 17 00:00:00 2001
From: Kushal Shukla <kushalshukla110@gmail.com>
Date: Sat, 30 Nov 2024 12:34:37 +0530
Subject: [PATCH] 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 <kushalshukla110@gmail.com>
---
 tools/load-generator/main.go | 52 +++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 24 deletions(-)

diff --git a/tools/load-generator/main.go b/tools/load-generator/main.go
index 823790a69..b4b49f93a 100644
--- a/tools/load-generator/main.go
+++ b/tools/load-generator/main.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"io"
 	"log"
@@ -100,8 +101,6 @@ type BucketConfig struct {
 
 type configState struct {
 	bucketConfig *BucketConfig
-	Err          error
-	absoluteTime int64
 }
 
 func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
@@ -133,14 +132,16 @@ func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
 }
 
 // Function to load `minTime` and `maxTime` from bucket-config.yml
-func loadKeyConfig() (*BucketConfig, error) {
-	filePath := "/config/bucket-config.yml"
-	_, err := os.Stat(filePath)
+func loadBucketConfig() (*BucketConfig, error) {
+	filePath := flag.String("bucketconfig-file", "/config/bucket-config.yml", "Path to the bucket configuration file")
+	flag.Parse()
+
+	_, err := os.Stat(*filePath)
 	if os.IsNotExist(err) {
-		return nil, fmt.Errorf("file not found: %s", filePath)
+		return nil, fmt.Errorf("file not found: %s", *filePath)
 	}
 
-	data, err := os.ReadFile(filePath)
+	data, err := os.ReadFile(*filePath)
 	if err != nil {
 		return nil, fmt.Errorf("error reading file: %w", err)
 	}
@@ -154,15 +155,13 @@ func loadKeyConfig() (*BucketConfig, error) {
 	return &bucketConfig, nil
 }
 
-func configstate(v *BucketConfig, err error) *configState {
-	var absolutetime int64
-	if v != nil {
-		absolutetime = v.MaxTime
+func setconfig(v *BucketConfig, err error) *configState {
+	// If there is an error in reading bucket-config.yml file then just return nil.
+	if err != nil {
+		return nil
 	}
 	return &configState{
 		bucketConfig: v,
-		Err:          err,
-		absoluteTime: absolutetime,
 	}
 }
 
@@ -173,16 +172,20 @@ func (q *Querier) run(wg *sync.WaitGroup, timeBound *configState) {
 
 	for {
 		start := time.Now()
+		// If timeBound is not nil, both the "absolute" and "current" blocks will run;
+		// otherwise, only the "current" block will execute. This execution pattern is used
+		// because if TSDB block data is present, both the head block and downloaded block
+		// need to be processed.
 		runBlockMode := "current"
 		for _, query := range q.queries {
 			if runBlockMode == "current" {
 				q.query(query.Expr, "current", nil)
-			} else if timeBound.Err == nil {
+			} else if runBlockMode == "absolute" {
 				q.query(query.Expr, "absolute", timeBound)
 			}
-			if runBlockMode == "current" && timeBound.Err == nil {
+			if runBlockMode == "current" && timeBound != nil {
 				runBlockMode = "absolute"
-			} else if timeBound.Err == nil {
+			} else if timeBound != nil {
 				runBlockMode = "current"
 			}
 		}
@@ -219,7 +222,7 @@ func (q *Querier) query(expr string, timeMode string, timeBound *configState) {
 			qParams.Set("step", q.step)
 		}
 	} else if timeMode == "absolute" {
-		blockinstime := time.Unix(0, timeBound.absoluteTime*int64(time.Millisecond))
+		blockinstime := time.Unix(0, timeBound.bucketConfig.MaxTime*int64(time.Millisecond))
 		qParams.Set("time", fmt.Sprintf("%d", int64(blockinstime.Unix())))
 	}
 	req.URL.RawQuery = qParams.Encode()
@@ -268,9 +271,13 @@ func main() {
 	}
 	prNumber := os.Args[2]
 
-	configFile, err := os.ReadFile("/etc/loadgen/config.yaml")
+	configPath := flag.String("config-file", "/etc/loadgen/config.yaml", "Path to the configuration file")
+	flag.Parse()
+
+	configFile, err := os.ReadFile(*configPath)
 	if err != nil {
-		log.Fatalf("Failed to load config: %v", err)
+		fmt.Printf("Error reading config file: %v\n", err)
+		return
 	}
 
 	var config struct {
@@ -286,11 +293,8 @@ func main() {
 
 	var wg sync.WaitGroup
 
-	bucketConfig, err := loadKeyConfig()
-	if err != nil {
-		fmt.Printf("bucket-config.yml file is not present: %v\n", err)
-	}
-	timeBound := configstate(bucketConfig, err)
+	bucketConfig, err := loadBucketConfig()
+	timeBound := setconfig(bucketConfig, err)
 
 	for i, group := range config.Querier.Groups {
 		wg.Add(1)