From 60ef2016d508545f0e8b575a954899cae37a96e2 Mon Sep 17 00:00:00 2001
From: Krasi Georgiev <krasi.root@gmail.com>
Date: Sun, 3 Dec 2017 17:14:08 +0000
Subject: [PATCH] add a cancel func to the scrape pool as it is needed in the
 scrape loop select block

---
 retrieval/scrape.go      | 6 +++++-
 retrieval/scrape_test.go | 1 +
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/retrieval/scrape.go b/retrieval/scrape.go
index 62caadade..d07663e9e 100644
--- a/retrieval/scrape.go
+++ b/retrieval/scrape.go
@@ -126,6 +126,7 @@ type scrapePool struct {
 	targets        map[uint64]*Target
 	droppedTargets []*Target
 	loops          map[uint64]loop
+	cancel         context.CancelFunc
 
 	// Constructor for new scrape loops. This is settable for testing convenience.
 	newLoop func(*Target, scraper) loop
@@ -148,7 +149,9 @@ func newScrapePool(cfg *config.ScrapeConfig, app Appendable, logger log.Logger)
 
 	buffers := pool.NewBytesPool(163, 100e6, 3)
 
+	ctx, cancel := context.WithCancel(context.Background())
 	sp := &scrapePool{
+		cancel:     cancel,
 		appendable: app,
 		config:     cfg,
 		client:     client,
@@ -158,7 +161,7 @@ func newScrapePool(cfg *config.ScrapeConfig, app Appendable, logger log.Logger)
 	}
 	sp.newLoop = func(t *Target, s scraper) loop {
 		return newScrapeLoop(
-			context.Background(),
+			ctx,
 			s,
 			log.With(logger, "target", t),
 			buffers,
@@ -173,6 +176,7 @@ func newScrapePool(cfg *config.ScrapeConfig, app Appendable, logger log.Logger)
 
 // stop terminates all scrape loops and returns after they all terminated.
 func (sp *scrapePool) stop() {
+	sp.cancel()
 	var wg sync.WaitGroup
 
 	sp.mtx.Lock()
diff --git a/retrieval/scrape_test.go b/retrieval/scrape_test.go
index cc2cfc121..b3a0466ab 100644
--- a/retrieval/scrape_test.go
+++ b/retrieval/scrape_test.go
@@ -75,6 +75,7 @@ func TestScrapePoolStop(t *testing.T) {
 	sp := &scrapePool{
 		targets: map[uint64]*Target{},
 		loops:   map[uint64]loop{},
+		cancel:  func() {},
 	}
 	var mtx sync.Mutex
 	stopped := map[uint64]bool{}