Query: adding stats to the remote engine

We are currently losing track of query stats because the remote engine does not transmit performance stats on gRPC calls. In this PR I am adding some fields to the Query API response to include some stats. Signed-off-by: Pedro Tanaka <[email protected]>
junotx · May 23, 2024 · 09f8d0b · 09f8d0b
1 parent 2538562
commit 09f8d0b
Show file tree

Hide file tree

Showing 6 changed files with 371 additions and 61 deletions.
diff --git a/go.mod b/go.mod
@@ -289,3 +289,5 @@ replace (
  k8s.io/klog => github.com/simonpasquier/klog-gokit v0.3.0
  k8s.io/klog/v2 => github.com/simonpasquier/klog-gokit/v3 v3.0.0
 )
+
+replace github.com/thanos-io/promql-engine => github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc
diff --git a/go.sum b/go.sum
@@ -1386,6 +1386,8 @@ github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FI
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y=
 github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
+github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc h1:24B4ydZHfM9IeJLn9MZtiieV5yW3GKY3v8i664dlZ3E=
+github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc/go.mod h1:FEPnabuTql1bDA4OUM41mwcZOJ20R436k8vq+xtGEG0=
 github.com/performancecopilot/speed v3.0.0+incompatible/go.mod h1:/CLtqpZ5gBg1M9iaPbIdPPGyKcA8hKdoy6hAWba7Yac=
 github.com/phpdave11/gofpdf v1.4.2/go.mod h1:zpO6xFn9yxo3YLyMvW8HcKWVdbNqgIfOOp2dXMnm1mY=
 github.com/phpdave11/gofpdi v1.0.12/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI=
@@ -1552,8 +1554,6 @@ github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e h1:f1
 github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e/go.mod h1:jXcofnrSln/cLI6/dhlBxPQZEEQHVPCcFaH75M+nSzM=
 github.com/thanos-io/objstore v0.0.0-20240309075357-e8336a5fd5f3 h1:Q0BjHI7FMe5KkKVXBFYto5VNASxiA/+AEhHup/IT7N0=
 github.com/thanos-io/objstore v0.0.0-20240309075357-e8336a5fd5f3/go.mod h1:ptMYNPgbyAR7a2Ab2t7zHA2/0be2ePyawVR7lp7fZtg=
-github.com/thanos-io/promql-engine v0.0.0-20240405095051-b7d0da367508 h1:4X0ThYb7/wTTKS73wT13ixw0lj5OJ87g45RWIZhPZDA=
-github.com/thanos-io/promql-engine v0.0.0-20240405095051-b7d0da367508/go.mod h1:FEPnabuTql1bDA4OUM41mwcZOJ20R436k8vq+xtGEG0=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab h1:7ZR3hmisBWw77ZpO1/o86g+JV3VKlk3d48jopJxzTjU=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab/go.mod h1:eheTFp954zcWZXCU8d0AT76ftsQOTo4DTqkN/h3k1MY=
 github.com/tklauser/go-sysconf v0.3.4/go.mod h1:Cl2c8ZRWfHD5IrfHo9VN+FX9kCFjIOyVklgXycLB6ek=

diff --git a/pkg/api/query/grpc.go b/pkg/api/query/grpc.go
@@ -9,6 +9,7 @@ import (
 
  "github.com/prometheus/prometheus/promql"
  "github.com/prometheus/prometheus/storage"
+ "github.com/thanos-io/promql-engine/engine"
  "github.com/thanos-io/promql-engine/logicalplan"
  "google.golang.org/grpc"
  "google.golang.org/grpc/codes"
@@ -109,12 +110,22 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
  }
  }
 
+ stats := &querypb.QueryStats{
+ SamplesTotal: 0,
+ PeakSamples: 0,
+ }
+ if explQry, ok := qry.(engine.ExplainableQuery); ok {
+ analyze := explQry.Analyze()
+ stats.SamplesTotal = analyze.TotalSamples()
+ stats.PeakSamples = analyze.PeakSamples()
+ }
+
  switch vector := result.Value.(type) {
  case promql.Scalar:
  series := &prompb.TimeSeries{
  Samples: []prompb.Sample{{Value: vector.V, Timestamp: vector.T}},
  }
- if err := server.Send(querypb.NewQueryResponse(series)); err != nil {
+ if err := server.Send(querypb.NewQueryResponse(series, stats)); err != nil {
  return err
  }
  case promql.Vector:
@@ -125,7 +136,7 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
  Samples: floats,
  Histograms: histograms,
  }
- if err := server.Send(querypb.NewQueryResponse(series)); err != nil {
+ if err := server.Send(querypb.NewQueryResponse(series, stats)); err != nil {
  return err
  }
  }