scrape pool doesn't rely on context as Stop() needs to be blocking to prevent...

scrape pool doesn't rely on context as Stop() needs to be blocking to prevent Scrape loops trying to write to a closed TSDB storage.

cmd/prometheus/main.go

@@ -231,16 +231,24 @@ func main() {
 	cfg.queryEngine.Logger = log.With(logger, "component", "query engine")
 	var (
-		notifier                      = notifier.New(&cfg.notifier, log.With(logger, "component", "notifier"))
-		ctxDiscovery, cancelDiscovery = context.WithCancel(context.Background())
-		discoveryManager              = discovery.NewDiscoveryManager(ctxDiscovery, log.With(logger, "component", "discovery manager"))
-		ctxScrape, cancelScrape       = context.WithCancel(context.Background())
-		scrapeManager                 = retrieval.NewScrapeManager(ctxScrape, log.With(logger, "component", "scrape manager"), fanoutStorage)
-		queryEngine                   = promql.NewEngine(fanoutStorage, &cfg.queryEngine)
 		ctxWeb, cancelWeb             = context.WithCancel(context.Background())
-		webHandler                    = web.New(log.With(logger, "component", "web"), &cfg.web)
+		ctxDiscovery, cancelDiscovery = context.WithCancel(context.Background())
+		ctxRule                       = context.Background()
+
+		notifier         = notifier.New(&cfg.notifier, log.With(logger, "component", "notifier"))
+		discoveryManager = discovery.NewDiscoveryManager(ctxDiscovery, log.With(logger, "component", "discovery manager"))
+		scrapeManager    = retrieval.NewScrapeManager(log.With(logger, "component", "scrape manager"), fanoutStorage)
+		queryEngine      = promql.NewEngine(fanoutStorage, &cfg.queryEngine)
+		ruleManager      = rules.NewManager(&rules.ManagerOptions{Appendable: fanoutStorage,
+			Notifier:    notifier,
+			QueryEngine: queryEngine,
+			Context:     ctxRule,
+			ExternalURL: cfg.web.ExternalURL,
+			Logger:      log.With(logger, "component", "rule manager"),
+		})
 	)
+<<<<<<< HEAD
 	ctx := context.Background()
 	ruleManager := rules.NewManager(&rules.ManagerOptions{
 		Appendable:  fanoutStorage,
@@ -253,6 +261,9 @@ func main() {
 	})
 	cfg.web.Context = ctx
+=======
+	cfg.web.Context = ctxWeb
+>>>>>>> 95b1dec3... scrape pool doesn't rely on context as Stop() needs to be blocking to prevent Scrape loops trying to write to a closed TSDB storage.
 	cfg.web.TSDB = localStorage.Get
 	cfg.web.Storage = fanoutStorage
 	cfg.web.QueryEngine = queryEngine
@@ -274,6 +285,9 @@ func main() {
 		cfg.web.Flags[f.Name] = f.Value.String()
 	}
+	// Depend on cfg.web.ScrapeManager so needs to be after cfg.web.ScrapeManager = scrapeManager
+	webHandler := web.New(log.With(logger, "component", "web"), &cfg.web)
+
 	// Monitor outgoing connections on default transport with conntrack.
 	http.DefaultTransport.(*http.Transport).DialContext = conntrack.NewDialContextFunc(
 		conntrack.DialWithTracing(),
@@ -310,17 +324,6 @@ func main() {
 	var g group.Group
 	{
-		g.Add(
-			func() error {
-				err := discoveryManager.Run()
-				level.Info(logger).Log("msg", "Discovery manager stopped")
-				return err
-			},
-			func(err error) {
-				level.Info(logger).Log("msg", "Stopping discovery manager...")
-				cancelDiscovery()
-			},
-		)
 		term := make(chan os.Signal)
 		signal.Notify(term, os.Interrupt, syscall.SIGTERM)
 		cancel := make(chan struct{})
@@ -341,6 +344,34 @@ func main() {
 			},
 		)
 	}
+	{
+		g.Add(
+			func() error {
+				err := discoveryManager.Run()
+				level.Info(logger).Log("msg", "Discovery manager stopped")
+				return err
+			},
+			func(err error) {
+				level.Info(logger).Log("msg", "Stopping discovery manager...")
+				cancelDiscovery()
+			},
+		)
+	}
+	{
+		g.Add(
+			func() error {
+				err := scrapeManager.Run(discoveryManager.SyncCh())
+				level.Info(logger).Log("msg", "Scrape manager stopped")
+				return err
+			},
+			func(err error) {
+				// Scrape manager needs to be stopped before closing the local TSDB
+				// so that it doesn't try to write samples to a closed storage.
+				level.Info(logger).Log("msg", "Stopping scrape manager...")
+				scrapeManager.Stop()
+			},
+		)
+	}
 	{
 		// Make sure that sighup handler is registered with a redirect to the channel before the potentially
 		// long and synchronous tsdb init.
@@ -482,19 +513,6 @@ func main() {
 			},
 		)
 	}
-	{
-		g.Add(
-			func() error {
-				err := scrapeManager.Run(discoveryManager.SyncCh())
-				level.Info(logger).Log("msg", "Scrape manager stopped")
-				return err
-			},
-			func(err error) {
-				level.Info(logger).Log("msg", "Stopping scrape manager...")
-				cancelScrape()
-			},
-		)
-	}
 	if err := g.Run(); err != nil {
 		level.Error(logger).Log("err", err)
 	}

discovery/manager.go

@@ -87,7 +87,6 @@ func (m *DiscoveryManager) Run() error {
 			return m.ctx.Err()
 		}
 	}
-
 }
 // SyncCh returns a read only channel used by all DiscoveryProviders targetSet updates

retrieval/manager.go

@@ -14,7 +14,6 @@
 package retrieval
 import (
-	"context"
 	"fmt"
 	"github.com/go-kit/kit/log"
@@ -30,27 +29,27 @@ type Appendable interface {
 }
 // NewScrapeManager is the ScrapeManager constructor
-func NewScrapeManager(ctx context.Context, logger log.Logger, app Appendable) *ScrapeManager {
+func NewScrapeManager(logger log.Logger, app Appendable) *ScrapeManager {
 	return &ScrapeManager{
-		ctx:           ctx,
 		append:        app,
 		logger:        logger,
 		actionCh:      make(chan func()),
 		scrapeConfigs: make(map[string]*config.ScrapeConfig),
 		scrapePools:   make(map[string]*scrapePool),
+		graceShut:     make(chan struct{}),
 	}
 }
 // ScrapeManager maintains a set of scrape pools and manages start/stop cicles
 // when receiving new target groups form the discovery manager.
 type ScrapeManager struct {
-	ctx           context.Context
 	logger        log.Logger
 	append        Appendable
 	scrapeConfigs map[string]*config.ScrapeConfig
 	scrapePools   map[string]*scrapePool
 	actionCh      chan func()
+	graceShut     chan struct{}
 }
 // Run starts background processing to handle target updates and reload the scraping loops.
@@ -63,12 +62,20 @@ func (m *ScrapeManager) Run(tsets <-chan map[string][]*config.TargetGroup) error
 			f()
 		case ts := <-tsets:
 			m.reload(ts)
-		case <-m.ctx.Done():
-			return m.ctx.Err()
+		case <-m.graceShut:
+			return nil
 		}
 	}
 }
+// Stop cancels all running scrape pools and blocks until all have exited.
+func (m *ScrapeManager) Stop() {
+	for _, sp := range m.scrapePools {
+		sp.stop()
+	}
+	close(m.graceShut)
+}
+
 // ApplyConfig resets the manager's target providers and job configurations as defined by the new cfg.
 func (m *ScrapeManager) ApplyConfig(cfg *config.Config) error {
 	done := make(chan struct{})
@@ -123,10 +130,10 @@ func (m *ScrapeManager) reload(t map[string][]*config.TargetGroup) error {
 			return fmt.Errorf("target set '%v' doesn't have valid config", tsetName)
 		}
-		// scrape pool doesn't exist so start a new one
+		// Scrape pool doesn't exist so start a new one.
 		existing, ok := m.scrapePools[tsetName]
 		if !ok {
-			sp := newScrapePool(m.ctx, scrapeConfig, m.append, log.With(m.logger, "scrape_pool", tsetName))
+			sp := newScrapePool(scrapeConfig, m.append, log.With(m.logger, "scrape_pool", tsetName))
 			m.scrapePools[tsetName] = sp
 			sp.Sync(tgroup)
@@ -134,7 +141,7 @@ func (m *ScrapeManager) reload(t map[string][]*config.TargetGroup) error {
 			existing.Sync(tgroup)
 		}
-		// cleanup - check config and cancel the scrape loops if it don't exist in the scrape config
+		// Cleanup - check the config and cancel the scrape loops if it don't exist in the scrape config.
 		jobs := make(map[string]struct{})
 		for k := range m.scrapeConfigs {

retrieval/scrape.go

@@ -117,7 +117,6 @@ func init() {
 type scrapePool struct {
 	appendable Appendable
 	logger     log.Logger
-	ctx        context.Context
 	mtx    sync.RWMutex
 	config *config.ScrapeConfig
@@ -136,7 +135,7 @@ const maxAheadTime = 10 * time.Minute
 type labelsMutator func(labels.Labels) labels.Labels
-func newScrapePool(ctx context.Context, cfg *config.ScrapeConfig, app Appendable, logger log.Logger) *scrapePool {
+func newScrapePool(cfg *config.ScrapeConfig, app Appendable, logger log.Logger) *scrapePool {
 	if logger == nil {
 		logger = log.NewNopLogger()
 	}
@@ -152,14 +151,15 @@ func newScrapePool(ctx context.Context, cfg *config.ScrapeConfig, app Appendable
 	sp := &scrapePool{
 		appendable: app,
 		config:     cfg,
-		ctx:        ctx,
 		client:     client,
 		targets:    map[uint64]*Target{},
 		loops:      map[uint64]loop{},
 		logger:     logger,
 	}
 	sp.newLoop = func(t *Target, s scraper) loop {
-		return newScrapeLoop(sp.ctx, s,
+		return newScrapeLoop(
+			context.Background(),
+			s,
 			log.With(logger, "target", t),
 			buffers,
 			func(l labels.Labels) labels.Labels { return sp.mutateSampleLabels(l, t) },
@@ -189,7 +189,6 @@ func (sp *scrapePool) stop() {
 		delete(sp.loops, fp)
 		delete(sp.targets, fp)
 	}
-
 	wg.Wait()
 }
@@ -582,8 +581,7 @@ func (c *scrapeCache) forEachStale(f func(labels.Labels) bool) {
 	}
 }
-func newScrapeLoop(
-	ctx context.Context,
+func newScrapeLoop(ctx context.Context,
 	sc scraper,
 	l log.Logger,
 	buffers *pool.BytesPool,
@@ -605,8 +603,8 @@ func newScrapeLoop(
 		sampleMutator:       sampleMutator,
 		reportSampleMutator: reportSampleMutator,
 		stopped:             make(chan struct{}),
-		ctx:                 ctx,
 		l:                   l,
+		ctx:                 ctx,
 	}
 	sl.scrapeCtx, sl.cancel = context.WithCancel(ctx)

retrieval/scrape_test.go

@@ -44,7 +44,7 @@ func TestNewScrapePool(t *testing.T) {
 	var (
 		app = &nopAppendable{}
 		cfg = &config.ScrapeConfig{}
-		sp  = newScrapePool(context.Background(), cfg, app, nil)
+		sp  = newScrapePool(cfg, app, nil)
 	)
 	if a, ok := sp.appendable.(*nopAppendable); !ok || a != app {
@@ -231,7 +231,7 @@ func TestScrapePoolReload(t *testing.T) {
 func TestScrapePoolAppender(t *testing.T) {
 	cfg := &config.ScrapeConfig{}
 	app := &nopAppendable{}
-	sp := newScrapePool(context.Background(), cfg, app, nil)
+	sp := newScrapePool(cfg, app, nil)
 	wrapped := sp.appender()