Fail in warmup directly and rely on sync.Once for warmup thread-safety without WaitForWarmupComplete.

Author: godwinpang

pkg/internal/controller/controller.go

@@ -83,9 +83,6 @@ type ControllerOptions[request comparable] struct {
 }
 
 // Controller implements controller.Controller.
-// WARNING: If directly instantiating a Controller vs. using the New method, ensure that the
-// warmupResultChan is instantiated as a buffered channel of size 1. Otherwise, the controller will
-// panic on having Warmup called.
 type Controller[request comparable] struct {
 	// Name is used to uniquely identify a Controller in tracing, logging and monitoring.  Name is required.
 	Name string
@@ -133,10 +130,6 @@ type Controller[request comparable] struct {
 	// didStartEventSourcesOnce is used to ensure that the event sources are only started once.
 	didStartEventSourcesOnce sync.Once
 
-	// warmupResultChan receives the result (nil / non-nil error) of the warmup method. It is
-	// consumed by the WaitForWarmupComplete method that the warmup has finished.
-	warmupResultChan chan error
-
 	// 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
@@ -169,7 +162,6 @@ func New[request comparable](options ControllerOptions[request]) *Controller[req
 		RecoverPanic:            options.RecoverPanic,
 		LeaderElected:           options.LeaderElected,
 		EnableWarmup:            options.EnableWarmup,
-		warmupResultChan:        make(chan error, 1),
 	}
 }
 
@@ -226,31 +218,7 @@ func (c *Controller[request]) Warmup(ctx context.Context) error {
 		return nil
 	}
 
-	// Hold the lock to avoid concurrent access to c.startWatches with Start() when calling
-	// startEventSources
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	err := c.startEventSources(ctx)
-	c.warmupResultChan <- err
-
-	return err
-}
-
-// WaitForWarmupComplete returns true if warmup has completed without error, and false if there was
-// an error during warmup. If context is cancelled, it returns true.
-func (c *Controller[request]) WaitForWarmupComplete(ctx context.Context) bool {
-	if c.EnableWarmup == nil || !*c.EnableWarmup {
-		return true
-	}
-
-	warmupError, ok := <-c.warmupResultChan
-	if !ok {
-		// channel closed unexpectedly
-		return false
-	}
-
-	return warmupError == nil
+	return c.startEventSources(ctx)
 }
 
 // Start implements controller.Controller.

pkg/internal/controller/controller_test.go

@@ -1121,12 +1121,7 @@ var _ = Describe("controller", func() {
 				}),
 			}
 
-			err := ctrl.Warmup(ctx)
-			Expect(err).NotTo(HaveOccurred())
-
-			// Verify WaitForWarmupComplete returns true for successful sync
-			result := ctrl.WaitForWarmupComplete(ctx)
-			Expect(result).To(BeTrue())
+			Expect(ctrl.Warmup(ctx)).To(Succeed())
 		})
 
 		It("should track warmup status correctly with unsuccessful sync", func() {
@@ -1144,10 +1139,6 @@ var _ = Describe("controller", func() {
 			err := ctrl.Warmup(ctx)
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("sync error"))
-
-			// Verify WaitForWarmupComplete returns false for unsuccessful sync
-			result := ctrl.WaitForWarmupComplete(ctx)
-			Expect(result).To(BeFalse())
 		})
 
 		It("should return true if context is cancelled while waiting for source to start", func() {
@@ -1163,27 +1154,19 @@ var _ = Describe("controller", func() {
 				}),
 			}
 
-			resultChan := make(chan bool)
-
 			// Wait for the goroutines to finish before returning to avoid racing with the
 			// assignment in BeforeEach block.
 			var wg sync.WaitGroup
 
-			// Invoked in goroutines because the Warmup / WaitForWarmupComplete will block forever.
-			wg.Add(2)
+			// Invoked in a goroutine because Warmup will block
+			wg.Add(1)
 			go func() {
 				defer GinkgoRecover()
 				defer wg.Done()
 				Expect(ctrl.Warmup(ctx)).To(Succeed())
 			}()
-			go func() {
-				defer GinkgoRecover()
-				defer wg.Done()
-				resultChan <- ctrl.WaitForWarmupComplete(ctx)
-			}()
 
 			cancel()
-			Expect(<-resultChan).To(BeTrue())
 			wg.Wait()
 		})
 
@@ -1196,16 +1179,15 @@ var _ = Describe("controller", func() {
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 
-			hasCtrlWatchStarted, hasNonWarmupCtrlWatchStarted := atomic.Bool{}, atomic.Bool{}
-
-			// ctrl watch will block from finishing until the channel is produced to
-			ctrlWatchBlockingChan := make(chan struct{})
+			By("Creating a channel to track execution order")
+			runnableExecutionOrderChan := make(chan string, 2)
+			const nonWarmupRunnableName = "nonWarmupRunnable"
+			const warmupRunnableName = "warmupRunnable"
 
 			ctrl.CacheSyncTimeout = time.Second
 			ctrl.startWatches = []source.TypedSource[reconcile.Request]{
 				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
-					hasCtrlWatchStarted.Store(true)
-					<-ctrlWatchBlockingChan
+					runnableExecutionOrderChan <- warmupRunnableName
 					return nil
 				}),
 			}
@@ -1225,7 +1207,7 @@ var _ = Describe("controller", func() {
 			})
 			nonWarmupCtrl.startWatches = []source.TypedSource[reconcile.Request]{
 				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
-					hasNonWarmupCtrlWatchStarted.Store(true)
+					runnableExecutionOrderChan <- nonWarmupRunnableName
 					return nil
 				}),
 			}
@@ -1249,13 +1231,9 @@ var _ = Describe("controller", func() {
 				Expect(m.Start(ctx)).To(Succeed())
 			}()
 
-			By("Waiting for the warmup controller to start")
-			Eventually(hasCtrlWatchStarted.Load).Should(BeTrue())
-			Expect(hasNonWarmupCtrlWatchStarted.Load()).To(BeFalse())
-
-			By("Unblocking the warmup controller source start")
-			close(ctrlWatchBlockingChan)
-			Eventually(hasNonWarmupCtrlWatchStarted.Load).Should(BeTrue())
+			<-m.Elected()
+			Expect(<-runnableExecutionOrderChan).To(Equal(warmupRunnableName))
+			Expect(<-runnableExecutionOrderChan).To(Equal(nonWarmupRunnableName))
 		})
 
 		It("should not race with Start and only start sources once", func() {
@@ -1326,58 +1304,6 @@ var _ = Describe("controller", func() {
 			Eventually(isSourceStarted.Load).Should(BeTrue())
 		})
 	})
-
-	Describe("WaitForWarmupComplete", func() {
-		It("should short circuit without blocking if warmup is disabled", func() {
-			ctrl.EnableWarmup = ptr.To(false)
-
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-
-			// Call WaitForWarmupComplete and expect it to return immediately
-			result := ctrl.WaitForWarmupComplete(ctx)
-			Expect(result).To(BeTrue())
-		})
-
-		It("should block until warmup is complete if warmup is enabled", func() {
-			ctrl.EnableWarmup = ptr.To(true)
-			// Setup controller with sources that complete successfully
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-
-			// Close the channel to signal watch completion
-			shouldWatchCompleteChan := make(chan struct{})
-
-			ctrl.CacheSyncTimeout = time.Second
-			ctrl.startWatches = []source.TypedSource[reconcile.Request]{
-				source.Func(func(ctx context.Context, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) error {
-					<-shouldWatchCompleteChan
-					return nil
-				}),
-			}
-
-			By("Starting a blocking warmup")
-			go func() {
-				defer GinkgoRecover()
-				Expect(ctrl.Warmup(ctx)).To(Succeed())
-			}()
-
-			// didWaitForWarmupCompleteReturn is true when the call to WaitForWarmupComplete returns
-			didWaitForWarmupCompleteReturn := atomic.Bool{}
-			go func() {
-				defer GinkgoRecover()
-				// Verify WaitForWarmupComplete returns true for successful sync
-				Expect(ctrl.WaitForWarmupComplete(ctx)).To(BeTrue())
-				didWaitForWarmupCompleteReturn.Store(true)
-			}()
-			Consistently(didWaitForWarmupCompleteReturn.Load).Should(BeFalse())
-
-			By("Unblocking the watch to simulate initial sync completion")
-			close(shouldWatchCompleteChan)
-			Eventually(didWaitForWarmupCompleteReturn.Load).Should(BeTrue())
-		})
-
-	})
 })
 
 var _ = Describe("ReconcileIDFromContext function", func() {

pkg/manager/manager.go

@@ -321,10 +321,6 @@ type LeaderElectionRunnable interface {
 type warmupRunnable interface {
 	// Warmup will be called when the manager is started but before it becomes leader.
 	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/manager_test.go

@@ -1930,11 +1930,10 @@ var _ = Describe("manger.Manager", func() {
 	})
 
 	It("should run warmup runnables before leader election is won", func() {
-		By("Creating channels to track execution order")
-		warmupCalled := atomic.Bool{}
-		leaderElectionRunnableStarted := atomic.Bool{}
-		// warmupRunnable's WarmupFunc will block until this channel is closed
-		warmupRunnableWarmupBlockingChan := make(chan struct{})
+		By("Creating a channel to track execution order")
+		runnableExecutionOrderChan := make(chan string, 2)
+		const leaderElectionRunnableName = "leaderElectionRunnable"
+		const warmupRunnableName = "warmupRunnable"
 
 		By("Creating a manager with leader election enabled")
 		m, err := New(cfg, Options{
@@ -1952,25 +1951,23 @@ var _ = Describe("manger.Manager", func() {
 		// Create a warmup runnable
 		warmupRunnable := newWarmupRunnableFunc(
 			func(ctx context.Context) error {
-				// This is the main runnable that will be executed after leader election
-				// Block forever
+				// This is the leader election runnable that will be executed after leader election
+				// It will block until context is done/cancelled
 				<-ctx.Done()
 				return nil
 			},
 			func(ctx context.Context) error {
 				// This should be called during startup before leader election
-				warmupCalled.Store(true)
-				<-warmupRunnableWarmupBlockingChan
+				runnableExecutionOrderChan <- warmupRunnableName
 				return nil
 			},
 		)
 		Expect(m.Add(warmupRunnable)).To(Succeed())
 
 		By("Creating a runnable that requires leader election")
-
 		leaderElectionRunnable := RunnableFunc(
 			func(ctx context.Context) error {
-				leaderElectionRunnableStarted.Store(true)
+				runnableExecutionOrderChan <- leaderElectionRunnableName
 				<-ctx.Done()
 				return nil
 			},
@@ -1985,18 +1982,9 @@ var _ = Describe("manger.Manager", func() {
 			Expect(m.Start(ctx)).To(Succeed())
 		}()
 
-		By("Waiting for the warmup runnable to be called")
-		Eventually(warmupCalled.Load).Should(BeTrue())
-		Expect(leaderElectionRunnableStarted.Load()).To(BeFalse())
-
-		By("Closing the channel to unblock the warmup runnable")
-		close(warmupRunnableWarmupBlockingChan)
-
-		By("Waiting for leader election to be won")
 		<-m.Elected()
-
-		By("Verifying the leader election runnable is called after election")
-		Eventually(leaderElectionRunnableStarted.Load).Should(BeTrue())
+		Expect(<-runnableExecutionOrderChan).To(Equal(warmupRunnableName))
+		Expect(<-runnableExecutionOrderChan).To(Equal(leaderElectionRunnableName))
 	})
 })
 

pkg/manager/runnable_group.go

@@ -69,12 +69,7 @@ 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),
-				func(ctx context.Context) bool {
-					return warmupRunnable.WaitForWarmupComplete(ctx)
-				},
-			); err != nil {
+			if err := r.Warmup.Add(RunnableFunc(warmupRunnable.Warmup), nil); err != nil {
 				return err
 			}
 		}

pkg/manager/runnable_group_test.go

@@ -350,9 +350,8 @@ func newWarmupRunnableFunc(
 	warmupFunc func(context.Context) error,
 ) *warmupRunnableFunc {
 	return &warmupRunnableFunc{
-		startFunc:           startFunc,
-		warmupFunc:          warmupFunc,
-		didWarmupFinishChan: make(chan struct{}),
+		startFunc:  startFunc,
+		warmupFunc: warmupFunc,
 	}
 }
 
@@ -361,28 +360,14 @@ func newWarmupRunnableFunc(
 type warmupRunnableFunc struct {
 	startFunc  func(context.Context) error
 	warmupFunc func(context.Context) error
-
-	// didWarmupFinishChan is closed when warmup is finished
-	didWarmupFinishChan chan struct{}
-
-	// didWarmupFinishSuccessfully is set to true if warmup was successful
-	didWarmupFinishSuccessfully atomic.Bool
 }
 
 func (r *warmupRunnableFunc) Start(ctx context.Context) error {
 	return r.startFunc(ctx)
 }
 
 func (r *warmupRunnableFunc) Warmup(ctx context.Context) error {
-	err := r.warmupFunc(ctx)
-	r.didWarmupFinishSuccessfully.Store(err == nil)
-	close(r.didWarmupFinishChan)
-	return err
-}
-
-func (r *warmupRunnableFunc) WaitForWarmupComplete(ctx context.Context) bool {
-	<-r.didWarmupFinishChan
-	return r.didWarmupFinishSuccessfully.Load()
+	return r.warmupFunc(ctx)
 }
 
 var _ LeaderElectionRunnable = &leaderElectionAndWarmupRunnable{}
@@ -401,9 +386,8 @@ func newLeaderElectionAndWarmupRunnable(
 ) *leaderElectionAndWarmupRunnable {
 	return &leaderElectionAndWarmupRunnable{
 		warmupRunnableFunc: &warmupRunnableFunc{
-			startFunc:           startFunc,
-			warmupFunc:          warmupFunc,
-			didWarmupFinishChan: make(chan struct{}),
+			startFunc:  startFunc,
+			warmupFunc: warmupFunc,
 		},
 		needLeaderElection: needLeaderElection,
 	}