Prometheus: Add support to make parallel queries (#90316)

* Add support for prometheus datasource to make parallel queries

* Incorporate review comments

* Update pkg/promlib/querydata/request.go

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>

* Fix lint

* Add parallel queries behind feature flag

* Fixing lint issue

* Update go.mod

* Update pkg/promlib/querydata/request.go

* Update pkg/promlib/querydata/request.go

---------

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>
Co-authored-by: Charandas <542168+charandas@users.noreply.github.com>

pkg/promlib/go.mod

@@ -3,6 +3,7 @@ module github.com/grafana/grafana/pkg/promlib
 go 1.22.4
 
 require (
+	github.com/grafana/dskit v0.0.0-20240805174438-dfa83b4ed2d3
 	github.com/grafana/grafana-plugin-sdk-go v0.241.0
 	github.com/json-iterator/go v1.1.12
 	github.com/patrickmn/go-cache v2.1.0+incompatible

pkg/promlib/go.sum

@@ -96,6 +96,8 @@ github.com/gopherjs/gopherjs v0.0.0-20181103185306-d547d1d9531e h1:JKmoR8x90Iww1
 github.com/gopherjs/gopherjs v0.0.0-20181103185306-d547d1d9531e/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
 github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
+github.com/grafana/dskit v0.0.0-20240805174438-dfa83b4ed2d3 h1:as4PmrFoYI1byS5JjsgPC7uSGTMh+SgS0ePv6hOyDGU=
+github.com/grafana/dskit v0.0.0-20240805174438-dfa83b4ed2d3/go.mod h1:lcjGB6SuaZ2o44A9nD6p/tR4QXSPbzViRY520Gy6pTQ=
 github.com/grafana/grafana-plugin-sdk-go v0.241.0 h1:zBcSW9xV9gA9hD8UN+HjJtD7tESMZcaQhA1BI76MTxM=
 github.com/grafana/grafana-plugin-sdk-go v0.241.0/go.mod h1:2HjNwzGCfaFAyR2HGoECTwAmq8vSIn2L1/1yOt4XRS4=
 github.com/grafana/otel-profiling-go v0.5.1 h1:stVPKAFZSa7eGiqbYuG25VcqYksR6iWvF3YH66t4qL8=

pkg/promlib/querydata/request.go

@@ -5,21 +5,21 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"sync"
 	"time"
 
+	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	"github.com/grafana/grafana-plugin-sdk-go/data"
-	"github.com/grafana/grafana-plugin-sdk-go/data/utils/maputil"
-	"go.opentelemetry.io/otel/trace"
-
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/tracing"
-
+	"github.com/grafana/grafana-plugin-sdk-go/data"
+	"github.com/grafana/grafana-plugin-sdk-go/data/utils/maputil"
 	"github.com/grafana/grafana/pkg/promlib/client"
 	"github.com/grafana/grafana/pkg/promlib/intervalv2"
 	"github.com/grafana/grafana/pkg/promlib/models"
 	"github.com/grafana/grafana/pkg/promlib/querydata/exemplar"
 	"github.com/grafana/grafana/pkg/promlib/utils"
+	"go.opentelemetry.io/otel/trace"
 )
 
 const legendFormatAuto = "__auto"
@@ -88,22 +88,49 @@ func (s *QueryData) Execute(ctx context.Context, req *backend.QueryDataRequest)
 		Responses: backend.Responses{},
 	}
 
-	cfg := backend.GrafanaConfigFromContext(ctx)
-	hasPromQLScopeFeatureFlag := cfg.FeatureToggles().IsEnabled("promQLScope")
-	hasPrometheusDataplaneFeatureFlag := cfg.FeatureToggles().IsEnabled("prometheusDataplane")
+	var (
+		cfg                               = backend.GrafanaConfigFromContext(ctx)
+		hasPromQLScopeFeatureFlag         = cfg.FeatureToggles().IsEnabled("promQLScope")
+		hasPrometheusDataplaneFeatureFlag = cfg.FeatureToggles().IsEnabled("prometheusDataplane")
+		hasPrometheusRunQueriesInParallel = cfg.FeatureToggles().IsEnabled("prometheusRunQueriesInParallel")
+	)
 
-	for _, q := range req.Queries {
-		r := s.handleQuery(ctx, q, fromAlert, hasPromQLScopeFeatureFlag, hasPrometheusDataplaneFeatureFlag)
-		if r == nil {
-			continue
+	if hasPrometheusRunQueriesInParallel {
+		var (
+			m sync.Mutex
+		)
+
+		concurrentQueryCount, err := req.PluginContext.GrafanaConfig.ConcurrentQueryCount()
+		if err != nil {
+			logger := s.log.FromContext(ctx)
+			logger.Debug(fmt.Sprintf("Concurrent Query Count read/parse error: %v", err), "prometheusRunQueriesInParallel")
+			concurrentQueryCount = 10
+		}
+
+		_ = concurrency.ForEachJob(ctx, len(req.Queries), concurrentQueryCount, func(ctx context.Context, idx int) error {
+			query := req.Queries[idx]
+			r := s.handleQuery(ctx, query, fromAlert, hasPromQLScopeFeatureFlag, hasPrometheusDataplaneFeatureFlag, true)
+			if r != nil {
+				m.Lock()
+				result.Responses[query.RefID] = *r
+				m.Unlock()
+			}
+			return nil
+		})
+	} else {
+		for _, q := range req.Queries {
+			r := s.handleQuery(ctx, q, fromAlert, hasPromQLScopeFeatureFlag, hasPrometheusDataplaneFeatureFlag, false)
+			if r != nil {
+				result.Responses[q.RefID] = *r
+			}
 		}
-		result.Responses[q.RefID] = *r
 	}
 
 	return &result, nil
 }
 
-func (s *QueryData) handleQuery(ctx context.Context, bq backend.DataQuery, fromAlert, hasPromQLScopeFeatureFlag, hasPrometheusDataplaneFeatureFlag bool) *backend.DataResponse {
+func (s *QueryData) handleQuery(ctx context.Context, bq backend.DataQuery, fromAlert,
+	hasPromQLScopeFeatureFlag, hasPrometheusDataplaneFeatureFlag, hasPrometheusRunQueriesInParallel bool) *backend.DataResponse {
 	traceCtx, span := s.tracer.Start(ctx, "datasource.prometheus")
 	defer span.End()
 	query, err := models.Parse(span, bq, s.TimeInterval, s.intervalCalculator, fromAlert, hasPromQLScopeFeatureFlag)
@@ -113,14 +140,15 @@ func (s *QueryData) handleQuery(ctx context.Context, bq backend.DataQuery, fromA
 		}
 	}
 
-	r := s.fetch(traceCtx, s.client, query, hasPrometheusDataplaneFeatureFlag)
+	r := s.fetch(traceCtx, s.client, query, hasPrometheusDataplaneFeatureFlag, hasPrometheusRunQueriesInParallel)
 	if r == nil {
 		s.log.FromContext(ctx).Debug("Received nil response from runQuery", "query", query.Expr)
 	}
 	return r
 }
 
-func (s *QueryData) fetch(traceCtx context.Context, client *client.Client, q *models.Query, enablePrometheusDataplane bool) *backend.DataResponse {
+func (s *QueryData) fetch(traceCtx context.Context, client *client.Client, q *models.Query,
+	enablePrometheusDataplane, hasPrometheusRunQueriesInParallel bool) *backend.DataResponse {
 	logger := s.log.FromContext(traceCtx)
 	logger.Debug("Sending query", "start", q.Start, "end", q.End, "step", q.Step, "query", q.Expr)
 
@@ -129,37 +157,69 @@ func (s *QueryData) fetch(traceCtx context.Context, client *client.Client, q *mo
 		Error:  nil,
 	}
 
+	var (
+		wg sync.WaitGroup
+		m  sync.Mutex
+	)
+
 	if q.InstantQuery {
-		res := s.instantQuery(traceCtx, client, q, enablePrometheusDataplane)
-		dr.Error = res.Error
-		dr.Frames = res.Frames
-		dr.Status = res.Status
+		if hasPrometheusRunQueriesInParallel {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				res := s.instantQuery(traceCtx, client, q, enablePrometheusDataplane)
+				m.Lock()
+				addDataResponse(&res, dr)
+				m.Unlock()
+			}()
+		} else {
+			res := s.instantQuery(traceCtx, client, q, enablePrometheusDataplane)
+			addDataResponse(&res, dr)
+		}
 	}
 
 	if q.RangeQuery {
-		res := s.rangeQuery(traceCtx, client, q, enablePrometheusDataplane)
-		if res.Error != nil {
-			if dr.Error == nil {
-				dr.Error = res.Error
-			} else {
-				dr.Error = fmt.Errorf("%v %w", dr.Error, res.Error)
-			}
-			// When both instant and range are true, we may overwrite the status code.
-			// To fix this (and other things) they should come in separate http requests.
-			dr.Status = res.Status
+		if hasPrometheusRunQueriesInParallel {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				res := s.rangeQuery(traceCtx, client, q, enablePrometheusDataplane)
+				m.Lock()
+				addDataResponse(&res, dr)
+				m.Unlock()
+			}()
+		} else {
+			res := s.rangeQuery(traceCtx, client, q, enablePrometheusDataplane)
+			addDataResponse(&res, dr)
 		}
-		dr.Frames = append(dr.Frames, res.Frames...)
 	}
 
 	if q.ExemplarQuery {
-		res := s.exemplarQuery(traceCtx, client, q, enablePrometheusDataplane)
-		if res.Error != nil {
-			// If exemplar query returns error, we want to only log it and
-			// continue with other results processing
-			logger.Error("Exemplar query failed", "query", q.Expr, "err", res.Error)
+		if hasPrometheusRunQueriesInParallel {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				res := s.exemplarQuery(traceCtx, client, q, enablePrometheusDataplane)
+				m.Lock()
+				if res.Error != nil {
+					// If exemplar query returns error, we want to only log it and
+					// continue with other results processing
+					logger.Error("Exemplar query failed", "query", q.Expr, "err", res.Error)
+				}
+				dr.Frames = append(dr.Frames, res.Frames...)
+				m.Unlock()
+			}()
+		} else {
+			res := s.exemplarQuery(traceCtx, client, q, enablePrometheusDataplane)
+			if res.Error != nil {
+				// If exemplar query returns error, we want to only log it and
+				// continue with other results processing
+				logger.Error("Exemplar query failed", "query", q.Expr, "err", res.Error)
+			}
+			dr.Frames = append(dr.Frames, res.Frames...)
 		}
-		dr.Frames = append(dr.Frames, res.Frames...)
 	}
+	wg.Wait()
 
 	return dr
 }
@@ -225,3 +285,15 @@ func (s *QueryData) exemplarQuery(ctx context.Context, c *client.Client, q *mode
 	}()
 	return s.parseResponse(ctx, q, res, enablePrometheusDataplaneFlag)
 }
+
+func addDataResponse(res *backend.DataResponse, dr *backend.DataResponse) {
+	if res.Error != nil {
+		if dr.Error == nil {
+			dr.Error = res.Error
+		} else {
+			dr.Error = fmt.Errorf("%v %w", dr.Error, res.Error)
+		}
+		dr.Status = res.Status
+	}
+	dr.Frames = append(dr.Frames, res.Frames...)
+}

pkg/services/featuremgmt/registry.go

@@ -319,6 +319,13 @@ var (
 			FrontendOnly: false,
 			Owner:        grafanaObservabilityMetricsSquad,
 		},
+		{
+			Name:         "prometheusRunQueriesInParallel",
+			Description:  "Enables running Prometheus queries in parallel",
+			Stage:        FeatureStagePrivatePreview,
+			FrontendOnly: false,
+			Owner:        grafanaObservabilityMetricsSquad,
+		},
 		{
 			Name:           "prometheusDataplane",
 			Description:    "Changes responses to from Prometheus to be compliant with the dataplane specification. In particular, when this feature toggle is active, the numeric `Field.Name` is set from 'Value' to the value of the `__name__` label.",

pkg/services/featuremgmt/toggles_gen.csv

@@ -40,6 +40,7 @@ prometheusMetricEncyclopedia,GA,@grafana/observability-metrics,false,false,true
 influxdbBackendMigration,GA,@grafana/observability-metrics,false,false,true
 influxqlStreamingParser,experimental,@grafana/observability-metrics,false,false,false
 influxdbRunQueriesInParallel,privatePreview,@grafana/observability-metrics,false,false,false
+prometheusRunQueriesInParallel,privatePreview,@grafana/observability-metrics,false,false,false
 prometheusDataplane,GA,@grafana/observability-metrics,false,false,false
 lokiMetricDataplane,GA,@grafana/observability-logs,false,false,false
 lokiLogsDataplane,experimental,@grafana/observability-logs,false,false,false

pkg/services/featuremgmt/toggles_gen.go

@@ -171,6 +171,10 @@ const (
 	// Enables running InfluxDB Influxql queries in parallel
 	FlagInfluxdbRunQueriesInParallel = "influxdbRunQueriesInParallel"
 
+	// FlagPrometheusRunQueriesInParallel
+	// Enables running Prometheus queries in parallel
+	FlagPrometheusRunQueriesInParallel = "prometheusRunQueriesInParallel"
+
 	// FlagPrometheusDataplane
 	// Changes responses to from Prometheus to be compliant with the dataplane specification. In particular, when this feature toggle is active, the numeric `Field.Name` is set from 'Value' to the value of the `__name__` label.
 	FlagPrometheusDataplane = "prometheusDataplane"

pkg/services/featuremgmt/toggles_gen.json

@@ -2183,6 +2183,18 @@
         "frontend": true
       }
     },
+    {
+      "metadata": {
+        "name": "prometheusRunQueriesInParallel",
+        "resourceVersion": "1720677541862",
+        "creationTimestamp": "2024-07-11T05:59:01Z"
+      },
+      "spec": {
+        "description": "Enables running Prometheus queries in parallel",
+        "stage": "privatePreview",
+        "codeowner": "@grafana/observability-metrics"
+      }
+    },
     {
       "metadata": {
         "name": "publicDashboards",