Chore: Fix timeout issues when gathering prometheus datasource stats (#74618)

* fix timeout issues when gathering prometheus flavor stats

* workaround data race in sdk tracing middleware

* cap concurrency at 10

---------

Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>

pkg/infra/usagestats/statscollector/prometheus_flavor.go

@@ -5,8 +5,11 @@ import (
 	"encoding/json"
 	"io"
 	"net/http"
+	"sync"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
 	"github.com/grafana/grafana/pkg/services/datasources"
 )
 
@@ -44,21 +47,43 @@ func (s *Service) detectPrometheusVariants(ctx context.Context) (map[string]int6
 		return nil, err
 	}
 
-	variants := map[string]int64{}
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(10)
+	flavors := sync.Map{}
+
 	for _, ds := range dataSources {
-		variant, err := s.detectPrometheusVariant(ctx, ds)
+		ds := ds
-		if err != nil {
+		g.Go(func() error {
-			return nil, err
+			variant, err := s.detectPrometheusVariant(ctx, ds)
+			if err != nil {
+				return err
+			}
+			flavors.Store(ds.UID, variant)
+			return nil
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+
+	variants := map[string]int64{}
+
+	flavors.Range(func(_, value any) bool {
+		variant, ok := value.(string)
+		if !ok {
+			return true
 		}
+
 		if variant == "" {
-			continue
+			return true
 		}
 
 		if _, exists := variants[variant]; !exists {
 			variants[variant] = 0
 		}
 		variants[variant] += 1
-	}
+		return true
+	})
 
 	s.promFlavorCache.variants = variants
 	s.promFlavorCache.memoized = time.Now()
@@ -66,13 +91,16 @@ func (s *Service) detectPrometheusVariants(ctx context.Context) (map[string]int6
 }
 
 func (s *Service) detectPrometheusVariant(ctx context.Context, ds *datasources.DataSource) (string, error) {
+	// 5s timeout
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
 	type buildInfo struct {
 		Data struct {
 			Application *string        `json:"application"`
 			Features    map[string]any `json:"features"`
 		} `json:"data"`
 	}
-
 	c, err := s.datasources.GetHTTPTransport(ctx, ds, s.httpClientProvider)
 	if err != nil {
 		s.log.Error("Failed to get HTTP client for Prometheus data source", "error", err)

pkg/infra/usagestats/statscollector/prometheus_flavor_test.go

@@ -42,6 +42,7 @@ func TestDetectPrometheusVariant(t *testing.T) {
 		withDatasources(mockDatasourceService{datasources: []*datasources.DataSource{
 			{
 				ID:      1,
+				UID:     "vanilla",
 				OrgID:   1,
 				Version: 1,
 				Name:    "Vanilla",
@@ -51,6 +52,7 @@ func TestDetectPrometheusVariant(t *testing.T) {
 			},
 			{
 				ID:      2,
+				UID:     "mimir",
 				OrgID:   1,
 				Version: 1,
 				Name:    "Mimir",
@@ -60,6 +62,7 @@ func TestDetectPrometheusVariant(t *testing.T) {
 			},
 			{
 				ID:      3,
+				UID:     "another-mimir",
 				OrgID:   1,
 				Version: 1,
 				Name:    "Another Mimir",
@@ -69,6 +72,7 @@ func TestDetectPrometheusVariant(t *testing.T) {
 			},
 			{
 				ID:      4,
+				UID:     "cortex",
 				OrgID:   1,
 				Version: 1,
 				Name:    "Cortex",

pkg/infra/usagestats/statscollector/service_test.go

@@ -384,7 +384,7 @@ func createService(t testing.TB, cfg *setting.Cfg, store db.DB, statsService sta
 		&pluginstore.FakePluginStore{},
 		featuremgmt.WithFeatures("feature1", "feature2"),
 		o.datasources,
-		httpclient.NewProvider(),
+		httpclient.NewProvider(sdkhttpclient.ProviderOptions{Middlewares: []sdkhttpclient.Middleware{}}),
 	)
 }