From 24af5049bbf665376d8da3d72404d670cdabfa9d Mon Sep 17 00:00:00 2001
From: gotjosh <josue@grafana.com>
Date: Thu, 16 Apr 2020 09:30:47 +0100
Subject: [PATCH] API: Allow `TargetRetriever` to receive a Context (#7125)

Fixes #7103

Signed-off-by: gotjosh <josue@grafana.com>
---
 web/api/v1/api.go      | 15 ++++++++-------
 web/api/v1/api_test.go | 10 ++++++++--
 web/web.go             |  5 ++++-
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/web/api/v1/api.go b/web/api/v1/api.go
index 97790365c..2bcccbba4 100644
--- a/web/api/v1/api.go
+++ b/web/api/v1/api.go
@@ -100,7 +100,8 @@ func (e *apiError) Error() string {
 	return fmt.Sprintf("%s: %s", e.typ, e.err)
 }
 
-type targetRetriever interface {
+// TargetRetriever provides the list of active/dropped targets to scrape or not.
+type TargetRetriever interface {
 	TargetsActive() map[string][]*scrape.Target
 	TargetsDropped() map[string][]*scrape.Target
 }
@@ -173,7 +174,7 @@ type API struct {
 	Queryable   storage.Queryable
 	QueryEngine *promql.Engine
 
-	targetRetriever       targetRetriever
+	targetRetriever       func(context.Context) TargetRetriever
 	alertmanagerRetriever alertmanagerRetriever
 	rulesRetriever        rulesRetriever
 	now                   func() time.Time
@@ -202,7 +203,7 @@ func init() {
 func NewAPI(
 	qe *promql.Engine,
 	q storage.Queryable,
-	tr targetRetriever,
+	tr func(context.Context) TargetRetriever,
 	ar alertmanagerRetriever,
 	configFunc func() config.Config,
 	flagsMap map[string]string,
@@ -659,7 +660,7 @@ func (api *API) targets(r *http.Request) apiFuncResult {
 	res := &TargetDiscovery{}
 
 	if showActive {
-		targetsActive := api.targetRetriever.TargetsActive()
+		targetsActive := api.targetRetriever(r.Context()).TargetsActive()
 		activeKeys, numTargets := sortKeys(targetsActive)
 		res.ActiveTargets = make([]*Target, 0, numTargets)
 
@@ -697,7 +698,7 @@ func (api *API) targets(r *http.Request) apiFuncResult {
 		res.ActiveTargets = []*Target{}
 	}
 	if showDropped {
-		tDropped := flatten(api.targetRetriever.TargetsDropped())
+		tDropped := flatten(api.targetRetriever(r.Context()).TargetsDropped())
 		res.DroppedTargets = make([]*DroppedTarget, 0, len(tDropped))
 		for _, t := range tDropped {
 			res.DroppedTargets = append(res.DroppedTargets, &DroppedTarget{
@@ -744,7 +745,7 @@ func (api *API) targetMetadata(r *http.Request) apiFuncResult {
 	metric := r.FormValue("metric")
 
 	res := []metricMetadata{}
-	for _, tt := range api.targetRetriever.TargetsActive() {
+	for _, tt := range api.targetRetriever(r.Context()).TargetsActive() {
 		for _, t := range tt {
 			if limit >= 0 && len(res) >= limit {
 				break
@@ -877,7 +878,7 @@ func (api *API) metricMetadata(r *http.Request) apiFuncResult {
 
 	metric := r.FormValue("metric")
 
-	for _, tt := range api.targetRetriever.TargetsActive() {
+	for _, tt := range api.targetRetriever(r.Context()).TargetsActive() {
 		for _, t := range tt {
 
 			if metric == "" {
diff --git a/web/api/v1/api_test.go b/web/api/v1/api_test.go
index 2c8a7c6ac..552fda467 100644
--- a/web/api/v1/api_test.go
+++ b/web/api/v1/api_test.go
@@ -163,6 +163,12 @@ func (t *testTargetRetriever) ResetMetadataStore() {
 	}
 }
 
+func (t *testTargetRetriever) toFactory() func(context.Context) TargetRetriever {
+	return func(context.Context) TargetRetriever {
+		return t
+	}
+}
+
 type testAlertmanagerRetriever struct{}
 
 func (t testAlertmanagerRetriever) Alertmanagers() []*url.URL {
@@ -311,7 +317,7 @@ func TestEndpoints(t *testing.T) {
 		api := &API{
 			Queryable:             suite.Storage(),
 			QueryEngine:           suite.QueryEngine(),
-			targetRetriever:       testTargetRetriever,
+			targetRetriever:       testTargetRetriever.toFactory(),
 			alertmanagerRetriever: testAlertmanagerRetriever{},
 			flagsMap:              sampleFlagMap,
 			now:                   func() time.Time { return now },
@@ -375,7 +381,7 @@ func TestEndpoints(t *testing.T) {
 		api := &API{
 			Queryable:             remote,
 			QueryEngine:           suite.QueryEngine(),
-			targetRetriever:       testTargetRetriever,
+			targetRetriever:       testTargetRetriever.toFactory(),
 			alertmanagerRetriever: testAlertmanagerRetriever{},
 			flagsMap:              sampleFlagMap,
 			now:                   func() time.Time { return now },
diff --git a/web/web.go b/web/web.go
index fadd119c3..7e377b49a 100644
--- a/web/web.go
+++ b/web/web.go
@@ -293,7 +293,10 @@ func New(logger log.Logger, o *Options) *Handler {
 		ready: 0,
 	}
 
-	h.apiV1 = api_v1.NewAPI(h.queryEngine, h.storage, h.scrapeManager, h.notifier,
+	factoryTr := func(_ context.Context) api_v1.TargetRetriever {
+		return h.scrapeManager
+	}
+	h.apiV1 = api_v1.NewAPI(h.queryEngine, h.storage, factoryTr, h.notifier,
 		func() config.Config {
 			h.mtx.RLock()
 			defer h.mtx.RUnlock()