From 840c86b160f1eb6a1d9f17c6fe43e48bc1d5b273 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 16:30:49 -0500
Subject: [PATCH 1/7] integrations-next: wait for integrations to exit after
 stopping them

---
 CHANGELOG.md                                  |   3 +
 pkg/integrations/v2/controller.go             | 156 +++++++++---------
 .../v2/controller_metricsintegration_test.go  |  39 ++++-
 3 files changed, 119 insertions(+), 79 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0d48aed73b14..42d8e27bb148 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,9 @@
 
 - [BUGFIX] Allow inlining credentials in remote_write url. (@tpaschalis)
 
+- [BUGFIX] integrations-next: Wait for integrations to stop when starting new
+  instances or shutting down (@rfratto).
+
 # v0.22.0 (2022-01-13)
 
 This release has deprecations. Please read [DEPRECATION] entries and consult
diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index 004af87111f4..9d2a24a782e9 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -57,64 +57,108 @@ func newController(l log.Logger, cfg controllerConfig, globals Globals) (*contro
 
 // run starts the controller and blocks until ctx is canceled.
 func (c *controller) run(ctx context.Context) {
+	type worker struct {
+		ci     *controlledIntegration
+		stop   context.CancelFunc
+		exited chan struct{}
+	}
+	var (
+		workersMut sync.Mutex
+		workers    = make(map[*controlledIntegration]worker)
+	)
+
+	runWorker := func(ctx context.Context, w worker) {
+		defer func() {
+			// Have a worker remove itself from the map when it exits. Doing this
+			// allows exited integrations to re-start when the config is reloaded.
+			workersMut.Lock()
+			defer workersMut.Unlock()
+			delete(workers, w.ci)
+		}()
+
+		defer close(w.exited)
+
+		w.ci.running.Store(true)
+		defer w.ci.running.Store(false)
+
+		err := w.ci.i.RunIntegration(ctx)
+		if err != nil {
+			level.Warn(c.logger).Log("msg", "integration exited with error", "instance", w.ci.id, "err", err)
+		}
+	}
+
 	defer func() {
 		level.Debug(c.logger).Log("msg", "stopping all integrations")
 
 		c.mut.Lock()
 		defer c.mut.Unlock()
 
-		for _, exist := range c.integrations {
-			exist.Stop()
+		workersMut.Lock()
+		defer workersMut.Unlock()
+
+		for _, w := range workers {
+			w.stop()
+		}
+		for key, w := range workers {
+			<-w.exited
+			delete(workers, key)
 		}
 	}()
 
-	var currentIntegrations []*controlledIntegration
-
 	updateIntegrations := func() {
 		// Lock the mutex to prevent another set of integrations from being
 		// loaded in.
 		c.mut.Lock()
 		defer c.mut.Unlock()
-		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(currentIntegrations), "new_count", len(c.integrations))
+
+		workersMut.Lock()
+		defer workersMut.Unlock()
+
+		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(workers), "new_count", len(c.integrations))
 
 		newIntegrations := c.integrations
 
-		// Shut down all old integrations. If the integration exists in
-		// newIntegrations but has a different gen number, then there's a new
-		// instance to launch.
-		for _, exist := range currentIntegrations {
+		// Shut down workers whose integrations have gone away.
+		var stopped []worker
+		for ci, w := range workers {
 			var found bool
 			for _, current := range newIntegrations {
-				if exist.id == current.id && current.gen == exist.gen {
+				if ci == current {
 					found = true
 					break
 				}
 			}
 			if !found {
-				exist.Stop()
+				w.stop()
+				stopped = append(stopped, w)
+				delete(workers, ci)
 			}
 		}
+		for _, w := range stopped {
+			// Wait for stopped integrations to fully exit. We do this in a separate
+			// loop so context cancellations can be handled simultaneously, allowing
+			// the wait to complete faster.
+			<-w.exited
+		}
 
-		var waitStarted sync.WaitGroup
-		waitStarted.Add(len(newIntegrations))
-
-		// Now all integrations can be launched.
+		// Spawn new workers.
 		for _, current := range newIntegrations {
-			go func(current *controlledIntegration) {
-				waitStarted.Done()
-
-				err := current.Run(ctx)
-				if err != nil && !errors.Is(err, errIntegrationRunning) {
-					level.Warn(c.logger).Log("msg", "integration exited with error", "instance", current.id, "err", err)
-				}
-			}(current)
-		}
+			w, ok := workers[current]
+			if ok {
+				continue
+			}
 
-		// Wait for all integration goroutines to have been scheduled at least once.
-		waitStarted.Wait()
+			// This integration doesn't have a worker yet; create a new one.
+			workerContext, workerCancel := context.WithCancel(ctx)
 
-		// Finally, store the current list of contolled integrations.
-		currentIntegrations = newIntegrations
+			w = worker{
+				ci:     current,
+				stop:   workerCancel,
+				exited: make(chan struct{}),
+			}
+			go runWorker(workerContext, w)
+			workers[current] = w
+		}
 	}
 
 	for {
@@ -131,59 +175,22 @@ func (c *controller) run(ctx context.Context) {
 	}
 }
 
-// controlledIntegration is a running Integration.
-// A running integration is identified uniquely by its id and gen.
+// controlledIntegration is a running Integration. A running integration is
+// identified uniquely by its id.
 type controlledIntegration struct {
-	id  integrationID
-	gen uint64
-
-	i Integration
-	c Config // Config that generated i. Used for changing to see if a config changed.
-
+	id      integrationID
+	i       Integration
+	c       Config // Config that generated i. Used for changing to see if a config changed.
 	running atomic.Bool
-
-	mut  sync.Mutex
-	stop context.CancelFunc
 }
 
 func (ci *controlledIntegration) Running() bool {
 	return ci.running.Load()
 }
 
-func (ci *controlledIntegration) Run(ctx context.Context) error {
-	updatedRunningState := ci.running.CAS(false, true)
-	if !updatedRunningState {
-		// The CAS will fail if our integration was already running.
-		return errIntegrationRunning
-	}
-	defer ci.running.Store(false)
-
-	ci.mut.Lock()
-	ctx, ci.stop = context.WithCancel(ctx)
-	ci.mut.Unlock()
-
-	// Early optimization: don't do anything if ctx has already been canceled
-	if ctx.Err() != nil {
-		return nil
-	}
-	return ci.i.RunIntegration(ctx)
-}
-
-var errIntegrationRunning = fmt.Errorf("already running")
-
-func (ci *controlledIntegration) Stop() {
-	ci.mut.Lock()
-	if ci.stop != nil {
-		ci.stop()
-	}
-	ci.mut.Unlock()
-}
-
 // integrationID uses a tuple of Name and Identifier to uniquely identify an
 // integration.
-type integrationID struct {
-	Name, Identifier string
-}
+type integrationID struct{ Name, Identifier string }
 
 func (id integrationID) String() string {
 	return fmt.Sprintf("%s/%s", id.Name, id.Identifier)
@@ -260,10 +267,9 @@ NextConfig:
 
 		// Create a new controlled integration.
 		integrations = append(integrations, &controlledIntegration{
-			id:  id,
-			gen: c.gen.Inc(),
-			i:   integration,
-			c:   ic,
+			id: id,
+			i:  integration,
+			c:  ic,
 		})
 	}
 
diff --git a/pkg/integrations/v2/controller_metricsintegration_test.go b/pkg/integrations/v2/controller_metricsintegration_test.go
index 6ef0507c06c7..21d7b4efab2e 100644
--- a/pkg/integrations/v2/controller_metricsintegration_test.go
+++ b/pkg/integrations/v2/controller_metricsintegration_test.go
@@ -1,6 +1,7 @@
 package integrations
 
 import (
+	"context"
 	"testing"
 
 	"github.com/go-kit/log"
@@ -21,7 +22,7 @@ import (
 func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 	integrationWithTarget := func(targetName string) Integration {
 		return mockMetricsIntegration{
-			Integration: NoOpIntegration,
+			Integration: newWaitStartedIntegration(),
 			TargetsFunc: func(Endpoint) []*targetgroup.Group {
 				return []*targetgroup.Group{{
 					Targets: []model.LabelSet{{model.AddressLabel: model.LabelValue(targetName)}},
@@ -40,6 +41,17 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 		}),
 	}
 
+	// waitIntegrations starts a controller and waits for all of its integrations
+	// to run.
+	waitIntegrations := func(t *testing.T, ctrl *controller) {
+		t.Helper()
+		_ = newSyncController(t, ctrl)
+		forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
+			wsi := ci.i.(mockMetricsIntegration).Integration.(*waitStartedIntegration)
+			wsi.trigger.WaitContext(context.Background())
+		})
+	}
+
 	t.Run("All", func(t *testing.T) {
 		ctrl, err := newController(
 			util.TestLogger(t),
@@ -47,7 +59,12 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 			Globals{},
 		)
 		require.NoError(t, err)
-		_ = newSyncController(t, ctrl)
+		waitIntegrations(t, ctrl)
+
+		forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
+			wsi := ci.i.(mockMetricsIntegration).Integration.(*waitStartedIntegration)
+			wsi.trigger.WaitContext(context.Background())
+		})
 
 		result := ctrl.Targets(Endpoint{Prefix: "/"}, TargetOptions{})
 		expect := []*targetGroup{
@@ -64,7 +81,7 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 			Globals{},
 		)
 		require.NoError(t, err)
-		_ = newSyncController(t, ctrl)
+		waitIntegrations(t, ctrl)
 
 		result := ctrl.Targets(Endpoint{Prefix: "/"}, TargetOptions{
 			Integrations: []string{"a", "b"},
@@ -83,7 +100,7 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 			Globals{},
 		)
 		require.NoError(t, err)
-		_ = newSyncController(t, ctrl)
+		waitIntegrations(t, ctrl)
 
 		result := ctrl.Targets(Endpoint{Prefix: "/"}, TargetOptions{
 			Integrations: []string{"a"},
@@ -139,6 +156,20 @@ func Test_controller_MetricsIntegration_ScrapeConfig(t *testing.T) {
 // Tests for controller's utilization of the MetricsIntegration interface.
 //
 
+type waitStartedIntegration struct {
+	trigger *util.WaitTrigger
+}
+
+func newWaitStartedIntegration() *waitStartedIntegration {
+	return &waitStartedIntegration{trigger: util.NewWaitTrigger()}
+}
+
+func (i *waitStartedIntegration) RunIntegration(ctx context.Context) error {
+	i.trigger.Trigger()
+	<-ctx.Done()
+	return nil
+}
+
 type mockMetricsIntegration struct {
 	Integration
 	TargetsFunc       func(ep Endpoint) []*targetgroup.Group

From a2ffaf0111e8c65af185c7afa1ca5b1822c8e960 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 16:42:40 -0500
Subject: [PATCH 2/7] fix lint errors

---
 pkg/integrations/v2/controller.go                      |  8 ++------
 .../v2/controller_metricsintegration_test.go           | 10 +++-------
 2 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index 9d2a24a782e9..def686bab669 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -33,9 +33,6 @@ type controller struct {
 	integrations       []*controlledIntegration // Integrations to run
 	reloadIntegrations chan struct{}            // Inform Controller.Run to re-read integrations
 
-	// Next generation value to use for an integration.
-	gen atomic.Uint64
-
 	// onUpdateDone is used for testing and will be invoked when integrations
 	// finish reloading.
 	onUpdateDone func()
@@ -143,15 +140,14 @@ func (c *controller) run(ctx context.Context) {
 
 		// Spawn new workers.
 		for _, current := range newIntegrations {
-			w, ok := workers[current]
-			if ok {
+			if _, workerExists := workers[current]; workerExists {
 				continue
 			}
 
 			// This integration doesn't have a worker yet; create a new one.
 			workerContext, workerCancel := context.WithCancel(ctx)
 
-			w = worker{
+			w := worker{
 				ci:     current,
 				stop:   workerCancel,
 				exited: make(chan struct{}),
diff --git a/pkg/integrations/v2/controller_metricsintegration_test.go b/pkg/integrations/v2/controller_metricsintegration_test.go
index 21d7b4efab2e..2f04d37899a4 100644
--- a/pkg/integrations/v2/controller_metricsintegration_test.go
+++ b/pkg/integrations/v2/controller_metricsintegration_test.go
@@ -46,10 +46,11 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 	waitIntegrations := func(t *testing.T, ctrl *controller) {
 		t.Helper()
 		_ = newSyncController(t, ctrl)
-		forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
+		err := forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
 			wsi := ci.i.(mockMetricsIntegration).Integration.(*waitStartedIntegration)
-			wsi.trigger.WaitContext(context.Background())
+			_ = wsi.trigger.WaitContext(context.Background())
 		})
+		require.NoError(t, err)
 	}
 
 	t.Run("All", func(t *testing.T) {
@@ -61,11 +62,6 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 		require.NoError(t, err)
 		waitIntegrations(t, ctrl)
 
-		forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
-			wsi := ci.i.(mockMetricsIntegration).Integration.(*waitStartedIntegration)
-			wsi.trigger.WaitContext(context.Background())
-		})
-
 		result := ctrl.Targets(Endpoint{Prefix: "/"}, TargetOptions{})
 		expect := []*targetGroup{
 			{Targets: []model.LabelSet{{model.AddressLabel: "a"}}},

From 22bcc76350b407e8df0da0322a571f5c69c670c5 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 17:02:33 -0500
Subject: [PATCH 3/7] minor refactor

---
 pkg/integrations/v2/controller.go | 90 +++++++++++++++++--------------
 1 file changed, 50 insertions(+), 40 deletions(-)

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index def686bab669..33287040fd97 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -60,48 +60,68 @@ func (c *controller) run(ctx context.Context) {
 		exited chan struct{}
 	}
 	var (
+		runningWorkers sync.WaitGroup
+
 		workersMut sync.Mutex
 		workers    = make(map[*controlledIntegration]worker)
 	)
 
-	runWorker := func(ctx context.Context, w worker) {
-		defer func() {
-			// Have a worker remove itself from the map when it exits. Doing this
-			// allows exited integrations to re-start when the config is reloaded.
-			workersMut.Lock()
-			defer workersMut.Unlock()
-			delete(workers, w.ci)
-		}()
-
-		defer close(w.exited)
-
-		w.ci.running.Store(true)
-		defer w.ci.running.Store(false)
-
-		err := w.ci.i.RunIntegration(ctx)
-		if err != nil {
-			level.Warn(c.logger).Log("msg", "integration exited with error", "instance", w.ci.id, "err", err)
-		}
-	}
-
+	// Shut down all workers on shutdown.
 	defer func() {
-		level.Debug(c.logger).Log("msg", "stopping all integrations")
-
-		c.mut.Lock()
-		defer c.mut.Unlock()
+		defer runningWorkers.Wait()
 
 		workersMut.Lock()
 		defer workersMut.Unlock()
 
+		level.Debug(c.logger).Log("msg", "stopping all integrations")
+
 		for _, w := range workers {
 			w.stop()
 		}
-		for key, w := range workers {
-			<-w.exited
-			delete(workers, key)
-		}
 	}()
 
+	// scheduleWorker starts a new worker for an integration in the background.
+	// The worker will be removed when the integration stops running.
+	//
+	// workersMut should be held while calling this.
+	scheduleWorker := func(ctx context.Context, ci *controlledIntegration) {
+		runningWorkers.Add(1)
+
+		ctx, cancel := context.WithCancel(ctx)
+
+		w := worker{
+			ci:     ci,
+			stop:   cancel,
+			exited: make(chan struct{}),
+		}
+		workers[ci] = w
+
+		go func() {
+			w.ci.running.Store(true)
+
+			// When the integration stops running, we want to free any of our
+			// resources that will notify watchers waiting for the worker to stop.
+			//
+			// Afterwards, we'll block until we remove ourselves from the map; having
+			// an worker remove itself on shutdown allows exited integrations to
+			// re-start when the config is reloaded.
+			defer func() {
+				w.ci.running.Store(false)
+				close(w.exited)
+				runningWorkers.Done()
+
+				workersMut.Lock()
+				defer workersMut.Unlock()
+				delete(workers, ci)
+			}()
+
+			err := ci.i.RunIntegration(ctx)
+			if err != nil {
+				level.Error(c.logger).Log("msg", "integration exited with error", "id", ci.id, "err", err)
+			}
+		}()
+	}
+
 	updateIntegrations := func() {
 		// Lock the mutex to prevent another set of integrations from being
 		// loaded in.
@@ -128,7 +148,6 @@ func (c *controller) run(ctx context.Context) {
 			if !found {
 				w.stop()
 				stopped = append(stopped, w)
-				delete(workers, ci)
 			}
 		}
 		for _, w := range stopped {
@@ -143,17 +162,8 @@ func (c *controller) run(ctx context.Context) {
 			if _, workerExists := workers[current]; workerExists {
 				continue
 			}
-
-			// This integration doesn't have a worker yet; create a new one.
-			workerContext, workerCancel := context.WithCancel(ctx)
-
-			w := worker{
-				ci:     current,
-				stop:   workerCancel,
-				exited: make(chan struct{}),
-			}
-			go runWorker(workerContext, w)
-			workers[current] = w
+			// This integration doesn't have an existing worker; schedule a new one.
+			scheduleWorker(ctx, current)
 		}
 	}
 

From e60270425097133e6605f7068662034861332adb Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 17:15:37 -0500
Subject: [PATCH 4/7] integrations-next: stop holding config mutex for entire
 reload

---
 pkg/integrations/v2/controller.go | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index 33287040fd97..a5d17eca43b6 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -122,19 +122,12 @@ func (c *controller) run(ctx context.Context) {
 		}()
 	}
 
-	updateIntegrations := func() {
-		// Lock the mutex to prevent another set of integrations from being
-		// loaded in.
-		c.mut.Lock()
-		defer c.mut.Unlock()
-
+	updateIntegrations := func(newIntegrations []*controlledIntegration) {
 		workersMut.Lock()
 		defer workersMut.Unlock()
 
 		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(workers), "new_count", len(c.integrations))
 
-		newIntegrations := c.integrations
-
 		// Shut down workers whose integrations have gone away.
 		var stopped []worker
 		for ci, w := range workers {
@@ -173,7 +166,11 @@ func (c *controller) run(ctx context.Context) {
 			level.Debug(c.logger).Log("msg", "controller exiting")
 			return
 		case <-c.reloadIntegrations:
-			updateIntegrations()
+			c.mut.Lock()
+			newIntegrations := c.integrations
+			c.mut.Unlock()
+
+			updateIntegrations(newIntegrations)
 			if c.onUpdateDone != nil {
 				c.onUpdateDone()
 			}

From 58944621b56079b444caafd27cb4df276e96335c Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 17:29:48 -0500
Subject: [PATCH 5/7] make controller.run authoritative over running
 integrations

---
 pkg/integrations/v2/controller.go             | 57 +++++++++----------
 .../v2/controller_metricsintegration_test.go  |  6 +-
 2 files changed, 29 insertions(+), 34 deletions(-)

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index a5d17eca43b6..4a9daeec94da 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -25,13 +25,14 @@ type controllerConfig []Config
 
 // controller manages a set of integrations.
 type controller struct {
-	logger  log.Logger
-	mut     sync.Mutex
-	cfg     controllerConfig
-	globals Globals
+	logger log.Logger
 
-	integrations       []*controlledIntegration // Integrations to run
-	reloadIntegrations chan struct{}            // Inform Controller.Run to re-read integrations
+	mut          sync.Mutex
+	cfg          controllerConfig
+	globals      Globals
+	integrations []*controlledIntegration // Running integrations
+
+	runIntegrations chan []*controlledIntegration // Schedule integrations to run
 
 	// onUpdateDone is used for testing and will be invoked when integrations
 	// finish reloading.
@@ -43,8 +44,8 @@ type controller struct {
 // integrations.
 func newController(l log.Logger, cfg controllerConfig, globals Globals) (*controller, error) {
 	c := &controller{
-		logger:             l,
-		reloadIntegrations: make(chan struct{}, 1),
+		logger:          l,
+		runIntegrations: make(chan []*controlledIntegration, 1),
 	}
 	if err := c.UpdateController(cfg, globals); err != nil {
 		return nil, err
@@ -150,7 +151,7 @@ func (c *controller) run(ctx context.Context) {
 			<-w.exited
 		}
 
-		// Spawn new workers.
+		// Spawn new workers for integrations that don't have them.
 		for _, current := range newIntegrations {
 			if _, workerExists := workers[current]; workerExists {
 				continue
@@ -158,6 +159,11 @@ func (c *controller) run(ctx context.Context) {
 			// This integration doesn't have an existing worker; schedule a new one.
 			scheduleWorker(ctx, current)
 		}
+
+		// Update the set of integrations we're running.
+		c.mut.Lock()
+		defer c.mut.Unlock()
+		c.integrations = newIntegrations
 	}
 
 	for {
@@ -165,11 +171,7 @@ func (c *controller) run(ctx context.Context) {
 		case <-ctx.Done():
 			level.Debug(c.logger).Log("msg", "controller exiting")
 			return
-		case <-c.reloadIntegrations:
-			c.mut.Lock()
-			newIntegrations := c.integrations
-			c.mut.Unlock()
-
+		case newIntegrations := <-c.runIntegrations:
 			updateIntegrations(newIntegrations)
 			if c.onUpdateDone != nil {
 				c.onUpdateDone()
@@ -276,9 +278,8 @@ NextConfig:
 		})
 	}
 
-	// Update integrations and inform
-	c.integrations = integrations
-	c.reloadIntegrations <- struct{}{}
+	// Schedule integrations to run
+	c.runIntegrations <- integrations
 
 	c.cfg = cfg
 	c.globals = globals
@@ -292,9 +293,6 @@ NextConfig:
 // Handler is expensive to compute and should only be done after reloading the
 // config.
 func (c *controller) Handler(prefix string) (http.Handler, error) {
-	c.mut.Lock()
-	defer c.mut.Unlock()
-
 	var firstErr error
 	saveFirstErr := func(err error) {
 		if firstErr == nil {
@@ -304,7 +302,7 @@ func (c *controller) Handler(prefix string) (http.Handler, error) {
 
 	r := mux.NewRouter()
 
-	err := forEachIntegration(c.integrations, prefix, func(ci *controlledIntegration, iprefix string) {
+	err := c.forEachIntegration(prefix, func(ci *controlledIntegration, iprefix string) {
 		id := ci.id
 
 		i, ok := ci.i.(HTTPIntegration)
@@ -340,20 +338,23 @@ func (c *controller) Handler(prefix string) (http.Handler, error) {
 
 // forEachIntegration calculates the prefix for each integration and calls f.
 // prefix will not end in /.
-func forEachIntegration(set []*controlledIntegration, basePrefix string, f func(ci *controlledIntegration, iprefix string)) error {
+func (c *controller) forEachIntegration(basePrefix string, f func(ci *controlledIntegration, iprefix string)) error {
+	c.mut.Lock()
+	defer c.mut.Unlock()
+
 	// Pre-populate a mapping of integration name -> identifier. If there are
 	// two instances of the same integration, we want to ensure unique routing.
 	//
 	// This special logic is done for backwards compatibility with the original
 	// design of integrations.
 	identifiersMap := map[string][]string{}
-	for _, i := range set {
+	for _, i := range c.integrations {
 		identifiersMap[i.id.Name] = append(identifiersMap[i.id.Name], i.id.Identifier)
 	}
 
 	usedPrefixes := map[string]struct{}{}
 
-	for _, ci := range set {
+	for _, ci := range c.integrations {
 		id := ci.id
 		multipleInstances := len(identifiersMap[id.Name]) > 1
 
@@ -388,8 +389,7 @@ func (c *controller) Targets(ep Endpoint, opts TargetOptions) []*targetGroup {
 	}
 	var mm []prefixedMetricsIntegration
 
-	c.mut.Lock()
-	err := forEachIntegration(c.integrations, ep.Prefix, func(ci *controlledIntegration, iprefix string) {
+	err := c.forEachIntegration(ep.Prefix, func(ci *controlledIntegration, iprefix string) {
 		// Best effort liveness check. They might stop running when we request
 		// their targets, which is fine, but we should save as much work as we
 		// can.
@@ -404,7 +404,6 @@ func (c *controller) Targets(ep Endpoint, opts TargetOptions) []*targetGroup {
 	if err != nil {
 		level.Warn(c.logger).Log("msg", "error when iterating over integrations to get targets", "err", err)
 	}
-	c.mut.Unlock()
 
 	var tgs []*targetGroup
 	for _, mi := range mm {
@@ -497,8 +496,7 @@ func (c *controller) ScrapeConfigs(prefix string, sdConfig *http_sd.SDConfig) []
 	}
 	var mm []prefixedMetricsIntegration
 
-	c.mut.Lock()
-	err := forEachIntegration(c.integrations, prefix, func(ci *controlledIntegration, iprefix string) {
+	err := c.forEachIntegration(prefix, func(ci *controlledIntegration, iprefix string) {
 		if mi, ok := ci.i.(MetricsIntegration); ok {
 			mm = append(mm, prefixedMetricsIntegration{id: ci.id, i: mi, prefix: iprefix})
 		}
@@ -506,7 +504,6 @@ func (c *controller) ScrapeConfigs(prefix string, sdConfig *http_sd.SDConfig) []
 	if err != nil {
 		level.Warn(c.logger).Log("msg", "error when iterating over integrations to get scrape configs", "err", err)
 	}
-	c.mut.Unlock()
 
 	var cfgs []*autoscrape.ScrapeConfig
 	for _, mi := range mm {
diff --git a/pkg/integrations/v2/controller_metricsintegration_test.go b/pkg/integrations/v2/controller_metricsintegration_test.go
index 2f04d37899a4..6302c21519fb 100644
--- a/pkg/integrations/v2/controller_metricsintegration_test.go
+++ b/pkg/integrations/v2/controller_metricsintegration_test.go
@@ -46,7 +46,7 @@ func Test_controller_MetricsIntegration_Targets(t *testing.T) {
 	waitIntegrations := func(t *testing.T, ctrl *controller) {
 		t.Helper()
 		_ = newSyncController(t, ctrl)
-		err := forEachIntegration(ctrl.integrations, "/", func(ci *controlledIntegration, _ string) {
+		err := ctrl.forEachIntegration("/", func(ci *controlledIntegration, _ string) {
 			wsi := ci.i.(mockMetricsIntegration).Integration.(*waitStartedIntegration)
 			_ = wsi.trigger.WaitContext(context.Background())
 		})
@@ -136,9 +136,7 @@ func Test_controller_MetricsIntegration_ScrapeConfig(t *testing.T) {
 		Globals{},
 	)
 	require.NoError(t, err)
-	// NOTE(rfratto): we explicitly don't run the controller here because
-	// ScrapeConfigs should return the list of scrape targets even when the
-	// integration isn't running.
+	_ = newSyncController(t, ctrl)
 
 	result := ctrl.ScrapeConfigs("/", &http.DefaultSDConfig)
 	expect := []*autoscrape.ScrapeConfig{

From 5a21cdc1ba8d7831d6d8dc8d15f8b6edd5a11e2e Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 17:32:29 -0500
Subject: [PATCH 6/7] fix log line

---
 pkg/integrations/v2/controller.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index 4a9daeec94da..251c670f57ca 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -127,7 +127,7 @@ func (c *controller) run(ctx context.Context) {
 		workersMut.Lock()
 		defer workersMut.Unlock()
 
-		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(workers), "new_count", len(c.integrations))
+		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(workers), "new_count", len(newIntegrations))
 
 		// Shut down workers whose integrations have gone away.
 		var stopped []worker

From 2b0709e449c7ae6747d187c3dd1444240e468eb4 Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Wed, 26 Jan 2022 17:50:27 -0500
Subject: [PATCH 7/7] move running integrations into a dedicated worker pool

---
 pkg/integrations/v2/controller.go      | 125 ++-----------------------
 pkg/integrations/v2/controller_test.go |  53 +++++------
 pkg/integrations/v2/workers.go         | 122 ++++++++++++++++++++++++
 3 files changed, 150 insertions(+), 150 deletions(-)
 create mode 100644 pkg/integrations/v2/workers.go

diff --git a/pkg/integrations/v2/controller.go b/pkg/integrations/v2/controller.go
index 251c670f57ca..16f80a20245d 100644
--- a/pkg/integrations/v2/controller.go
+++ b/pkg/integrations/v2/controller.go
@@ -33,10 +33,6 @@ type controller struct {
 	integrations []*controlledIntegration // Running integrations
 
 	runIntegrations chan []*controlledIntegration // Schedule integrations to run
-
-	// onUpdateDone is used for testing and will be invoked when integrations
-	// finish reloading.
-	onUpdateDone func()
 }
 
 // newController creates a new Controller. Controller is intended to be
@@ -55,116 +51,8 @@ func newController(l log.Logger, cfg controllerConfig, globals Globals) (*contro
 
 // run starts the controller and blocks until ctx is canceled.
 func (c *controller) run(ctx context.Context) {
-	type worker struct {
-		ci     *controlledIntegration
-		stop   context.CancelFunc
-		exited chan struct{}
-	}
-	var (
-		runningWorkers sync.WaitGroup
-
-		workersMut sync.Mutex
-		workers    = make(map[*controlledIntegration]worker)
-	)
-
-	// Shut down all workers on shutdown.
-	defer func() {
-		defer runningWorkers.Wait()
-
-		workersMut.Lock()
-		defer workersMut.Unlock()
-
-		level.Debug(c.logger).Log("msg", "stopping all integrations")
-
-		for _, w := range workers {
-			w.stop()
-		}
-	}()
-
-	// scheduleWorker starts a new worker for an integration in the background.
-	// The worker will be removed when the integration stops running.
-	//
-	// workersMut should be held while calling this.
-	scheduleWorker := func(ctx context.Context, ci *controlledIntegration) {
-		runningWorkers.Add(1)
-
-		ctx, cancel := context.WithCancel(ctx)
-
-		w := worker{
-			ci:     ci,
-			stop:   cancel,
-			exited: make(chan struct{}),
-		}
-		workers[ci] = w
-
-		go func() {
-			w.ci.running.Store(true)
-
-			// When the integration stops running, we want to free any of our
-			// resources that will notify watchers waiting for the worker to stop.
-			//
-			// Afterwards, we'll block until we remove ourselves from the map; having
-			// an worker remove itself on shutdown allows exited integrations to
-			// re-start when the config is reloaded.
-			defer func() {
-				w.ci.running.Store(false)
-				close(w.exited)
-				runningWorkers.Done()
-
-				workersMut.Lock()
-				defer workersMut.Unlock()
-				delete(workers, ci)
-			}()
-
-			err := ci.i.RunIntegration(ctx)
-			if err != nil {
-				level.Error(c.logger).Log("msg", "integration exited with error", "id", ci.id, "err", err)
-			}
-		}()
-	}
-
-	updateIntegrations := func(newIntegrations []*controlledIntegration) {
-		workersMut.Lock()
-		defer workersMut.Unlock()
-
-		level.Debug(c.logger).Log("msg", "updating running integrations", "prev_count", len(workers), "new_count", len(newIntegrations))
-
-		// Shut down workers whose integrations have gone away.
-		var stopped []worker
-		for ci, w := range workers {
-			var found bool
-			for _, current := range newIntegrations {
-				if ci == current {
-					found = true
-					break
-				}
-			}
-			if !found {
-				w.stop()
-				stopped = append(stopped, w)
-			}
-		}
-		for _, w := range stopped {
-			// Wait for stopped integrations to fully exit. We do this in a separate
-			// loop so context cancellations can be handled simultaneously, allowing
-			// the wait to complete faster.
-			<-w.exited
-		}
-
-		// Spawn new workers for integrations that don't have them.
-		for _, current := range newIntegrations {
-			if _, workerExists := workers[current]; workerExists {
-				continue
-			}
-			// This integration doesn't have an existing worker; schedule a new one.
-			scheduleWorker(ctx, current)
-		}
-
-		// Update the set of integrations we're running.
-		c.mut.Lock()
-		defer c.mut.Unlock()
-		c.integrations = newIntegrations
-	}
+	pool := newWorkerPool(ctx, c.logger)
+	defer pool.Close()
 
 	for {
 		select {
@@ -172,10 +60,11 @@ func (c *controller) run(ctx context.Context) {
 			level.Debug(c.logger).Log("msg", "controller exiting")
 			return
 		case newIntegrations := <-c.runIntegrations:
-			updateIntegrations(newIntegrations)
-			if c.onUpdateDone != nil {
-				c.onUpdateDone()
-			}
+			pool.Reload(newIntegrations)
+
+			c.mut.Lock()
+			c.integrations = newIntegrations
+			c.mut.Unlock()
 		}
 	}
 }
diff --git a/pkg/integrations/v2/controller_test.go b/pkg/integrations/v2/controller_test.go
index 46b1706de47e..ab342c821955 100644
--- a/pkg/integrations/v2/controller_test.go
+++ b/pkg/integrations/v2/controller_test.go
@@ -134,57 +134,46 @@ func Test_controller_ConfigChanges(t *testing.T) {
 }
 
 type syncController struct {
-	inner   *controller
-	applyWg sync.WaitGroup
-
-	stop     context.CancelFunc
-	exitedCh chan struct{}
+	inner *controller
+	pool  *workerPool
 }
 
-// newSyncController makes calls to Controller synchronous. newSyncController
-// will start running the inner controller and wait for it to update.
+// newSyncController pairs an unstarted controller with a manually managed
+// worker pool to synchronously apply integrations.
 func newSyncController(t *testing.T, inner *controller) *syncController {
 	t.Helper()
 
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(func() {
-		cancel()
-	})
-
 	sc := &syncController{
-		inner:    inner,
-		stop:     cancel,
-		exitedCh: make(chan struct{}),
+		inner: inner,
+		pool:  newWorkerPool(context.Background(), inner.logger),
 	}
-	inner.onUpdateDone = sc.applyWg.Done // Inform WG whenever an apply finishes
 
-	// There's always immediately ony applied queued from any successfully created controller.
-	sc.applyWg.Add(1)
+	// There's always immediately one queued integration set from any
+	// successfully created controller.
+	sc.refresh()
+	return sc
+}
 
-	go func() {
-		inner.run(ctx)
-		close(sc.exitedCh)
-	}()
+func (sc *syncController) refresh() {
+	sc.inner.mut.Lock()
+	defer sc.inner.mut.Unlock()
 
-	sc.applyWg.Wait()
-	return sc
+	newIntegrations := <-sc.inner.runIntegrations
+	sc.pool.Reload(newIntegrations)
+	sc.inner.integrations = newIntegrations
 }
 
 func (sc *syncController) UpdateController(c controllerConfig, g Globals) error {
-	sc.applyWg.Add(1)
-
-	if err := sc.inner.UpdateController(c, g); err != nil {
-		sc.applyWg.Done() // The wg won't ever be finished now
+	err := sc.inner.UpdateController(c, g)
+	if err != nil {
 		return err
 	}
-
-	sc.applyWg.Wait()
+	sc.refresh()
 	return nil
 }
 
 func (sc *syncController) Stop() {
-	sc.stop()
-	<-sc.exitedCh
+	sc.pool.Close()
 }
 
 const mockIntegrationName = "mock"
diff --git a/pkg/integrations/v2/workers.go b/pkg/integrations/v2/workers.go
new file mode 100644
index 000000000000..805cf37a45c0
--- /dev/null
+++ b/pkg/integrations/v2/workers.go
@@ -0,0 +1,122 @@
+package integrations
+
+import (
+	"context"
+	"sync"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+type workerPool struct {
+	log       log.Logger
+	parentCtx context.Context
+
+	mut     sync.Mutex
+	workers map[*controlledIntegration]worker
+
+	runningWorkers sync.WaitGroup
+}
+
+type worker struct {
+	ci     *controlledIntegration
+	stop   context.CancelFunc
+	exited chan struct{}
+}
+
+func newWorkerPool(ctx context.Context, l log.Logger) *workerPool {
+	return &workerPool{
+		log:       l,
+		parentCtx: ctx,
+
+		workers: make(map[*controlledIntegration]worker),
+	}
+}
+
+func (p *workerPool) Reload(newIntegrations []*controlledIntegration) {
+	p.mut.Lock()
+	defer p.mut.Unlock()
+
+	level.Debug(p.log).Log("msg", "updating running integrations", "prev_count", len(p.workers), "new_count", len(newIntegrations))
+
+	// Shut down workers whose integrations have gone away.
+	var stopped []worker
+	for ci, w := range p.workers {
+		var found bool
+		for _, current := range newIntegrations {
+			if ci == current {
+				found = true
+				break
+			}
+		}
+		if !found {
+			w.stop()
+			stopped = append(stopped, w)
+		}
+	}
+	for _, w := range stopped {
+		// Wait for stopped integrations to fully exit. We do this in a separate
+		// loop so context cancellations can be handled simultaneously, allowing
+		// the wait to complete faster.
+		<-w.exited
+	}
+
+	// Spawn new workers for integrations that don't have them.
+	for _, current := range newIntegrations {
+		if _, workerExists := p.workers[current]; workerExists {
+			continue
+		}
+		// This integration doesn't have an existing worker; schedule a new one.
+		p.scheduleWorker(current)
+	}
+}
+
+func (p *workerPool) Close() {
+	p.mut.Lock()
+	defer p.mut.Unlock()
+
+	level.Debug(p.log).Log("msg", "stopping all integrations")
+
+	defer p.runningWorkers.Wait()
+	for _, w := range p.workers {
+		w.stop()
+	}
+}
+
+func (p *workerPool) scheduleWorker(ci *controlledIntegration) {
+	p.runningWorkers.Add(1)
+
+	ctx, cancel := context.WithCancel(p.parentCtx)
+
+	w := worker{
+		ci:     ci,
+		stop:   cancel,
+		exited: make(chan struct{}),
+	}
+	p.workers[ci] = w
+
+	go func() {
+		ci.running.Store(true)
+
+		// When the integration stops running, we want to free any of our
+		// resources that will notify watchers waiting for the worker to stop.
+		//
+		// Afterwards, we'll block until we remove ourselves from the map; having
+		// an worker remove itself on shutdown allows exited integrations to
+		// re-start when the config is reloaded.
+		defer func() {
+			ci.running.Store(false)
+			close(w.exited)
+			p.runningWorkers.Done()
+
+			p.mut.Lock()
+			defer p.mut.Unlock()
+			delete(p.workers, ci)
+		}()
+
+		err := ci.i.RunIntegration(ctx)
+		if err != nil {
+			level.Error(p.log).Log("msg", "integration exited with error", "id", ci.id, "err", err)
+		}
+	}()
+}