Datasources: Refactor mixed datasource support to improve concurrency and error handling (#58163)

* recover from panic inside mixed ds query loop

* remove accidental commit

* add messages from panics and errors to the query response

* refactor based on PR comments

* quick update to unit test to verify mixed errors and successes

* reduce concurrency limit

pkg/services/query/query.go

@@ -21,10 +21,10 @@ import (
 	"github.com/grafana/grafana/pkg/tsdb/legacydata"
 	"github.com/grafana/grafana/pkg/util/errutil"
 	"github.com/grafana/grafana/pkg/util/proxyutil"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
-	"golang.org/x/sync/errgroup"
 )
 
 const (
@@ -91,35 +91,60 @@ func (s *Service) QueryData(ctx context.Context, user *user.SignedInUser, skipCa
 		return s.handleQuerySingleDatasource(ctx, user, parsedReq)
 	}
 	// If there are multiple datasources, handle their queries concurrently and return the aggregate result
-	byDataSource := parsedReq.parsedQueries
-	resp := backend.NewQueryDataResponse()
+	return s.executeConcurrentQueries(ctx, user, skipCache, reqDTO, parsedReq.parsedQueries)
+}
 
+// executeConcurrentQueries executes queries to multiple datasources concurrently and returns the aggregate result.
+func (s *Service) executeConcurrentQueries(ctx context.Context, user *user.SignedInUser, skipCache bool, reqDTO dtos.MetricRequest, queriesbyDs map[string][]parsedQuery) (*backend.QueryDataResponse, error) {
 	g, ctx := errgroup.WithContext(ctx)
-	results := make([]backend.Responses, len(byDataSource))
+	g.SetLimit(8) // arbitrary limit to prevent too many concurrent requests
+	rchan := make(chan backend.Responses, len(queriesbyDs))
 
-	for _, queries := range byDataSource {
+	// Create panic recovery function for loop below
+	recoveryFn := func(queries []*simplejson.Json) {
+		if r := recover(); r != nil {
+			var err error
+			s.log.Error("query datasource panic", "error", r, "stack", log.Stack(1))
+			if theErr, ok := r.(error); ok {
+				err = theErr
+			} else if theErrString, ok := r.(string); ok {
+				err = fmt.Errorf(theErrString)
+			} else {
+				err = fmt.Errorf("unexpected error, see the server log for details")
+			}
+			// Due to the panic, there is no valid response for any query for this datasource. Append an error for each one.
+			rchan <- buildErrorResponses(err, queries)
+		}
+	}
+
+	// Query each datasource concurrently
+	for _, queries := range queriesbyDs {
 		rawQueries := make([]*simplejson.Json, len(queries))
 		for i := 0; i < len(queries); i++ {
 			rawQueries[i] = queries[i].rawQuery
 		}
 		g.Go(func() error {
 			subDTO := reqDTO.CloneWithQueries(rawQueries)
+			// Handle panics in the datasource qery
+			defer recoveryFn(subDTO.Queries)
 
 			subResp, err := s.QueryData(ctx, user, skipCache, subDTO)
-
 			if err == nil {
-				results = append(results, subResp.Responses)
+				rchan <- subResp.Responses
+			} else {
+				// If there was an error, return an error response for each query for this datasource
+				rchan <- buildErrorResponses(err, subDTO.Queries)
 			}
-
-			return err
+			return nil
 		})
 	}
 
 	if err := g.Wait(); err != nil {
 		return nil, err
 	}
-
-	for _, result := range results {
+	close(rchan)
+	resp := backend.NewQueryDataResponse()
+	for result := range rchan {
 		for refId, dataResponse := range result {
 			resp.Responses[refId] = dataResponse
 		}
@@ -128,6 +153,17 @@ func (s *Service) QueryData(ctx context.Context, user *user.SignedInUser, skipCa
 	return resp, nil
 }
 
+// buildErrorResponses applies the provided error to each query response in the list. These queries should all belong to the same datasource.
+func buildErrorResponses(err error, queries []*simplejson.Json) backend.Responses {
+	er := backend.Responses{}
+	for _, query := range queries {
+		er[query.Get("refId").MustString("A")] = backend.DataResponse{
+			Error: err,
+		}
+	}
+	return er
+}
+
 // handleExpressions handles POST /api/ds/query when there is an expression.
 func (s *Service) handleExpressions(ctx context.Context, user *user.SignedInUser, parsedReq *parsedRequest) (*backend.QueryDataResponse, error) {
 	exprReq := expr.Request{

pkg/services/query/query_test.go

@@ -406,7 +406,8 @@ func TestQueryDataMultipleSources(t *testing.T) {
 				"datasource": {
 					"type": "mysql",
 					"uid": "ds1"
-				}
+				},
+				"refId": "A"
 			}
 		`))
 		require.NoError(t, err)
@@ -415,7 +416,8 @@ func TestQueryDataMultipleSources(t *testing.T) {
 				"datasource": {
 					"type": "mysql",
 					"uid": "ds2"
-				}
+				},
+				"refId": "B"
 			}
 		`))
 		require.NoError(t, err)
@@ -436,7 +438,6 @@ func TestQueryDataMultipleSources(t *testing.T) {
 
 	t.Run("can query multiple datasources with an expression present", func(t *testing.T) {
 		tc := setup(t)
-		// refId does get set if not included, but better to include it explicitly here
 		query1, err := simplejson.NewJson([]byte(`
 			{
 				"datasource": {
@@ -452,7 +453,8 @@ func TestQueryDataMultipleSources(t *testing.T) {
 				"datasource": {
 					"type": "mysql",
 					"uid": "ds2"
-				}
+				},
+				"refId": "B"
 			}
 		`))
 		require.NoError(t, err)
@@ -493,7 +495,7 @@ func TestQueryDataMultipleSources(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	t.Run("error is returned when one of the queries fails", func(t *testing.T) {
+	t.Run("error is returned in query when one of the queries fails", func(t *testing.T) {
 		tc := setup(t)
 
 		query1, _ := simplejson.NewJson([]byte(`
@@ -501,7 +503,8 @@ func TestQueryDataMultipleSources(t *testing.T) {
 				"datasource": {
 					"type": "mysql",
 					"uid": "ds1"
-				}
+				},
+				"refId": "A"
 			}
 		`))
 		query2, _ := simplejson.NewJson([]byte(`
@@ -510,6 +513,7 @@ func TestQueryDataMultipleSources(t *testing.T) {
 					"type": "prometheus",
 					"uid": "ds2"
 				},
+				"refId": "B",
 				"queryType": "FAIL"
 			}
 		`))
@@ -525,9 +529,12 @@ func TestQueryDataMultipleSources(t *testing.T) {
 			HTTPRequest:                nil,
 		}
 
-		_, err := tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
+		res, err := tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
 
-		require.Error(t, err)
+		require.NoError(t, err)
+		require.Error(t, res.Responses["B"].Error)
+		// Responses aren't mocked, so a "healthy" query will just return an empty response
+		require.NotContains(t, res.Responses, "A")
 	})
 }