[Warm Replicas] Implement warm replica support for controllers.

Author: godwinpang

Diff for: pkg/config/controller.go

@@ -60,6 +60,10 @@ type Controller struct {
 	// Defaults to true, which means the controller will use leader election.
 	NeedLeaderElection *bool
 
+	// NeedWarmUp indicates whether the controller needs to use warm up.
+	// Defaults to false, which means the controller will not use warm up.
+	NeedWarmUp *bool
+
 	// UsePriorityQueue configures the controllers queue to use the controller-runtime provided
 	// priority queue.
 	//

Diff for: pkg/controller/controller.go

@@ -93,6 +93,12 @@ type TypedOptions[request comparable] struct {
 	//
 	// Note: This flag is disabled by default until a future version. It's currently in beta.
 	UsePriorityQueue *bool
+
+	// ShouldWarmupWithoutLeadership specifies whether the controller should start its sources
+	// when the manager is not the leader.
+	// Defaults to false, which means that the controller will wait for leader election to start
+	// before starting sources.
+	ShouldWarmupWithoutLeadership *bool
 }
 
 // DefaultFromConfig defaults the config from a config.Controller
@@ -244,15 +250,16 @@ func NewTypedUnmanaged[request comparable](name string, options TypedOptions[req
 
 	// Create controller with dependencies set
 	return &controller.Controller[request]{
-		Do:                      options.Reconciler,
-		RateLimiter:             options.RateLimiter,
-		NewQueue:                options.NewQueue,
-		MaxConcurrentReconciles: options.MaxConcurrentReconciles,
-		CacheSyncTimeout:        options.CacheSyncTimeout,
-		Name:                    name,
-		LogConstructor:          options.LogConstructor,
-		RecoverPanic:            options.RecoverPanic,
-		LeaderElected:           options.NeedLeaderElection,
+		Do:                            options.Reconciler,
+		RateLimiter:                   options.RateLimiter,
+		NewQueue:                      options.NewQueue,
+		MaxConcurrentReconciles:       options.MaxConcurrentReconciles,
+		CacheSyncTimeout:              options.CacheSyncTimeout,
+		Name:                          name,
+		LogConstructor:                options.LogConstructor,
+		RecoverPanic:                  options.RecoverPanic,
+		LeaderElected:                 options.NeedLeaderElection,
+		ShouldWarmupWithoutLeadership: options.ShouldWarmupWithoutLeadership,
 	}, nil
 }
 

Diff for: pkg/controller/controller_test.go

@@ -474,5 +474,78 @@ var _ = Describe("controller.Controller", func() {
 			_, ok = q.(priorityqueue.PriorityQueue[reconcile.Request])
 			Expect(ok).To(BeFalse())
 		})
+
+		It("should set ShouldWarmupWithoutLeadership correctly", func() {
+			m, err := manager.New(cfg, manager.Options{})
+			Expect(err).NotTo(HaveOccurred())
+
+			// Test with ShouldWarmupWithoutLeadership set to true
+			ctrlWithWarmup, err := controller.New("warmup-enabled-ctrl", m, controller.Options{
+				Reconciler:                    reconcile.Func(nil),
+				ShouldWarmupWithoutLeadership: ptr.To(true),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlWithWarmup, ok := ctrlWithWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlWithWarmup.ShouldWarmupWithoutLeadership).NotTo(BeNil())
+			Expect(*internalCtrlWithWarmup.ShouldWarmupWithoutLeadership).To(BeTrue())
+
+			// Test with ShouldWarmupWithoutLeadership set to false
+			ctrlWithoutWarmup, err := controller.New("warmup-disabled-ctrl", m, controller.Options{
+				Reconciler:                    reconcile.Func(nil),
+				ShouldWarmupWithoutLeadership: ptr.To(false),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlWithoutWarmup, ok := ctrlWithoutWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlWithoutWarmup.ShouldWarmupWithoutLeadership).NotTo(BeNil())
+			Expect(*internalCtrlWithoutWarmup.ShouldWarmupWithoutLeadership).To(BeFalse())
+
+			// Test with ShouldWarmupWithoutLeadership not set (should default to nil)
+			ctrlWithDefaultWarmup, err := controller.New("warmup-default-ctrl", m, controller.Options{
+				Reconciler: reconcile.Func(nil),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlWithDefaultWarmup, ok := ctrlWithDefaultWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlWithDefaultWarmup.ShouldWarmupWithoutLeadership).To(BeNil())
+		})
+
+		It("should inherit ShouldWarmupWithoutLeadership from manager config", func() {
+			// Test with manager default setting ShouldWarmupWithoutLeadership to true
+			managerWithWarmup, err := manager.New(cfg, manager.Options{
+				Controller: config.Controller{
+					NeedWarmUp: ptr.To(true),
+				},
+			})
+			Expect(err).NotTo(HaveOccurred())
+			ctrlInheritingWarmup, err := controller.New("inherit-warmup-enabled", managerWithWarmup, controller.Options{
+				Reconciler: reconcile.Func(nil),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlInheritingWarmup, ok := ctrlInheritingWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			// Note: This test will fail until the DefaultFromConfig method is updated to set
+			// ShouldWarmupWithoutLeadership from config.Controller.NeedWarmUp
+			// This test demonstrates that the feature needs to be completed
+			Expect(internalCtrlInheritingWarmup.ShouldWarmupWithoutLeadership).NotTo(BeNil())
+			Expect(*internalCtrlInheritingWarmup.ShouldWarmupWithoutLeadership).To(BeTrue())
+
+			// Test that explicit controller setting overrides manager setting
+			ctrlOverridingWarmup, err := controller.New("override-warmup-disabled", managerWithWarmup, controller.Options{
+				Reconciler:                    reconcile.Func(nil),
+				ShouldWarmupWithoutLeadership: ptr.To(false),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlOverridingWarmup, ok := ctrlOverridingWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlOverridingWarmup.ShouldWarmupWithoutLeadership).NotTo(BeNil())
+			Expect(*internalCtrlOverridingWarmup.ShouldWarmupWithoutLeadership).To(BeFalse())
+		})
 	})
 })

Diff for: pkg/internal/controller/controller.go

@@ -83,6 +83,9 @@ type Controller[request comparable] struct {
 	// startWatches maintains a list of sources, handlers, and predicates to start when the controller is started.
 	startWatches []source.TypedSource[request]
 
+	// didStartEventSources is used to indicate whether the event sources have been started.
+	didStartEventSources atomic.Bool
+
 	// LogConstructor is used to construct a logger to then log messages to users during reconciliation,
 	// or for example when a watch is started.
 	// Note: LogConstructor has to be able to handle nil requests as we are also using it
@@ -95,6 +98,12 @@ type Controller[request comparable] struct {
 
 	// LeaderElected indicates whether the controller is leader elected or always running.
 	LeaderElected *bool
+
+	// ShouldWarmupWithoutLeadership specifies whether the controller should start its sources
+	// when the manager is not the leader.
+	// Defaults to false, which means that the controller will wait for leader election to start
+	// before starting sources.
+	ShouldWarmupWithoutLeadership *bool
 }
 
 // Reconcile implements reconcile.Reconciler.
@@ -144,6 +153,15 @@ func (c *Controller[request]) NeedLeaderElection() bool {
 	return *c.LeaderElected
 }
 
+// Warmup implements the manager.WarmupRunnable interface.
+func (c *Controller[request]) Warmup(ctx context.Context) error {
+	if c.ShouldWarmupWithoutLeadership == nil || !*c.ShouldWarmupWithoutLeadership {
+		return nil
+	}
+
+	return c.startEventSources(ctx)
+}
+
 // Start implements controller.Controller.
 func (c *Controller[request]) Start(ctx context.Context) error {
 	// use an IIFE to get proper lock handling
@@ -221,6 +239,13 @@ func (c *Controller[request]) Start(ctx context.Context) error {
 // startEventSources launches all the sources registered with this controller and waits
 // for them to sync. It returns an error if any of the sources fail to start or sync.
 func (c *Controller[request]) startEventSources(ctx context.Context) error {
+	// CAS returns false if value is already true, so early exit since another goroutine must have
+	// called startEventSources previously
+	if !c.didStartEventSources.CompareAndSwap(false, true) {
+		c.LogConstructor(nil).Info("Skipping starting event sources since it was already started")
+		return nil
+	}
+
 	errGroup := &errgroup.Group{}
 	for _, watch := range c.startWatches {
 		log := c.LogConstructor(nil)

Diff for: pkg/internal/controller/controller_test.go

@@ -1014,6 +1014,127 @@ var _ = Describe("controller", func() {
 			})
 		})
 	})
+
+	Describe("Warmup", func() {
+		It("should start event sources when ShouldWarmupWithoutLeadership is true", func() {
+			// Setup
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// Create a mock source that we can verify was started
+			sourceStarted := false
+			mockSource := source.Func(func(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+				sourceStarted = true
+				return nil
+			})
+
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+			ctrl.ShouldWarmupWithoutLeadership = ptr.To(true)
+
+			// Act
+			err := ctrl.Warmup(ctx)
+
+			// Assert
+			Expect(err).NotTo(HaveOccurred())
+			Expect(sourceStarted).To(BeTrue(), "Event source should have been started")
+			Expect(ctrl.didStartEventSources.Load()).To(BeTrue(), "didStartEventSources flag should be set")
+		})
+
+		It("should not start event sources when ShouldWarmupWithoutLeadership is false", func() {
+			// Setup
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// Create a mock source that should not be started
+			sourceStarted := false
+			mockSource := source.Func(func(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+				sourceStarted = true
+				return nil
+			})
+
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+			ctrl.ShouldWarmupWithoutLeadership = ptr.To(false)
+
+			// Act
+			err := ctrl.Warmup(ctx)
+
+			// Assert
+			Expect(err).NotTo(HaveOccurred())
+			Expect(sourceStarted).To(BeFalse(), "Event source should not have been started")
+			Expect(ctrl.didStartEventSources.Load()).To(BeFalse(), "didStartEventSources flag should not be set")
+		})
+
+		It("should not start event sources when ShouldWarmupWithoutLeadership is nil", func() {
+			// Setup
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// Create a mock source that should not be started
+			sourceStarted := false
+			mockSource := source.Func(func(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+				sourceStarted = true
+				return nil
+			})
+
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+			ctrl.ShouldWarmupWithoutLeadership = nil
+
+			// Act
+			err := ctrl.Warmup(ctx)
+
+			// Assert
+			Expect(err).NotTo(HaveOccurred())
+			Expect(sourceStarted).To(BeFalse(), "Event source should not have been started")
+			Expect(ctrl.didStartEventSources.Load()).To(BeFalse(), "didStartEventSources flag should not be set")
+		})
+
+		It("should not start event sources twice when called multiple times", func() {
+			// Setup
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// Create a mock source that counts how many times it's started
+			startCount := 0
+			mockSource := source.Func(func(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+				startCount++
+				return nil
+			})
+
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+			ctrl.ShouldWarmupWithoutLeadership = ptr.To(true)
+
+			// Act
+			err1 := ctrl.Warmup(ctx)
+			err2 := ctrl.Warmup(ctx)
+
+			// Assert
+			Expect(err1).NotTo(HaveOccurred())
+			Expect(err2).NotTo(HaveOccurred())
+			Expect(startCount).To(Equal(1), "Event source should have been started only once")
+			Expect(ctrl.didStartEventSources.Load()).To(BeTrue(), "didStartEventSources flag should be set")
+		})
+
+		It("should propagate errors from event sources", func() {
+			// Setup
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			// Create a mock source that returns an error
+			expectedErr := errors.New("test error")
+			mockSource := source.Func(func(ctx context.Context, queue workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+				return expectedErr
+			})
+
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+			ctrl.ShouldWarmupWithoutLeadership = ptr.To(true)
+
+			// Act
+			err := ctrl.Warmup(ctx)
+
+			// Assert
+			Expect(err).To(MatchError(expectedErr))
+		})
+	})
 })
 
 var _ = Describe("ReconcileIDFromContext function", func() {

Diff for: pkg/manager/internal.go

@@ -439,6 +439,11 @@ func (cm *controllerManager) Start(ctx context.Context) (err error) {
 		return fmt.Errorf("failed to start other runnables: %w", err)
 	}
 
+	// Start and wait for sources to start.
+	if err := cm.runnables.Warmup.Start(cm.internalCtx); err != nil {
+		return fmt.Errorf("failed to start warmup runnables: %w", err)
+	}
+
 	// Start the leader election and all required runnables.
 	{
 		ctx, cancel := context.WithCancel(context.Background())
@@ -544,6 +549,10 @@ func (cm *controllerManager) engageStopProcedure(stopComplete <-chan struct{}) e
 		cm.runnables.LeaderElection.startOnce.Do(func() {})
 		cm.runnables.LeaderElection.StopAndWait(cm.shutdownCtx)
 
+		// Stop the warmup runnables
+		cm.logger.Info("Stopping and waiting for warmup runnables")
+		cm.runnables.Warmup.StopAndWait(cm.shutdownCtx)
+
 		// Stop the caches before the leader election runnables, this is an important
 		// step to make sure that we don't race with the reconcilers by receiving more events
 		// from the API servers and enqueueing them.

Diff for: pkg/manager/manager.go

@@ -314,6 +314,12 @@ type LeaderElectionRunnable interface {
 	NeedLeaderElection() bool
 }
 
+// WarmupRunnable knows if a Runnable should be a warmup runnable.
+type WarmupRunnable interface {
+	// Warmup returns true if the Runnable should be run as warmup.
+	Warmup(context.Context) error
+}
+
 // New returns a new Manager for creating Controllers.
 // Note that if ContentType in the given config is not set, "application/vnd.kubernetes.protobuf"
 // will be used for all built-in resources of Kubernetes, and "application/json" is for other types