API: Allow AlertmanagerRetriever and RulesRetriever to receive a Context (#7256)

* API: Allow AlertmanagerRetriever and RulesRetriever to receive a
Context

Signed-off-by: gotjosh <josue@grafana.com>

* Use single line functions where possible

Signed-off-by: gotjosh <josue@grafana.com>

web/api/v1/api.go

@@ -105,12 +105,14 @@ type TargetRetriever interface {
 	TargetsDropped() map[string][]*scrape.Target
 }
 
-type alertmanagerRetriever interface {
+// AlertmanagerRetriever provides a list of all/dropped AlertManager URLs.
+type AlertmanagerRetriever interface {
 	Alertmanagers() []*url.URL
 	DroppedAlertmanagers() []*url.URL
 }
 
-type rulesRetriever interface {
+// RulesRetriever provides a list of active rules and alerts.
+type RulesRetriever interface {
 	RuleGroups() []*rules.Group
 	AlertingRules() []*rules.AlertingRule
 }
@@ -174,8 +176,8 @@ type API struct {
 	QueryEngine *promql.Engine
 
 	targetRetriever       func(context.Context) TargetRetriever
-	alertmanagerRetriever alertmanagerRetriever
-	rulesRetriever        rulesRetriever
+	alertmanagerRetriever func(context.Context) AlertmanagerRetriever
+	rulesRetriever        func(context.Context) RulesRetriever
 	now                   func() time.Time
 	config                func() config.Config
 	flagsMap              map[string]string
@@ -204,7 +206,7 @@ func NewAPI(
 	qe *promql.Engine,
 	q storage.Queryable,
 	tr func(context.Context) TargetRetriever,
-	ar alertmanagerRetriever,
+	ar func(context.Context) AlertmanagerRetriever,
 	configFunc func() config.Config,
 	flagsMap map[string]string,
 	globalURLOptions GlobalURLOptions,
@@ -213,7 +215,7 @@ func NewAPI(
 	dbDir string,
 	enableAdmin bool,
 	logger log.Logger,
-	rr rulesRetriever,
+	rr func(context.Context) RulesRetriever,
 	remoteReadSampleLimit int,
 	remoteReadConcurrencyLimit int,
 	remoteReadMaxBytesInFrame int,
@@ -847,8 +849,8 @@ type AlertmanagerTarget struct {
 }
 
 func (api *API) alertmanagers(r *http.Request) apiFuncResult {
-	urls := api.alertmanagerRetriever.Alertmanagers()
-	droppedURLS := api.alertmanagerRetriever.DroppedAlertmanagers()
+	urls := api.alertmanagerRetriever(r.Context()).Alertmanagers()
+	droppedURLS := api.alertmanagerRetriever(r.Context()).DroppedAlertmanagers()
 	ams := &AlertmanagerDiscovery{ActiveAlertmanagers: make([]*AlertmanagerTarget, len(urls)), DroppedAlertmanagers: make([]*AlertmanagerTarget, len(droppedURLS))}
 	for i, url := range urls {
 		ams.ActiveAlertmanagers[i] = &AlertmanagerTarget{URL: url.String()}
@@ -874,7 +876,7 @@ type Alert struct {
 }
 
 func (api *API) alerts(r *http.Request) apiFuncResult {
-	alertingRules := api.rulesRetriever.AlertingRules()
+	alertingRules := api.rulesRetriever(r.Context()).AlertingRules()
 	alerts := []*Alert{}
 
 	for _, alertingRule := range alertingRules {
@@ -1021,7 +1023,7 @@ type recordingRule struct {
 }
 
 func (api *API) rules(r *http.Request) apiFuncResult {
-	ruleGroups := api.rulesRetriever.RuleGroups()
+	ruleGroups := api.rulesRetriever(r.Context()).RuleGroups()
 	res := &RuleDiscovery{RuleGroups: make([]*RuleGroup, len(ruleGroups))}
 	typeParam := strings.ToLower(r.URL.Query().Get("type"))
 

web/api/v1/api_test.go

@@ -164,9 +164,7 @@ func (t *testTargetRetriever) ResetMetadataStore() {
 }
 
 func (t *testTargetRetriever) toFactory() func(context.Context) TargetRetriever {
-	return func(context.Context) TargetRetriever {
-		return t
-	}
+	return func(context.Context) TargetRetriever { return t }
 }
 
 type testAlertmanagerRetriever struct{}
@@ -191,6 +189,10 @@ func (t testAlertmanagerRetriever) DroppedAlertmanagers() []*url.URL {
 	}
 }
 
+func (t testAlertmanagerRetriever) toFactory() func(context.Context) AlertmanagerRetriever {
+	return func(context.Context) AlertmanagerRetriever { return t }
+}
+
 type rulesRetrieverMock struct {
 	testing *testing.T
 }
@@ -276,6 +278,10 @@ func (m rulesRetrieverMock) RuleGroups() []*rules.Group {
 	return []*rules.Group{group}
 }
 
+func (m rulesRetrieverMock) toFactory() func(context.Context) RulesRetriever {
+	return func(context.Context) RulesRetriever { return m }
+}
+
 var samplePrometheusCfg = config.Config{
 	GlobalConfig:       config.GlobalConfig{},
 	AlertingConfig:     config.AlertingConfig{},
@@ -318,12 +324,12 @@ func TestEndpoints(t *testing.T) {
 			Queryable:             suite.Storage(),
 			QueryEngine:           suite.QueryEngine(),
 			targetRetriever:       testTargetRetriever.toFactory(),
-			alertmanagerRetriever: testAlertmanagerRetriever{},
+			alertmanagerRetriever: testAlertmanagerRetriever{}.toFactory(),
 			flagsMap:              sampleFlagMap,
 			now:                   func() time.Time { return now },
 			config:                func() config.Config { return samplePrometheusCfg },
 			ready:                 func(f http.HandlerFunc) http.HandlerFunc { return f },
-			rulesRetriever:        algr,
+			rulesRetriever:        algr.toFactory(),
 		}
 
 		testEndpoints(t, api, testTargetRetriever, true)
@@ -382,12 +388,12 @@ func TestEndpoints(t *testing.T) {
 			Queryable:             remote,
 			QueryEngine:           suite.QueryEngine(),
 			targetRetriever:       testTargetRetriever.toFactory(),
-			alertmanagerRetriever: testAlertmanagerRetriever{},
+			alertmanagerRetriever: testAlertmanagerRetriever{}.toFactory(),
 			flagsMap:              sampleFlagMap,
 			now:                   func() time.Time { return now },
 			config:                func() config.Config { return samplePrometheusCfg },
 			ready:                 func(f http.HandlerFunc) http.HandlerFunc { return f },
-			rulesRetriever:        algr,
+			rulesRetriever:        algr.toFactory(),
 		}
 
 		testEndpoints(t, api, testTargetRetriever, false)

web/web.go

@@ -297,10 +297,11 @@ func New(logger log.Logger, o *Options) *Handler {
 		ready: 0,
 	}
 
-	factoryTr := func(_ context.Context) api_v1.TargetRetriever {
-		return h.scrapeManager
-	}
-	h.apiV1 = api_v1.NewAPI(h.queryEngine, h.storage, factoryTr, h.notifier,
+	factoryTr := func(_ context.Context) api_v1.TargetRetriever { return h.scrapeManager }
+	factoryAr := func(_ context.Context) api_v1.AlertmanagerRetriever { return h.notifier }
+	FactoryRr := func(_ context.Context) api_v1.RulesRetriever { return h.ruleManager }
+
+	h.apiV1 = api_v1.NewAPI(h.queryEngine, h.storage, factoryTr, factoryAr,
 		func() config.Config {
 			h.mtx.RLock()
 			defer h.mtx.RUnlock()
@@ -317,7 +318,7 @@ func New(logger log.Logger, o *Options) *Handler {
 		h.options.TSDBDir,
 		h.options.EnableAdminAPI,
 		logger,
-		h.ruleManager,
+		FactoryRr,
 		h.options.RemoteReadSampleLimit,
 		h.options.RemoteReadConcurrencyLimit,
 		h.options.RemoteReadBytesInFrame,