Support shutdown controllers and watches dynamically

Co-authored-by: FillZpp <[email protected]>
Signed-off-by: Tim Ramlot <[email protected]>

Author: inteon

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/prometheus/client_model v0.3.0
 	go.uber.org/goleak v1.2.0
 	go.uber.org/zap v1.24.0
+	golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4
 	golang.org/x/sys v0.4.0
 	golang.org/x/time v0.3.0
 	gomodules.xyz/jsonpatch/v2 v2.2.0

go.sum

@@ -371,6 +371,8 @@ golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4 h1:uVc8UZUe6tr40fFVnUP5Oj+veunVezqYl9z7DYw9xzw=
+golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

pkg/controller/controllertest/util.go

@@ -17,6 +17,8 @@ limitations under the License.
 package controllertest
 
 import (
+	"fmt"
+	"sync"
 	"time"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -33,7 +35,12 @@ type FakeInformer struct {
 	// RunCount is incremented each time RunInformersAndControllers is called
 	RunCount int
 
-	handlers []cache.ResourceEventHandler
+	handlers []*listenHandler
+	mu       sync.RWMutex
+}
+
+type listenHandler struct {
+	cache.ResourceEventHandler
 }
 
 // AddIndexers does nothing.  TODO(community): Implement this.
@@ -58,8 +65,11 @@ func (f *FakeInformer) HasSynced() bool {
 
 // AddEventHandler implements the Informer interface.  Adds an EventHandler to the fake Informers. TODO(community): Implement Registration.
 func (f *FakeInformer) AddEventHandler(handler cache.ResourceEventHandler) (cache.ResourceEventHandlerRegistration, error) {
-	f.handlers = append(f.handlers, handler)
-	return nil, nil
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	lh := &listenHandler{ResourceEventHandler: handler}
+	f.handlers = append(f.handlers, lh)
+	return lh, nil
 }
 
 // Run implements the Informer interface.  Increments f.RunCount.
@@ -69,20 +79,26 @@ func (f *FakeInformer) Run(<-chan struct{}) {
 
 // Add fakes an Add event for obj.
 func (f *FakeInformer) Add(obj metav1.Object) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
 	for _, h := range f.handlers {
 		h.OnAdd(obj)
 	}
 }
 
 // Update fakes an Update event for obj.
 func (f *FakeInformer) Update(oldObj, newObj metav1.Object) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
 	for _, h := range f.handlers {
 		h.OnUpdate(oldObj, newObj)
 	}
 }
 
 // Delete fakes an Delete event for obj.
 func (f *FakeInformer) Delete(obj metav1.Object) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
 	for _, h := range f.handlers {
 		h.OnDelete(obj)
 	}
@@ -95,6 +111,21 @@ func (f *FakeInformer) AddEventHandlerWithResyncPeriod(handler cache.ResourceEve
 
 // RemoveEventHandler does nothing.  TODO(community): Implement this.
 func (f *FakeInformer) RemoveEventHandler(handle cache.ResourceEventHandlerRegistration) error {
+	lh, ok := handle.(*listenHandler)
+	if !ok {
+		return fmt.Errorf("invalid key type %t", handle)
+	}
+
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	handlers := make([]*listenHandler, 0, len(f.handlers))
+	for _, h := range f.handlers {
+		if h == lh {
+			continue
+		}
+		handlers = append(handlers, h)
+	}
+	f.handlers = handlers
 	return nil
 }
 

pkg/internal/controller/controller.go

@@ -18,7 +18,6 @@ package controller
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -58,12 +57,18 @@ type Controller struct {
 	// the Queue for processing
 	Queue workqueue.RateLimitingInterface
 
+	// startedSources maintains a list of sources that have already started.
+	startedSources []source.Source
+
 	// mu is used to synchronize Controller setup
 	mu sync.Mutex
 
 	// Started is true if the Controller has been Started
 	Started bool
 
+	// Stopped is true if the Controller has been Stopped
+	Stopped bool
+
 	// ctx is the context that was passed to Start() and used when starting watches.
 	//
 	// According to the docs, contexts should not be stored in a struct: https://golang.org/pkg/context,
@@ -123,6 +128,10 @@ func (c *Controller) Watch(src source.Source, evthdler handler.EventHandler, prc
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.Stopped {
+		return fmt.Errorf("can not start watch in a stopped controller")
+	}
+
 	// Controller hasn't started yet, store the watches locally and return.
 	//
 	// These watches are going to be held on the controller struct until the manager or user calls Start(...).
@@ -132,7 +141,11 @@ func (c *Controller) Watch(src source.Source, evthdler handler.EventHandler, prc
 	}
 
 	c.LogConstructor(nil).Info("Starting EventSource", "source", src)
-	return src.Start(c.ctx, evthdler, c.Queue, prct...)
+	err := src.Start(c.ctx, evthdler, c.Queue, prct...)
+	if err == nil {
+		c.startedSources = append(c.startedSources, src)
+	}
+	return err
 }
 
 // NeedLeaderElection implements the manager.LeaderElectionRunnable interface.
@@ -148,23 +161,23 @@ func (c *Controller) Start(ctx context.Context) error {
 	// use an IIFE to get proper lock handling
 	// but lock outside to get proper handling of the queue shutdown
 	c.mu.Lock()
+	if c.Stopped {
+		c.mu.Unlock()
+		return fmt.Errorf("can not restart a stopped controller, you should create a new one")
+	}
 	if c.Started {
-		return errors.New("controller was started more than once. This is likely to be caused by being added to a manager multiple times")
+		c.mu.Unlock()
+		return fmt.Errorf("controller was started more than once. This is likely to be caused by being added to a manager multiple times")
 	}
 
 	c.initMetrics()
 
-	// Set the internal context.
+	// Set the internal context, which is used to start watches after the controller is started.
 	c.ctx = ctx
 
 	c.Queue = c.MakeQueue()
-	go func() {
-		<-ctx.Done()
-		c.Queue.ShutDown()
-	}()
-
 	wg := &sync.WaitGroup{}
-	err := func() error {
+	startErr := func() error {
 		defer c.mu.Unlock()
 
 		// TODO(pwittrock): Reconsider HandleCrash
@@ -179,6 +192,7 @@ func (c *Controller) Start(ctx context.Context) error {
 			if err := watch.src.Start(ctx, watch.handler, c.Queue, watch.predicates...); err != nil {
 				return err
 			}
+			c.startedSources = append(c.startedSources, watch.src)
 		}
 
 		// Start the SharedIndexInformer factories to begin populating the SharedIndexInformer caches
@@ -231,14 +245,31 @@ func (c *Controller) Start(ctx context.Context) error {
 		c.Started = true
 		return nil
 	}()
-	if err != nil {
-		return err
+
+	defer func() {
+		c.LogConstructor(nil).Info("Shutdown signal received, waiting for all workers to finish")
+		c.Queue.ShutDown()
+		wg.Wait()
+		c.LogConstructor(nil).Info("All workers finished")
+
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		c.Stopped = true
+
+		for _, src := range c.startedSources {
+			if err := src.Stop(); err != nil {
+				c.LogConstructor(nil).Error(err, "Failed to stop watch source when controller stopping", "source", src)
+			}
+		}
+		c.LogConstructor(nil).Info("All watch sources finished")
+	}()
+
+	if startErr != nil {
+		return startErr
 	}
 
 	<-ctx.Done()
-	c.LogConstructor(nil).Info("Shutdown signal received, waiting for all workers to finish")
-	wg.Wait()
-	c.LogConstructor(nil).Info("All workers finished")
+
 	return nil
 }
 

pkg/internal/controller/controller_test.go

@@ -304,13 +304,47 @@ var _ = Describe("controller", func() {
 		It("should return an error if it gets started more than once", func() {
 			// Use a cancelled context so Start doesn't block
 			ctx, cancel := context.WithCancel(context.Background())
-			cancel()
-			Expect(ctrl.Start(ctx)).To(BeNil())
+			started, done := make(chan struct{}), make(chan struct{})
+			go func() {
+				defer close(done)
+				close(started)
+				Expect(ctrl.Start(ctx)).To(BeNil())
+			}()
+			defer func() { <-done }()
+			<-started
+			defer cancel()
 			err := ctrl.Start(ctx)
 			Expect(err).NotTo(BeNil())
 			Expect(err.Error()).To(Equal("controller was started more than once. This is likely to be caused by being added to a manager multiple times"))
 		})
 
+		It("should return an error if it gets start after stopped", func() {
+			stoppedChan := make(chan struct{})
+			controllerCtx, controllerCancel := context.WithCancel(context.Background())
+			go func() {
+				Expect(ctrl.Start(controllerCtx)).To(BeNil())
+				close(stoppedChan)
+			}()
+
+			// wait for started
+			var started bool
+			for !started {
+				func() {
+					ctrl.mu.Lock()
+					defer ctrl.mu.Unlock()
+					started = ctrl.Started
+				}()
+			}
+
+			controllerCancel()
+			<-stoppedChan
+			Expect(ctrl.Stopped).To(Equal(true))
+
+			err := ctrl.Start(context.TODO())
+			Expect(err).NotTo(BeNil())
+			Expect(err.Error()).To(Equal("can not restart a stopped controller, you should create a new one"))
+		})
+
 	})
 
 	Describe("Processing queue items from a Controller", func() {