From 9490db344717138a575c6a2b8d34b4fabf6ac915 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tsenart@gmail.com>
Date: Tue, 18 Jun 2024 19:50:10 +0200
Subject: [PATCH] add iter

---
 axiom/run.go             | 7 +++++--
 axiom/server_versions.go | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/axiom/run.go b/axiom/run.go
index 83160df86..d5c81e3fd 100644
--- a/axiom/run.go
+++ b/axiom/run.go
@@ -98,7 +98,7 @@ func gitSha() string {
 
 func benchmark(ctx context.Context, cli *axiomClient, id int, query string, iters int, noCache, noColumnFilters bool, enc *json.Encoder) error {
 	for i := 1; i <= iters; i++ {
-		result, err := cli.Query(ctx, id, query, noCache, noColumnFilters)
+		result, err := cli.Query(ctx, i, id, query, noCache, noColumnFilters)
 		if err != nil {
 			return fmt.Errorf("failed query #%d, iter %d: %w", id, i, err)
 		}
@@ -184,6 +184,8 @@ type QueryResult struct {
 	Error string `json:"error"`
 	// Cost is the cost of the query in GBms
 	Cost int `json:"cost"`
+	// Iter is the iteration number of the query within a run
+	Iter int `json:"iter"`
 }
 
 type httpError struct {
@@ -258,7 +260,7 @@ func (c *axiomClient) query(ctx context.Context, id int, aplQuery string, noCach
 	return &r, resp, nil
 }
 
-func (c *axiomClient) Query(ctx context.Context, id int, aplQuery string, noCache, noColumnFilters bool) (*QueryResult, error) {
+func (c *axiomClient) Query(ctx context.Context, iter, id int, aplQuery string, noCache, noColumnFilters bool) (*QueryResult, error) {
 	began := time.Now().UTC()
 
 	result := &QueryResult{
@@ -266,6 +268,7 @@ func (c *axiomClient) Query(ctx context.Context, id int, aplQuery string, noCach
 		ID:      id,
 		Time:    began,
 		Version: c.version,
+		Iter:    iter,
 	}
 
 	var httpErr *httpError
diff --git a/axiom/server_versions.go b/axiom/server_versions.go
index e93220d58..a23aa4a43 100644
--- a/axiom/server_versions.go
+++ b/axiom/server_versions.go
@@ -60,7 +60,7 @@ func serverVersions(apiURL, traceURL, org, token, label string, failfast bool) e
 		earliest time.Time
 	)
 
-	sc.Buffer(make([]byte, 0, 1024*1024), 0)
+	sc.Buffer(make([]byte, 0, 10*1024*1024), 0)
 
 	for sc.Scan() {
 		var r QueryResult