From c4f8020dca39f5f80c8f60c81460b1947ed02269 Mon Sep 17 00:00:00 2001
From: Goutham Veeramachaneni <gouthamve@gmail.com>
Date: Wed, 16 Mar 2022 09:45:15 +0100
Subject: [PATCH] Embed MetadaStore in scrape context (#10450)

This will allow downstream users to easily access metadata required.

Signed-off-by: Goutham Veeramachaneni <gouthamve@gmail.com>
---
 scrape/scrape.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/scrape/scrape.go b/scrape/scrape.go
index 33cb3592fa..68c392587f 100644
--- a/scrape/scrape.go
+++ b/scrape/scrape.go
@@ -303,8 +303,11 @@ func newScrapePool(cfg *config.ScrapeConfig, app storage.Appendable, jitterSeed
 		}
 		opts.target.SetMetadataStore(cache)
 
+		// Store the cache in the context.
+		loopCtx := ContextWithMetricMetadataStore(ctx, cache)
+
 		return newScrapeLoop(
-			ctx,
+			loopCtx,
 			opts.scraper,
 			log.With(logger, "target", opts.target),
 			buffers,
@@ -1801,3 +1804,21 @@ func reusableCache(r, l *config.ScrapeConfig) bool {
 	}
 	return reflect.DeepEqual(zeroConfig(r), zeroConfig(l))
 }
+
+// CtxKey is a dedicated type for keys of context-embedded values propagated
+// with the scrape context.
+type ctxKey int
+
+// Valid CtxKey values.
+const (
+	ctxKeyMetadata ctxKey = iota + 1
+)
+
+func ContextWithMetricMetadataStore(ctx context.Context, s MetricMetadataStore) context.Context {
+	return context.WithValue(ctx, ctxKeyMetadata, s)
+}
+
+func MetricMetadataStoreFromContext(ctx context.Context) (MetricMetadataStore, bool) {
+	s, ok := ctx.Value(ctxKeyMetadata).(MetricMetadataStore)
+	return s, ok
+}