Enhance Load Generator to Query Pre-Generated Data Blocks

prombench/manifests/prombench/benchmark/6_loadgen.yaml

@@ -125,6 +125,23 @@ spec:
       labels:
         app: loadgen-querier
     spec:
+      initContainers:      
+      - name: download-key
+        image: docker.io/prominfra/prometheus-builder:master
+        imagePullPolicy: Always
+        command: [ "/go/src/github.com/key.sh" ]
+        env:
+        - name: PR_NUMBER
+          value: "{{ .PR_NUMBER }}"
+        - name: GITHUB_ORG
+          value: "{{ .GITHUB_ORG }}"
+        - name: GITHUB_REPO
+          value: "{{ .GITHUB_REPO }}"
+        - name: STORAGE
+          value: "/storage-paths"                    
+        volumeMounts:
+        - name: key
+          mountPath: /storage-paths    
       containers:
       - name: prom-load-generator
         image: docker.io/prominfra/load-generator:master
@@ -138,13 +155,17 @@ spec:
         volumeMounts:
         - name: config-volume
           mountPath: /etc/loadgen
+        - name: key
+          mountPath: /config   
         ports:
         - name: loadgen-port
           containerPort: 8080
       volumes:
       - name: config-volume
         configMap:
           name: prometheus-loadgen
+      - name: key
+        emptyDir: {}     
       nodeSelector:
         node-name: nodes-{{ .PR_NUMBER }}
         isolation: none

tools/load-generator/main.go

@@ -92,6 +92,18 @@ type QueryGroup struct {
 	Step     string  `yaml:"step,omitempty"`
 }
 
+type BucketConfig struct {
+	Path    string `yaml:"path"`
+	MinTime int64  `yaml:"minTime"`
+	MaxTime int64  `yaml:"maxTime"`
+}
+
+type configState struct {
+	bucketConfig *BucketConfig
+	Err          error
+	absoluteTime int64
+}
+
 func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
 	qtype := qg.Type
 	if qtype == "" {
@@ -120,16 +132,59 @@ func NewQuerier(groupID int, target, prNumber string, qg QueryGroup) *Querier {
 	}
 }
 
-func (q *Querier) run(wg *sync.WaitGroup) {
+// Function to load `minTime` and `maxTime` from bucket-config.yml
+func loadKeyConfig() (*BucketConfig, error) {
+	filePath := "/config/bucket-config.yml"
+	_, err := os.Stat(filePath)
+	if os.IsNotExist(err) {
+		return nil, fmt.Errorf("file not found: %s", filePath)
+	}
+
+	data, err := os.ReadFile(filePath)
+	if err != nil {
+		return nil, fmt.Errorf("error reading file: %w", err)
+	}
+
+	var bucketConfig BucketConfig
+	err = yaml.Unmarshal(data, &bucketConfig)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing YAML: %w", err)
+	}
+
+	return &bucketConfig, nil
+}
+
+func configstate(v *BucketConfig, err error) *configState {
+	var absolutetime int64
+	if v != nil {
+		absolutetime = v.MaxTime
+	}
+	return &configState{
+		bucketConfig: v,
+		Err:          err,
+		absoluteTime: absolutetime,
+	}
+}
+
+func (q *Querier) run(wg *sync.WaitGroup, timeBound *configState) {
 	defer wg.Done()
 	fmt.Printf("Running querier %s %s for %s\n", q.target, q.name, q.url)
 	time.Sleep(20 * time.Second)
 
 	for {
 		start := time.Now()
-
+		runBlockMode := "current"
 		for _, query := range q.queries {
-			q.query(query.Expr)
+			if runBlockMode == "current" {
+				q.query(query.Expr, "current", nil)
+			} else if timeBound.Err == nil {
+				q.query(query.Expr, "absolute", timeBound)
+			}
+			if runBlockMode == "current" && timeBound.Err == nil {
+				runBlockMode = "absolute"
+			} else if timeBound.Err == nil {
+				runBlockMode = "current"
+			}
 		}
 
 		wait := q.interval - time.Since(start)
@@ -139,7 +194,7 @@ func (q *Querier) run(wg *sync.WaitGroup) {
 	}
 }
 
-func (q *Querier) query(expr string) {
+func (q *Querier) query(expr string, timeMode string, timeBound *configState) {
 	queryCount.WithLabelValues(q.target, q.name, expr, q.qtype).Inc()
 	start := time.Now()
 
@@ -153,9 +208,19 @@ func (q *Querier) query(expr string) {
 	qParams := req.URL.Query()
 	qParams.Set("query", expr)
 	if q.qtype == "range" {
-		qParams.Set("start", fmt.Sprintf("%d", int64(time.Now().Add(-q.start).Unix())))
-		qParams.Set("end", fmt.Sprintf("%d", int64(time.Now().Add(-q.end).Unix())))
-		qParams.Set("step", q.step)
+		if timeMode == "current" {
+			qParams.Set("start", fmt.Sprintf("%d", int64(time.Now().Add(-q.start).Unix())))
+			qParams.Set("end", fmt.Sprintf("%d", int64(time.Now().Add(-q.end).Unix())))
+			qParams.Set("step", q.step)
+		} else {
+			endTime := time.Unix(0, timeBound.bucketConfig.MaxTime*int64(time.Millisecond))
+			qParams.Set("start", fmt.Sprintf("%d", int64(endTime.Add(-q.start).Unix())))
+			qParams.Set("end", fmt.Sprintf("%d", int64(endTime.Add(-q.end).Unix())))
+			qParams.Set("step", q.step)
+		}
+	} else if timeMode == "absolute" {
+		blockinstime := time.Unix(0, timeBound.absoluteTime*int64(time.Millisecond))
+		qParams.Set("time", fmt.Sprintf("%d", int64(blockinstime.Unix())))
 	}
 	req.URL.RawQuery = qParams.Encode()
 
@@ -221,11 +286,17 @@ 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)
+
 	for i, group := range config.Querier.Groups {
 		wg.Add(1)
-		go NewQuerier(i, "pr", prNumber, group).run(&wg)
+		go NewQuerier(i, "pr", prNumber, group).run(&wg, timeBound)
 		wg.Add(1)
-		go NewQuerier(i, "release", prNumber, group).run(&wg)
+		go NewQuerier(i, "release", prNumber, group).run(&wg, timeBound)
 	}
 
 	prometheus.MustRegister(queryDuration, queryCount, queryFailCount)