Usage stats: Detect Prometheus flavors (#47942)

* Naïve Prometheus flavor detector

* Add concurrency and memoization

* Remove concurrency

* Fix tests

* close response body

* Add tests

pkg/infra/usagestats/statscollector/prometheus_flavor.go

@@ -0,0 +1,119 @@
+package statscollector
+
+import (
+	"context"
+	"encoding/json"
+	"io/ioutil"
+	"net/http"
+	"time"
+
+	"github.com/grafana/grafana/pkg/models"
+)
+
+const promFlavorCacheLifetime = time.Hour
+
+type memoPrometheusFlavor struct {
+	variants map[string]int64
+
+	memoized time.Time
+}
+
+func (s *Service) detectPrometheusVariants(ctx context.Context) (map[string]int64, error) {
+	if s.promFlavorCache.memoized.Add(promFlavorCacheLifetime).After(time.Now()) &&
+		s.promFlavorCache.variants != nil {
+		return s.promFlavorCache.variants, nil
+	}
+
+	dsProm := &models.GetDataSourcesByTypeQuery{Type: "prometheus"}
+	err := s.datasources.GetDataSourcesByType(ctx, dsProm)
+	if err != nil {
+		s.log.Error("Failed to read all Prometheus data sources", "error", err)
+		return nil, err
+	}
+
+	variants := map[string]int64{}
+	for _, ds := range dsProm.Result {
+		variant, err := s.detectPrometheusVariant(ctx, ds)
+		if err != nil {
+			return nil, err
+		}
+		if variant == "" {
+			continue
+		}
+
+		if _, exists := variants[variant]; !exists {
+			variants[variant] = 0
+		}
+		variants[variant] += 1
+	}
+
+	s.promFlavorCache.variants = variants
+	s.promFlavorCache.memoized = time.Now()
+	return variants, nil
+}
+
+func (s *Service) detectPrometheusVariant(ctx context.Context, ds *models.DataSource) (string, error) {
+	type buildInfo struct {
+		Data struct {
+			Application *string                `json:"application"`
+			Features    map[string]interface{} `json:"features"`
+		} `json:"data"`
+	}
+
+	c, err := s.datasources.GetHTTPTransport(ds, s.httpClientProvider)
+	if err != nil {
+		s.log.Error("Failed to get HTTP client for Prometheus data source", "error", err)
+		return "", err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ds.Url+"/api/v1/status/buildinfo", nil)
+	if err != nil {
+		s.log.Error("Failed to create Prometheus build info request", "error", err)
+		return "", err
+	}
+
+	resp, err := c.RoundTrip(req)
+	if err != nil {
+		// Possibly configuration error, the risk of a false positive is
+		// too high.
+		s.log.Debug("Failed to send Prometheus build info request", "error", err)
+		return "", nil
+	}
+	defer func() {
+		err := resp.Body.Close()
+		if err != nil {
+			s.log.Error("Got error while closing response body")
+		}
+	}()
+
+	if resp.StatusCode == 404 {
+		return "cortex-like", nil
+	}
+
+	if resp.StatusCode != 200 {
+		return "unknown", nil
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		s.log.Error("Failed to read Prometheus build info", "error", err)
+		return "", err
+	}
+
+	bi := &buildInfo{}
+	err = json.Unmarshal(body, bi)
+	if err != nil {
+		s.log.Warn("Failed to read Prometheus build info JSON", "error", err)
+		return "", err
+	}
+
+	if bi.Data.Application != nil && *bi.Data.Application == "Grafana Mimir" {
+		return "mimir", nil
+	}
+
+	if bi.Data.Features != nil {
+		return "mimir-like", nil
+	}
+
+	return "vanilla", nil
+}

pkg/infra/usagestats/statscollector/prometheus_flavor_test.go

@@ -0,0 +1,86 @@
+package statscollector
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/sqlstore/mockstore"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+func TestDetectPrometheusVariant(t *testing.T) {
+	vanilla := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = fmt.Fprint(w, `{"status":"success","data":{"version":"","revision":"","branch":"","buildUser":"","buildDate":"","goVersion":"go1.17.6"}}`)
+	}))
+	t.Cleanup(vanilla.Close)
+
+	mimir := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = fmt.Fprint(w, `{"status":"success","data":{"application":"Grafana Mimir","version":"2.0.0","revision":"9fd2da5","branch":"HEAD","goVersion":"go1.17.8","features":{"ruler_config_api":"true","alertmanager_config_api":"true","query_sharding":"false","federated_rules":"false"}}}`)
+	}))
+	t.Cleanup(mimir.Close)
+
+	cortex := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	t.Cleanup(cortex.Close)
+
+	sqlStore := mockstore.NewSQLStoreMock()
+	s := createService(
+		t,
+		setting.NewCfg(),
+		sqlStore,
+		withDatasources(mockDatasourceService{datasources: []*models.DataSource{
+			{
+				Id:      1,
+				OrgId:   1,
+				Version: 1,
+				Name:    "Vanilla",
+				Type:    "prometheus",
+				Access:  "proxy",
+				Url:     vanilla.URL,
+			},
+			{
+				Id:      2,
+				OrgId:   1,
+				Version: 1,
+				Name:    "Mimir",
+				Type:    "prometheus",
+				Access:  "proxy",
+				Url:     mimir.URL,
+			},
+			{
+				Id:      3,
+				OrgId:   1,
+				Version: 1,
+				Name:    "Another Mimir",
+				Type:    "prometheus",
+				Access:  "proxy",
+				Url:     mimir.URL,
+			},
+			{
+				Id:      4,
+				OrgId:   1,
+				Version: 1,
+				Name:    "Cortex",
+				Type:    "prometheus",
+				Access:  "proxy",
+				Url:     cortex.URL,
+			},
+		}}),
+	)
+
+	flavors, err := s.detectPrometheusVariants(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, int64(2), flavors["mimir"])
+	assert.Equal(t, int64(1), flavors["vanilla"])
+	assert.Equal(t, int64(1), flavors["cortex-like"])
+}

pkg/infra/usagestats/statscollector/service.go

@@ -6,31 +6,34 @@ import (
 	"strings"
 	"time"
 
-	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/infra/httpclient"
-
-	"github.com/grafana/grafana/pkg/infra/usagestats"
-
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/metrics"
+	"github.com/grafana/grafana/pkg/infra/usagestats"
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/setting"
 )
 
 type Service struct {
-	cfg        *setting.Cfg
+	cfg                *setting.Cfg
-	sqlstore   sqlstore.Store
+	sqlstore           sqlstore.Store
-	plugins    plugins.Store
+	plugins            plugins.Store
-	social     social.Service
+	social             social.Service
-	usageStats usagestats.Service
+	usageStats         usagestats.Service
-	features   *featuremgmt.FeatureManager
+	features           *featuremgmt.FeatureManager
+	datasources        datasources.DataSourceService
+	httpClientProvider httpclient.Provider
 
 	log log.Logger
 
 	startTime                time.Time
 	concurrentUserStatsCache memoConcurrentUserStats
+	promFlavorCache          memoPrometheusFlavor
 }
 
 func ProvideService(
@@ -40,14 +43,18 @@ func ProvideService(
 	social social.Service,
 	plugins plugins.Store,
 	features *featuremgmt.FeatureManager,
+	datasourceService datasources.DataSourceService,
+	httpClientProvider httpclient.Provider,
 ) *Service {
 	s := &Service{
-		cfg:        cfg,
+		cfg:                cfg,
-		sqlstore:   store,
+		sqlstore:           store,
-		plugins:    plugins,
+		plugins:            plugins,
-		social:     social,
+		social:             social,
-		usageStats: usagestats,
+		usageStats:         usagestats,
-		features:   features,
+		features:           features,
+		datasources:        datasourceService,
+		httpClientProvider: httpClientProvider,
 
 		startTime: time.Now(),
 		log:       log.New("infra.usagestats.collector"),
@@ -189,6 +196,15 @@ func (s *Service) collect(ctx context.Context) (map[string]interface{}, error) {
 		return nil, err
 	}
 
+	variants, err := s.detectPrometheusVariants(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for variant, count := range variants {
+		m["stats.ds.prometheus.flavor."+variant+".count"] = count
+	}
+
 	// send access counters for each data source
 	// but ignore any custom data sources
 	// as sending that name could be sensitive information

pkg/infra/usagestats/statscollector/service_test.go

@@ -3,10 +3,13 @@ package statscollector
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"testing"
 	"time"
 
-	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+
+	"github.com/grafana/grafana/pkg/infra/httpclient"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -16,6 +19,8 @@ import (
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/services/sqlstore/mockstore"
 	"github.com/grafana/grafana/pkg/setting"
@@ -348,9 +353,15 @@ func (pr fakePluginStore) Plugins(_ context.Context, pluginTypes ...plugins.Type
 	return result
 }
 
-func createService(t testing.TB, cfg *setting.Cfg, store sqlstore.Store) *Service {
+func createService(t testing.TB, cfg *setting.Cfg, store sqlstore.Store, opts ...func(*serviceOptions)) *Service {
 	t.Helper()
 
+	o := &serviceOptions{datasources: mockDatasourceService{}}
+
+	for _, opt := range opts {
+		opt(o)
+	}
+
 	return ProvideService(
 		&usagestats.UsageStatsMock{},
 		cfg,
@@ -358,5 +369,32 @@ func createService(t testing.TB, cfg *setting.Cfg, store sqlstore.Store) *Servic
 		&mockSocial{},
 		&fakePluginStore{},
 		featuremgmt.WithFeatures("feature1", "feature2"),
+		o.datasources,
+		httpclient.NewProvider(),
 	)
 }
+
+type serviceOptions struct {
+	datasources datasources.DataSourceService
+}
+
+func withDatasources(ds datasources.DataSourceService) func(*serviceOptions) {
+	return func(options *serviceOptions) {
+		options.datasources = ds
+	}
+}
+
+type mockDatasourceService struct {
+	datasources.DataSourceService
+
+	datasources []*models.DataSource
+}
+
+func (s mockDatasourceService) GetDataSourcesByType(ctx context.Context, query *models.GetDataSourcesByTypeQuery) error {
+	query.Result = s.datasources
+	return nil
+}
+
+func (s mockDatasourceService) GetHTTPTransport(ds *models.DataSource, provider httpclient.Provider, customMiddlewares ...sdkhttpclient.Middleware) (http.RoundTripper, error) {
+	return provider.GetTransport()
+}