Support shutdown watches dynamically

Co-authored-by: FillZpp <FillZpp.pub@gmail.com> Signed-off-by: Tim Ramlot <42113979+inteon@users.noreply.github.com>

Author: inteon

pkg/controller/controller_test.go

@@ -77,7 +77,7 @@ var _ = Describe("controller.Controller", func() {
 
 ctx, cancel := context.WithCancel(context.Background())
 watchChan := make(chan event.GenericEvent, 1)
-watch := &source.Channel{Source: watchChan}
+watch := &source.Channel{Broadcaster: source.NewChannelBroadcaster(watchChan)}
 watchChan <- event.GenericEvent{Object: &corev1.Pod{}}
 
 reconcileStarted := make(chan struct{})

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,8 @@ type FakeInformer struct {
 // RunCount is incremented each time RunInformersAndControllers is called
 RunCount int
 
-handlers []eventHandlerWrapper
+mu sync.RWMutex
+handlers []*eventHandlerWrapper
 }
 
 type modernResourceEventHandler interface {
@@ -51,7 +54,8 @@ type legacyResourceEventHandler interface {
 // eventHandlerWrapper wraps a ResourceEventHandler in a manner that is compatible with client-go 1.27+ and older.
 // The interface was changed in these versions.
 type eventHandlerWrapper struct {
-handler any
+handler any
+hasSynced bool
 }
 
 func (e eventHandlerWrapper) OnAdd(obj interface{}) {
@@ -78,6 +82,10 @@ func (e eventHandlerWrapper) OnDelete(obj interface{}) {
 e.handler.(legacyResourceEventHandler).OnDelete(obj)
 }
 
+func (e eventHandlerWrapper) HasSynced() bool {
+return e.hasSynced
+}
+
 // AddIndexers does nothing. TODO(community): Implement this.
 func (f *FakeInformer) AddIndexers(indexers cache.Indexers) error {
 return nil
@@ -98,10 +106,13 @@ func (f *FakeInformer) HasSynced() bool {
 return f.Synced
 }
 
-// AddEventHandler implements the Informer interface. Adds an EventHandler to the fake Informers. TODO(community): Implement Registration.
+// AddEventHandler implements the Informer interface. Adds an EventHandler to the fake Informers.
 func (f *FakeInformer) AddEventHandler(handler cache.ResourceEventHandler) (cache.ResourceEventHandlerRegistration, error) {
-f.handlers = append(f.handlers, eventHandlerWrapper{handler})
-return nil, nil
+f.mu.Lock()
+defer f.mu.Unlock()
+eh := &eventHandlerWrapper{handler, true}
+f.handlers = append(f.handlers, eh)
+return eh, nil
 }
 
 // Run implements the Informer interface. Increments f.RunCount.
@@ -111,20 +122,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)
 }
@@ -135,8 +152,25 @@ func (f *FakeInformer) AddEventHandlerWithResyncPeriod(handler cache.ResourceEve
 return nil, nil
 }
 
-// RemoveEventHandler does nothing. TODO(community): Implement this.
+// RemoveEventHandler removes an EventHandler to the fake Informers.
 func (f *FakeInformer) RemoveEventHandler(handle cache.ResourceEventHandlerRegistration) error {
+eh, ok := handle.(*eventHandlerWrapper)
+if !ok {
+return fmt.Errorf("invalid registration type %t", handle)
+}
+
+f.mu.Lock()
+defer f.mu.Unlock()
+
+handlers := make([]*eventHandlerWrapper, 0, len(f.handlers))
+for _, h := range f.handlers {
+if h == eh {
+continue
+}
+handlers = append(handlers, h)
+}
+f.handlers = handlers
+
 return nil
 }
 

pkg/internal/controller/controller.go

@@ -59,12 +59,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,
@@ -124,6 +130,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(...).
@@ -133,7 +143,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.
@@ -149,22 +163,43 @@ 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 {
+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.initMetrics()
 
 // Set the internal context.
-c.ctx = ctx
+var cancel context.CancelFunc
+c.ctx, cancel = context.WithCancel(ctx)
+
+wg := &sync.WaitGroup{}
 
 c.Queue = c.MakeQueue()
-go func() {
-<-ctx.Done()
-c.Queue.ShutDown()
+defer func() {
+var startedSources []source.Source
+c.mu.Lock()
+c.Stopped = true
+startedSources = c.startedSources
+c.mu.Unlock()
+
+c.Queue.ShutDown() // Stop receiving new items in the queue
+
+cancel() // cancel the context to stop all the sources
+for _, src := range startedSources {
+if err := src.Shutdown(); 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")
+
+wg.Wait() // Wait for workers to finish
+c.LogConstructor(nil).Info("All workers finished")
 }()
 
-wg := &sync.WaitGroup{}
 err := func() error {
 defer c.mu.Unlock()
 
@@ -180,6 +215,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
@@ -238,8 +274,6 @@ func (c *Controller) Start(ctx context.Context) error {
 
 <-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

@@ -151,7 +151,7 @@ var _ = Describe("controller", func() {
 
 err = ctrl.Start(context.TODO())
 Expect(err).To(HaveOccurred())
-Expect(err.Error()).To(ContainSubstring("failed to wait for testcontroller caches to sync: timed out waiting for cache to be synced"))
+Expect(err.Error()).To(ContainSubstring("timed out trying to get an informer from cache and waiting for cache to be synced for Kind *v1.Deployment"))
 })
 
 It("should not error when context cancelled", func() {
@@ -226,7 +226,7 @@ var _ = Describe("controller", func() {
 Object: p,
 }
 
-ins := &source.Channel{Source: ch}
+ins := &source.Channel{Broadcaster: source.NewChannelBroadcaster(ch)}
 ins.DestBufferSize = 1
 
 // send the event to the channel
@@ -260,7 +260,7 @@ var _ = Describe("controller", func() {
 
 e := ctrl.Start(ctx)
 Expect(e).To(HaveOccurred())
-Expect(e.Error()).To(ContainSubstring("must specify Channel.Source"))
+Expect(e.Error()).To(ContainSubstring("must create Channel with a non-nil Broadcaster"))
 })
 
 It("should call Start on sources with the appropriate EventHandler, Queue, and Predicates", func() {
@@ -308,7 +308,7 @@ var _ = Describe("controller", func() {
 Expect(ctrl.Start(ctx)).To(Succeed())
 err := ctrl.Start(ctx)
 Expect(err).To(HaveOccurred())
-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"))
+Expect(err.Error()).To(Equal("can not restart a stopped controller, you should create a new one"))
 })
 
 })