✨ Implement warm replica support for controllers (#3192)

* [Warm Replicas] Implement warm replica support for controllers.

* Remove irrelevant runnable_group.go code.

* Rename ShouldWarmup.

* fmt

* Change to atomic.Bool to avoid race in test.

* Address comments.

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

* Keep test helper structs private.

* Address comments.

* Fix lint.

* Address naming + comments from sbueringer.

* Refactor tests to use HaveValue.

* Document + add UT for WaitForWarmupComplete behavior on ctx cancellation.

* Add unit test that exercises controller warmup integration with manager.

* Add UT that verifies WaitForWarmupComplete blocking / non-blocking behavior.

* Verify r.Others.startQueue in runnables test cases.

* Fix UT to verify runnable ordering.

* Fix UT for WaitForWarmupComplete blocking.

* Document !NeedLeaderElection+NeedWarmup behavior

* Fix test race.

* Cleanup test wrapper runnables.

* Make didStartEventSources run once with sync.Once + UT.

* Rewrite Warmup to avoid polling.

* Rename NeedWarmup to EnableWarmup.

* Clarify comment on Warmup.

* Move reset watches critical section inside of startEventSources.

* Add test to assert startEventSources blocking behavior.

* Make Start threadsafe with Warmup + UT.

* Change warmup to use buffered error channel and add New method.

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

* Sync controller EnableWarmup comments.

* Rename to startEventSourcesLocked and lock with c.mu

* Address edge case for watch added after warmup completes.

* Fix test description and set leaderelection==true

* Fix lint.

* Change shutdown order to shutdown warmup runnables in parallel with other runnables.

* Fix test races by ensuring goroutines do not outlive their It blocks.

* Block on source start on context cancel.

* Guard access to c.Queue explicitly.

* Initialize queue in warmup with test.

* Fix watch comment.

* Add warmup to manager and controller integration tests.

* fmt + lint.

* Add tests for Warmup for parity with Start.

* golangci-lint

* Increase test coverage for start with warmup race.

* Remove synchronization code around c.Queue concurrent access.

* Assert startedEventSourcesAndQueue is set in tests.

* fmt

* Clarify namespace predicate vs. cleanup in controller integration test.

* Add option for fake leader elector package and assert that warmup runnables are started before leader election.

* fmt + rename

* Address PR comments.

Author: godwinpang

pkg/config/controller.go

@@ -60,6 +60,19 @@ type Controller struct {
 	// Defaults to true, which means the controller will use leader election.
 	NeedLeaderElection *bool
 
+	// EnableWarmup specifies whether the controller should start its sources when the manager is not
+	// the leader. This is useful for cases where sources take a long time to start, as it allows
+	// for the controller to warm up its caches even before it is elected as the leader. This
+	// improves leadership failover time, as the caches will be prepopulated before the controller
+	// transitions to be leader.
+	//
+	// Setting EnableWarmup to true and NeedLeaderElection to true means the controller will start its
+	// sources without waiting to become leader.
+	// Setting EnableWarmup to true and NeedLeaderElection to false is a no-op as controllers without
+	// leader election do not wait on leader election to start their sources.
+	// Defaults to false.
+	EnableWarmup *bool
+
 	// UsePriorityQueue configures the controllers queue to use the controller-runtime provided
 	// priority queue.
 	//

pkg/controller/controller.go

@@ -93,6 +93,19 @@ type TypedOptions[request comparable] struct {
 	//
 	// Note: This flag is disabled by default until a future version. It's currently in beta.
 	UsePriorityQueue *bool
+
+	// EnableWarmup specifies whether the controller should start its sources when the manager is not
+	// the leader. This is useful for cases where sources take a long time to start, as it allows
+	// for the controller to warm up its caches even before it is elected as the leader. This
+	// improves leadership failover time, as the caches will be prepopulated before the controller
+	// transitions to be leader.
+	//
+	// Setting EnableWarmup to true and NeedLeaderElection to true means the controller will start its
+	// sources without waiting to become leader.
+	// Setting EnableWarmup to true and NeedLeaderElection to false is a no-op as controllers without
+	// leader election do not wait on leader election to start their sources.
+	// Defaults to false.
+	EnableWarmup *bool
 }
 
 // DefaultFromConfig defaults the config from a config.Controller
@@ -124,6 +137,10 @@ func (options *TypedOptions[request]) DefaultFromConfig(config config.Controller
 	if options.NeedLeaderElection == nil {
 		options.NeedLeaderElection = config.NeedLeaderElection
 	}
+
+	if options.EnableWarmup == nil {
+		options.EnableWarmup = config.EnableWarmup
+	}
 }
 
 // Controller implements an API. A Controller manages a work queue fed reconcile.Requests
@@ -243,7 +260,7 @@ func NewTypedUnmanaged[request comparable](name string, options TypedOptions[req
 	}
 
 	// Create controller with dependencies set
-	return &controller.Controller[request]{
+	return controller.New[request](controller.Options[request]{
 		Do:                      options.Reconciler,
 		RateLimiter:             options.RateLimiter,
 		NewQueue:                options.NewQueue,
@@ -253,7 +270,8 @@ func NewTypedUnmanaged[request comparable](name string, options TypedOptions[req
 		LogConstructor:          options.LogConstructor,
 		RecoverPanic:            options.RecoverPanic,
 		LeaderElected:           options.NeedLeaderElection,
-	}, nil
+		EnableWarmup:            options.EnableWarmup,
+	}), nil
 }
 
 // ReconcileIDFromContext gets the reconcileID from the current context.

pkg/controller/controller_integration_test.go

@@ -18,16 +18,21 @@ package controller_test
 
 import (
 	"context"
+	"fmt"
+	"strconv"
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/cache"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllertest"
 	"sigs.k8s.io/controller-runtime/pkg/handler"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	"sigs.k8s.io/controller-runtime/pkg/source"
 
@@ -48,30 +53,48 @@ var _ = Describe("controller", func() {
 	Describe("controller", func() {
 		// TODO(directxman12): write a whole suite of controller-client interaction tests
 
-		It("should reconcile", func() {
+		// The watches in this test are setup with a namespace predicate to avoid each table entry
+		// from interfering with the others. We cannot add a delete call for the pods created in the
+		// test, as it causes flakes with the api-server termination timing out.
+		// See https://github.com/kubernetes-sigs/controller-runtime/issues/1571 for a description
+		// of the issue, and a discussion here: https://github.com/kubernetes-sigs/controller-runtime/pull/3192#discussion_r2186967799
+		DescribeTable("should reconcile", func(enableWarmup bool) {
 			By("Creating the Manager")
 			cm, err := manager.New(cfg, manager.Options{})
 			Expect(err).NotTo(HaveOccurred())
 
 			By("Creating the Controller")
-			instance, err := controller.New("foo-controller", cm, controller.Options{
-				Reconciler: reconcile.Func(
-					func(_ context.Context, request reconcile.Request) (reconcile.Result, error) {
-						reconciled <- request
-						return reconcile.Result{}, nil
-					}),
-			})
+			instance, err := controller.New(
+				fmt.Sprintf("foo-controller-%t", enableWarmup),
+				cm,
+				controller.Options{
+					Reconciler: reconcile.Func(
+						func(_ context.Context, request reconcile.Request) (reconcile.Result, error) {
+							reconciled <- request
+							return reconcile.Result{}, nil
+						}),
+					EnableWarmup: ptr.To(enableWarmup),
+				},
+			)
 			Expect(err).NotTo(HaveOccurred())
 
+			testNamespace := strconv.FormatBool(enableWarmup)
+
 			By("Watching Resources")
 			err = instance.Watch(
 				source.Kind(cm.GetCache(), &appsv1.ReplicaSet{},
 					handler.TypedEnqueueRequestForOwner[*appsv1.ReplicaSet](cm.GetScheme(), cm.GetRESTMapper(), &appsv1.Deployment{}),
+					makeNamespacePredicate[*appsv1.ReplicaSet](testNamespace),
 				),
 			)
 			Expect(err).NotTo(HaveOccurred())
 
-			err = instance.Watch(source.Kind(cm.GetCache(), &appsv1.Deployment{}, &handler.TypedEnqueueRequestForObject[*appsv1.Deployment]{}))
+			err = instance.Watch(
+				source.Kind(cm.GetCache(), &appsv1.Deployment{},
+					&handler.TypedEnqueueRequestForObject[*appsv1.Deployment]{},
+					makeNamespacePredicate[*appsv1.Deployment](testNamespace),
+				),
+			)
 			Expect(err).NotTo(HaveOccurred())
 
 			err = cm.GetClient().Get(ctx, types.NamespacedName{Name: "foo"}, &corev1.Namespace{})
@@ -110,19 +133,25 @@ var _ = Describe("controller", func() {
 				},
 			}
 			expectedReconcileRequest := reconcile.Request{NamespacedName: types.NamespacedName{
-				Namespace: "default",
+				Namespace: testNamespace,
 				Name:      "deployment-name",
 			}}
 
+			By("Creating the test namespace")
+			_, err = clientset.CoreV1().Namespaces().Create(ctx, &corev1.Namespace{
+				ObjectMeta: metav1.ObjectMeta{Name: testNamespace},
+			}, metav1.CreateOptions{})
+			Expect(err).NotTo(HaveOccurred())
+
 			By("Invoking Reconciling for Create")
-			deployment, err = clientset.AppsV1().Deployments("default").Create(ctx, deployment, metav1.CreateOptions{})
+			deployment, err = clientset.AppsV1().Deployments(testNamespace).Create(ctx, deployment, metav1.CreateOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
 
 			By("Invoking Reconciling for Update")
 			newDeployment := deployment.DeepCopy()
 			newDeployment.Labels = map[string]string{"foo": "bar"}
-			_, err = clientset.AppsV1().Deployments("default").Update(ctx, newDeployment, metav1.UpdateOptions{})
+			_, err = clientset.AppsV1().Deployments(testNamespace).Update(ctx, newDeployment, metav1.UpdateOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
 
@@ -145,24 +174,24 @@ var _ = Describe("controller", func() {
 					Template: deployment.Spec.Template,
 				},
 			}
-			replicaset, err = clientset.AppsV1().ReplicaSets("default").Create(ctx, replicaset, metav1.CreateOptions{})
+			replicaset, err = clientset.AppsV1().ReplicaSets(testNamespace).Create(ctx, replicaset, metav1.CreateOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
 
 			By("Invoking Reconciling for an OwnedObject when it is updated")
 			newReplicaset := replicaset.DeepCopy()
 			newReplicaset.Labels = map[string]string{"foo": "bar"}
-			_, err = clientset.AppsV1().ReplicaSets("default").Update(ctx, newReplicaset, metav1.UpdateOptions{})
+			_, err = clientset.AppsV1().ReplicaSets(testNamespace).Update(ctx, newReplicaset, metav1.UpdateOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
 
 			By("Invoking Reconciling for an OwnedObject when it is deleted")
-			err = clientset.AppsV1().ReplicaSets("default").Delete(ctx, replicaset.Name, metav1.DeleteOptions{})
+			err = clientset.AppsV1().ReplicaSets(testNamespace).Delete(ctx, replicaset.Name, metav1.DeleteOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
 
 			By("Invoking Reconciling for Delete")
-			err = clientset.AppsV1().Deployments("default").
+			err = clientset.AppsV1().Deployments(testNamespace).
 				Delete(ctx, "deployment-name", metav1.DeleteOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
@@ -174,7 +203,12 @@ var _ = Describe("controller", func() {
 
 			By("Invoking Reconciling for a pod when it is created when adding watcher dynamically")
 			// Add new watcher dynamically
-			err = instance.Watch(source.Kind(cm.GetCache(), &corev1.Pod{}, &handler.TypedEnqueueRequestForObject[*corev1.Pod]{}))
+			err = instance.Watch(
+				source.Kind(cm.GetCache(), &corev1.Pod{},
+					&handler.TypedEnqueueRequestForObject[*corev1.Pod]{},
+					makeNamespacePredicate[*corev1.Pod](testNamespace),
+				),
+			)
 			Expect(err).NotTo(HaveOccurred())
 
 			pod := &corev1.Pod{
@@ -194,16 +228,27 @@ var _ = Describe("controller", func() {
 				},
 			}
 			expectedReconcileRequest = reconcile.Request{NamespacedName: types.NamespacedName{
-				Namespace: "default",
+				Namespace: testNamespace,
 				Name:      "pod-name",
 			}}
-			_, err = clientset.CoreV1().Pods("default").Create(ctx, pod, metav1.CreateOptions{})
+			_, err = clientset.CoreV1().Pods(testNamespace).Create(ctx, pod, metav1.CreateOptions{})
 			Expect(err).NotTo(HaveOccurred())
 			Expect(<-reconciled).To(Equal(expectedReconcileRequest))
-		})
+		},
+			Entry("with controller warmup enabled", true),
+			Entry("with controller warmup not enabled", false),
+		)
 	})
 })
 
+// makeNamespacePredicate returns a predicate that filters out all objects not in the passed in
+// namespace.
+func makeNamespacePredicate[object client.Object](namespace string) predicate.TypedPredicate[object] {
+	return predicate.NewTypedPredicateFuncs[object](func(obj object) bool {
+		return obj.GetNamespace() == namespace
+	})
+}
+
 func truePtr() *bool {
 	t := true
 	return &t

pkg/controller/controller_test.go

@@ -474,5 +474,71 @@ var _ = Describe("controller.Controller", func() {
 			_, ok = q.(priorityqueue.PriorityQueue[reconcile.Request])
 			Expect(ok).To(BeFalse())
 		})
+
+		It("should set EnableWarmup correctly", func() {
+			m, err := manager.New(cfg, manager.Options{})
+			Expect(err).NotTo(HaveOccurred())
+
+			// Test with EnableWarmup set to true
+			ctrlWithWarmup, err := controller.New("warmup-enabled-ctrl", m, controller.Options{
+				Reconciler:   reconcile.Func(nil),
+				EnableWarmup: ptr.To(true),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlWithWarmup, ok := ctrlWithWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlWithWarmup.EnableWarmup).To(HaveValue(BeTrue()))
+
+			// Test with EnableWarmup set to false
+			ctrlWithoutWarmup, err := controller.New("warmup-disabled-ctrl", m, controller.Options{
+				Reconciler:   reconcile.Func(nil),
+				EnableWarmup: ptr.To(false),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlWithoutWarmup, ok := ctrlWithoutWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlWithoutWarmup.EnableWarmup).To(HaveValue(BeFalse()))
+
+			// Test with EnableWarmup 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.EnableWarmup).To(BeNil())
+		})
+
+		It("should inherit EnableWarmup from manager config", func() {
+			// Test with manager default setting EnableWarmup to true
+			managerWithWarmup, err := manager.New(cfg, manager.Options{
+				Controller: config.Controller{
+					EnableWarmup: 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())
+			Expect(internalCtrlInheritingWarmup.EnableWarmup).To(HaveValue(BeTrue()))
+
+			// Test that explicit controller setting overrides manager setting
+			ctrlOverridingWarmup, err := controller.New("override-warmup-disabled", managerWithWarmup, controller.Options{
+				Reconciler:   reconcile.Func(nil),
+				EnableWarmup: ptr.To(false),
+			})
+			Expect(err).NotTo(HaveOccurred())
+
+			internalCtrlOverridingWarmup, ok := ctrlOverridingWarmup.(*internalcontroller.Controller[reconcile.Request])
+			Expect(ok).To(BeTrue())
+			Expect(internalCtrlOverridingWarmup.EnableWarmup).To(HaveValue(BeFalse()))
+		})
 	})
 })