Add ready check to block controller startup until warmup is complete.

Author: godwinpang

pkg/internal/controller/controller.go

@@ -29,7 +29,9 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/apimachinery/pkg/util/uuid"
+	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/util/workqueue"
+	"k8s.io/utils/ptr"
 
 	"sigs.k8s.io/controller-runtime/pkg/controller/priorityqueue"
 	ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/internal/controller/metrics"
@@ -38,6 +40,11 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/source"
 )
 
+const (
+	// syncedPollPeriod is the period to poll for cache sync
+	syncedPollPeriod = 100 * time.Millisecond
+)
+
 // Controller implements controller.Controller.
 type Controller[request comparable] struct {
 	// Name is used to uniquely identify a Controller in tracing, logging and monitoring.  Name is required.
@@ -86,6 +93,13 @@ type Controller[request comparable] struct {
 	// didStartEventSources is used to indicate whether the event sources have been started.
 	didStartEventSources atomic.Bool
 
+	// didEventSourcesFinishSync is used to indicate whether the event sources have finished
+	// successfully. It stores a *bool where
+	// - nil: not finished syncing
+	// - true: finished syncing without error
+	// - false: finished syncing with error
+	didEventSourcesFinishSync atomic.Value
+
 	// 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
@@ -156,12 +170,37 @@ func (c *Controller[request]) NeedLeaderElection() bool {
 // Warmup implements the manager.WarmupRunnable interface.
 func (c *Controller[request]) Warmup(ctx context.Context) error {
 	if c.NeedWarmup == nil || !*c.NeedWarmup {
+		c.didEventSourcesFinishSync.Store(ptr.To(true))
 		return nil
 	}
 
 	return c.startEventSources(ctx)
 }
 
+// DidFinishWarmup implements the manager.WarmupRunnable interface.
+func (c *Controller[request]) DidFinishWarmup(ctx context.Context) bool {
+	err := wait.PollUntilContextCancel(ctx, syncedPollPeriod, true, func(ctx context.Context) (bool, error) {
+		didFinishSync, ok := c.didEventSourcesFinishSync.Load().(*bool)
+		if !ok {
+			return false, errors.New("unexpected error: didEventSourcesFinishSync is not a bool pointer")
+		}
+
+		if didFinishSync == nil {
+			// event sources not finished syncing
+			return false, nil
+		}
+
+		if !*didFinishSync {
+			// event sources finished syncing with an error
+			return true, errors.New("event sources did not finish syncing succesfully")
+		}
+
+		return true, nil
+	})
+
+	return err == nil
+}
+
 // Start implements controller.Controller.
 func (c *Controller[request]) Start(ctx context.Context) error {
 	// use an IIFE to get proper lock handling
@@ -299,7 +338,11 @@ func (c *Controller[request]) startEventSources(ctx context.Context) error {
 			}
 		})
 	}
-	return errGroup.Wait()
+	err := errGroup.Wait()
+
+	c.didEventSourcesFinishSync.Store(ptr.To(err == nil))
+
+	return err
 }
 
 // processNextWorkItem will read a single work item off the workqueue and

pkg/internal/controller/controller_test.go

@@ -1016,126 +1016,78 @@ var _ = Describe("controller", func() {
 	})
 
 	Describe("Warmup", func() {
-		It("should start event sources when NeedWarmup 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.CacheSyncTimeout = time.Second
-			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
+		JustBeforeEach(func() {
 			ctrl.NeedWarmup = 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 NeedWarmup is false", func() {
-			// Setup
+		It("should track warmup status correctly with successful sync", func() {
+			// Setup controller with sources that complete successfully
 			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.NeedWarmup = ptr.To(false)
+			ctrl.CacheSyncTimeout = time.Second
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{
+				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+					return 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")
+
+			// Verify DidFinishWarmup returns true for successful sync
+			result := ctrl.DidFinishWarmup(ctx)
+			Expect(result).To(BeTrue())
 		})
 
-		It("should not start event sources when NeedWarmup is nil", func() {
-			// Setup
+		It("should track warmup status correctly with unsuccessful sync", func() {
+			// Setup controller with sources that complete with error
 			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.NeedWarmup = nil
+			ctrl.CacheSyncTimeout = time.Second
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{
+				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+					return errors.New("sync error")
+				}),
+			}
 
-			// Act
 			err := ctrl.Warmup(ctx)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("sync error"))
 
-			// 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")
+			// Verify DidFinishWarmup returns false for unsuccessful sync
+			result := ctrl.DidFinishWarmup(ctx)
+			Expect(result).To(BeFalse())
 		})
+	})
 
-		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.CacheSyncTimeout = time.Second
-			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
-			ctrl.NeedWarmup = 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")
+	Describe("Warmup without warmup enabled", func() {
+		JustBeforeEach(func() {
+			ctrl.NeedWarmup = ptr.To(false)
 		})
 
-		It("should propagate errors from event sources", func() {
-			// Setup
+		It("should not start sources if warmup is disabled.", func() {
+			// Setup controller with sources that complete successfully
 			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.CacheSyncTimeout = time.Second
-			ctrl.startWatches = []source.TypedSource[reconcile.Request]{mockSource}
-			ctrl.NeedWarmup = ptr.To(true)
+			isSourceStarted := false
+			ctrl.startWatches = []source.TypedSource[reconcile.Request]{
+				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
+					isSourceStarted = true
+					return nil
+				}),
+			}
 
-			// Act
 			err := ctrl.Warmup(ctx)
+			Expect(err).NotTo(HaveOccurred())
+
+			result := ctrl.DidFinishWarmup(ctx)
+			Expect(result).To(BeTrue())
 
-			// Assert
-			Expect(err).To(MatchError(expectedErr))
+			Expect(isSourceStarted).To(BeFalse())
 		})
 	})
 })

pkg/manager/manager.go

@@ -316,10 +316,13 @@ type LeaderElectionRunnable interface {
 
 // WarmupRunnable knows if a Runnable should be a warmup runnable. A warmup runnable is a runnable
 // that should be run when the manager is started, but before the leader election is acquired.
-// The expectation is that it will block startup until the warmup runnables have completed.
 type WarmupRunnable interface {
-	// Warmup is called when the manager is started, but before the leader election is acquired.
+	// Warmup is the implementation of the warmup runnable.
 	Warmup(context.Context) error
+
+	// WaitForWarmupComplete is a blocking function that waits for the warmup to be completed. It
+	// returns false if it could not successfully finish warmup.
+	WaitForWarmupComplete(context.Context) bool
 }
 
 // New returns a new Manager for creating Controllers.

pkg/manager/runnable_group.go

@@ -69,7 +69,12 @@ func (r *runnables) Add(fn Runnable) error {
 		return r.Webhooks.Add(fn, nil)
 	case WarmupRunnable, LeaderElectionRunnable:
 		if warmupRunnable, ok := fn.(WarmupRunnable); ok {
-			if err := r.Warmup.Add(RunnableFunc(warmupRunnable.Warmup), nil); err != nil {
+			if err := r.Warmup.Add(
+				RunnableFunc(warmupRunnable.Warmup),
+				func(ctx context.Context) bool {
+					return warmupRunnable.WaitForWarmupComplete(ctx)
+				},
+			); err != nil {
 				return err
 			}
 		}