diff --git a/pkg/infra/usagestats/statscollector/prometheus_flavor.go b/pkg/infra/usagestats/statscollector/prometheus_flavor.go index 9bba50c17da..8cd2ed31318 100644 --- a/pkg/infra/usagestats/statscollector/prometheus_flavor.go +++ b/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) - if err != nil { - return nil, err + ds := ds + g.Go(func() error { + 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) diff --git a/pkg/infra/usagestats/statscollector/prometheus_flavor_test.go b/pkg/infra/usagestats/statscollector/prometheus_flavor_test.go index d0ad4c4d46b..ed64ecb7319 100644 --- a/pkg/infra/usagestats/statscollector/prometheus_flavor_test.go +++ b/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", diff --git a/pkg/infra/usagestats/statscollector/service_test.go b/pkg/infra/usagestats/statscollector/service_test.go index 2f102ddc6cc..f3525e1c08d 100644 --- a/pkg/infra/usagestats/statscollector/service_test.go +++ b/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{}}), ) }