From 216b6e96a907bd6f15aad0b93b296adc30d7616a Mon Sep 17 00:00:00 2001
From: Will Assis <35489495+gassiss@users.noreply.github.com>
Date: Wed, 19 Mar 2025 09:15:13 -0300
Subject: [PATCH] fix(unified-storage): update instrumentation_server metric
 gatherer (#102399)

* update instrumentation_server used by the module server to use prometheus gatherer from wire instead of DefaultRegisterer
---
 pkg/server/instrumentation_service.go      | 16 +++++++++-------
 pkg/server/instrumentation_service_test.go |  2 +-
 pkg/server/module_server.go                | 12 ++++++++----
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/pkg/server/instrumentation_service.go b/pkg/server/instrumentation_service.go
index 7eca2a54fda..9a200ec7f38 100644
--- a/pkg/server/instrumentation_service.go
+++ b/pkg/server/instrumentation_service.go
@@ -9,19 +9,21 @@ import (
 	"github.com/grafana/dskit/services"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 type instrumentationService struct {
 	*services.BasicService
-	cfg      *setting.Cfg
-	httpServ *http.Server
-	log      log.Logger
-	errChan  chan error
+	cfg          *setting.Cfg
+	httpServ     *http.Server
+	log          log.Logger
+	errChan      chan error
+	promGatherer prometheus.Gatherer
 }
 
-func NewInstrumentationService(log log.Logger, cfg *setting.Cfg) (*instrumentationService, error) {
-	s := &instrumentationService{log: log, cfg: cfg}
+func NewInstrumentationService(log log.Logger, cfg *setting.Cfg, promGatherer prometheus.Gatherer) (*instrumentationService, error) {
+	s := &instrumentationService{log: log, cfg: cfg, promGatherer: promGatherer}
 	s.BasicService = services.NewBasicService(s.start, s.running, s.stop)
 	return s, nil
 }
@@ -56,7 +58,7 @@ func (s *instrumentationService) stop(failureReason error) error {
 
 func (s *instrumentationService) newInstrumentationServer(ctx context.Context) *http.Server {
 	router := http.NewServeMux()
-	router.Handle("/metrics", promhttp.Handler())
+	router.Handle("/metrics", promhttp.HandlerFor(s.promGatherer, promhttp.HandlerOpts{EnableOpenMetrics: true}))
 
 	srv := &http.Server{
 		// 5s timeout for header reads to avoid Slowloris attacks (https://thetooth.io/blog/slowloris-attack/)
diff --git a/pkg/server/instrumentation_service_test.go b/pkg/server/instrumentation_service_test.go
index 51af5ee4032..1dc8ffa8349 100644
--- a/pkg/server/instrumentation_service_test.go
+++ b/pkg/server/instrumentation_service_test.go
@@ -18,7 +18,7 @@ import (
 func TestRunInstrumentationService(t *testing.T) {
 	cfg := setting.NewCfg()
 	cfg.HTTPPort = "3001"
-	s, err := NewInstrumentationService(log.New("test-logger"), cfg)
+	s, err := NewInstrumentationService(log.New("test-logger"), cfg, prometheus.DefaultGatherer)
 	require.NoError(t, err)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
diff --git a/pkg/server/module_server.go b/pkg/server/module_server.go
index 82fd314f5e1..11db7c204d3 100644
--- a/pkg/server/module_server.go
+++ b/pkg/server/module_server.go
@@ -18,12 +18,13 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/grafana/grafana/pkg/storage/unified/sql"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 // NewModule returns an instance of a ModuleServer, responsible for managing
 // dskit modules (services).
-func NewModule(opts Options, apiOpts api.ServerOptions, features featuremgmt.FeatureToggles, cfg *setting.Cfg, storageMetrics *resource.StorageMetrics, indexMetrics *resource.BleveIndexMetrics) (*ModuleServer, error) {
-	s, err := newModuleServer(opts, apiOpts, features, cfg, storageMetrics, indexMetrics)
+func NewModule(opts Options, apiOpts api.ServerOptions, features featuremgmt.FeatureToggles, cfg *setting.Cfg, storageMetrics *resource.StorageMetrics, indexMetrics *resource.BleveIndexMetrics, promGatherer prometheus.Gatherer) (*ModuleServer, error) {
+	s, err := newModuleServer(opts, apiOpts, features, cfg, storageMetrics, indexMetrics, promGatherer)
 	if err != nil {
 		return nil, err
 	}
@@ -35,7 +36,7 @@ func NewModule(opts Options, apiOpts api.ServerOptions, features featuremgmt.Fea
 	return s, nil
 }
 
-func newModuleServer(opts Options, apiOpts api.ServerOptions, features featuremgmt.FeatureToggles, cfg *setting.Cfg, storageMetrics *resource.StorageMetrics, indexMetrics *resource.BleveIndexMetrics) (*ModuleServer, error) {
+func newModuleServer(opts Options, apiOpts api.ServerOptions, features featuremgmt.FeatureToggles, cfg *setting.Cfg, storageMetrics *resource.StorageMetrics, indexMetrics *resource.BleveIndexMetrics, promGatherer prometheus.Gatherer) (*ModuleServer, error) {
 	rootCtx, shutdownFn := context.WithCancel(context.Background())
 
 	s := &ModuleServer{
@@ -53,6 +54,7 @@ func newModuleServer(opts Options, apiOpts api.ServerOptions, features featuremg
 		buildBranch:      opts.BuildBranch,
 		storageMetrics:   storageMetrics,
 		indexMetrics:     indexMetrics,
+		promGatherer:     promGatherer,
 	}
 
 	return s, nil
@@ -81,6 +83,8 @@ type ModuleServer struct {
 	version     string
 	commit      string
 	buildBranch string
+
+	promGatherer prometheus.Gatherer
 }
 
 // init initializes the server and its services.
@@ -119,7 +123,7 @@ func (s *ModuleServer) Run() error {
 		if m.IsModuleEnabled(modules.All) || m.IsModuleEnabled(modules.Core) {
 			return services.NewBasicService(nil, nil, nil).WithName(modules.InstrumentationServer), nil
 		}
-		return NewInstrumentationService(s.log, s.cfg)
+		return NewInstrumentationService(s.log, s.cfg, s.promGatherer)
 	})
 
 	m.RegisterModule(modules.Core, func() (services.Service, error) {