From f93e8f8544077ae743a3354c276113f7df9d9aa6 Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Sun, 8 Dec 2024 11:05:52 +0300
Subject: [PATCH] fix(otelbench): properly handle lack of reports

---
 cmd/otelbench/chtracker/chtracker.go   |  9 ++-------
 cmd/otelbench/logqlbench/logqlbench.go | 12 +++++++++++-
 cmd/otelbench/logqlbench/result.go     |  1 +
 cmd/otelbench/promql_bench.go          |  6 +++++-
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/cmd/otelbench/chtracker/chtracker.go b/cmd/otelbench/chtracker/chtracker.go
index 5f16422a..aa2497be 100644
--- a/cmd/otelbench/chtracker/chtracker.go
+++ b/cmd/otelbench/chtracker/chtracker.go
@@ -3,7 +3,6 @@ package chtracker
 
 import (
 	"context"
-	"fmt"
 	"net/http"
 	"sync"
 	"time"
@@ -93,7 +92,7 @@ func (t *Tracker[Q]) Track(ctx context.Context, meta Q, cb func(context.Context,
 }
 
 // Report iterates over tracked queries.
-func (t *Tracker[Q]) Report(ctx context.Context, cb func(context.Context, TrackedQuery[Q], []QueryReport) error) error {
+func (t *Tracker[Q]) Report(ctx context.Context, cb func(context.Context, TrackedQuery[Q], []QueryReport, error) error) error {
 	if err := t.Flush(ctx); err != nil {
 		return err
 	}
@@ -125,11 +124,7 @@ func (t *Tracker[Q]) Report(ctx context.Context, cb func(context.Context, Tracke
 	for i, result := range queries {
 		tq := t.queries[i]
 
-		if result.Err != nil {
-			fmt.Printf("Failed to retrieve reports for %q: %s\n", tq.TraceID, result.Err)
-		}
-
-		if err := cb(ctx, tq, result.Reports); err != nil {
+		if err := cb(ctx, tq, result.Reports, result.Err); err != nil {
 			return errors.Wrap(err, "report callback")
 		}
 	}
diff --git a/cmd/otelbench/logqlbench/logqlbench.go b/cmd/otelbench/logqlbench/logqlbench.go
index 48c774cb..62c9cf93 100644
--- a/cmd/otelbench/logqlbench/logqlbench.go
+++ b/cmd/otelbench/logqlbench/logqlbench.go
@@ -132,7 +132,16 @@ func (p *LogQLBenchmark) Run(ctx context.Context) error {
 
 	var reports []LogQLReportQuery
 	if err := p.tracker.Report(ctx,
-		func(ctx context.Context, tq chtracker.TrackedQuery[Query], queries []chtracker.QueryReport) error {
+		func(ctx context.Context, tq chtracker.TrackedQuery[Query], queries []chtracker.QueryReport, retriveErr error) error {
+			var errMsg string
+			if retriveErr != nil {
+				if errors.Is(retriveErr, context.DeadlineExceeded) {
+					errMsg = "no queries"
+				} else {
+					errMsg = retriveErr.Error()
+				}
+			}
+
 			header := tq.Meta.Header()
 			reports = append(reports, LogQLReportQuery{
 				ID:            header.ID,
@@ -144,6 +153,7 @@ func (p *LogQLBenchmark) Run(ctx context.Context) error {
 				DurationNanos: tq.Duration.Nanoseconds(),
 				Queries:       queries,
 				Timeout:       tq.Timeout,
+				ReportError:   errMsg,
 			})
 			return nil
 		},
diff --git a/cmd/otelbench/logqlbench/result.go b/cmd/otelbench/logqlbench/result.go
index 06e2393e..1a07b36a 100644
--- a/cmd/otelbench/logqlbench/result.go
+++ b/cmd/otelbench/logqlbench/result.go
@@ -16,4 +16,5 @@ type LogQLReportQuery struct {
 	Matchers      []string                `yaml:"matchers,omitempty"`
 	Queries       []chtracker.QueryReport `yaml:"queries,omitempty"`
 	Timeout       bool                    `yaml:"timeout,omitempty"`
+	ReportError   string                  `yaml:"report_error,omitempty"`
 }
diff --git a/cmd/otelbench/promql_bench.go b/cmd/otelbench/promql_bench.go
index 01d2a8df..0fc38df3 100644
--- a/cmd/otelbench/promql_bench.go
+++ b/cmd/otelbench/promql_bench.go
@@ -419,7 +419,11 @@ func (p *PromQL) Run(ctx context.Context) error {
 
 	var reports []PromQLReportQuery
 	if err := p.tracker.Report(ctx,
-		func(ctx context.Context, tq chtracker.TrackedQuery[promQLQuery], queries []chtracker.QueryReport) error {
+		func(ctx context.Context, tq chtracker.TrackedQuery[promQLQuery], queries []chtracker.QueryReport, retriveErr error) error {
+			if retriveErr != nil {
+				return retriveErr
+			}
+
 			entry := PromQLReportQuery{
 				ID:            tq.Meta.ID,
 				DurationNanos: tq.Duration.Nanoseconds(),