feat(manager): add prestart hook support

When implementing a controller that uses leader election, there maybe be
work that needs to be done after winning the election but before
processing enqueued requests. For example, a controller may need to
build up an internal mapping of the current state of the cluster before
it can begin reconciling.

This changeset adds support for adding prestart hooks to
controller-runtime's manager implementation. This hook runs after the
manager has been elected leader, immediately before the leader election
controllers are started.

Related #607

Author: terinjokes

pkg/manager/internal.go

@@ -173,6 +173,10 @@ type controllerManager struct {
 	// internalProceduresStop channel is used internally to the manager when coordinating
 	// the proper shutdown of servers. This channel is also used for dependency injection.
 	internalProceduresStop chan struct{}
+
+	// prestartHooks are functions that are run immediately before calling the Start functions
+	// of the leader election runnables.
+	prestartHooks []func(ctx context.Context) error
 }
 
 type hasCache interface {
@@ -606,6 +610,15 @@ func (cm *controllerManager) engageStopProcedure(stopComplete <-chan struct{}) e
 }
 
 func (cm *controllerManager) startLeaderElectionRunnables() error {
+	for _, hook := range cm.prestartHooks {
+		if err := hook(cm.internalCtx); err != nil {
+			return err
+		}
+	}
+
+	// All the prestart hooks have been run, clear the slice to free the underlying resources.
+	cm.prestartHooks = nil
+
 	return cm.runnables.LeaderElection.Start(cm.internalCtx)
 }
 

pkg/manager/manager.go

@@ -305,6 +305,10 @@ type Options struct {
 	// stopped with the manager.
 	makeBroadcaster intrec.EventBroadcasterProducer
 
+	// PrestartHooks are functions called by the manager after being elected leader,
+	// immediately before calling the leader election controllers.
+	PrestartHooks []func(ctx context.Context) error
+
 	// Dependency injection for testing
 	newRecorderProvider    func(config *rest.Config, scheme *runtime.Scheme, logger logr.Logger, makeBroadcaster intrec.EventBroadcasterProducer) (*intrec.Provider, error)
 	newResourceLock        func(config *rest.Config, recorderProvider recorder.Provider, options leaderelection.Options) (resourcelock.Interface, error)
@@ -449,6 +453,7 @@ func New(config *rest.Config, options Options) (Manager, error) {
 		internalProceduresStop:        make(chan struct{}),
 		leaderElectionStopped:         make(chan struct{}),
 		leaderElectionReleaseOnCancel: options.LeaderElectionReleaseOnCancel,
+		prestartHooks:                 options.PrestartHooks,
 	}, nil
 }
 
@@ -651,5 +656,9 @@ func setOptionsDefaults(options Options) Options {
 		options.BaseContext = defaultBaseContext
 	}
 
+	if options.PrestartHooks == nil {
+		options.PrestartHooks = []func(context.Context) error{}
+	}
+
 	return options
 }

pkg/manager/manager_test.go

@@ -1076,6 +1076,75 @@ var _ = Describe("manger.Manager", func() {
 				<-runnableStopped
 			})
 
+			It("should run prestart hooks before calling Start on leader election runnables", func() {
+				var m Manager
+				var err error
+
+				runnableRan := make(chan struct{}, 0)
+
+				options.PrestartHooks = []func(ctx context.Context) error{
+					func(ctx context.Context) error {
+						defer GinkgoRecover()
+						Expect(m.Elected()).ShouldNot(BeClosed())
+						Consistently(runnableRan).ShouldNot(BeClosed())
+						return nil
+					},
+				}
+
+				m, err = New(cfg, options)
+				Expect(err).NotTo(HaveOccurred())
+				for _, cb := range callbacks {
+					cb(m)
+				}
+
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					close(runnableRan)
+					return nil
+				})))
+
+				ctx, cancel := context.WithCancel(context.Background())
+				defer cancel()
+				go func() {
+					defer GinkgoRecover()
+					Expect(m.Elected()).ShouldNot(BeClosed())
+					Expect(m.Start(ctx)).NotTo(HaveOccurred())
+				}()
+
+				<-m.Elected()
+			})
+
+			It("should not run leader election runnables if prestart hooks fail", func() {
+				var m Manager
+				var err error
+
+				runnableRan := make(chan struct{}, 0)
+
+				options.PrestartHooks = []func(ctx context.Context) error{
+					func(ctx context.Context) error {
+						defer GinkgoRecover()
+						Expect(m.Elected()).ShouldNot(BeClosed())
+						Consistently(runnableRan).ShouldNot(BeClosed())
+						return errors.New("prestart hook failed")
+					},
+				}
+
+				m, err = New(cfg, options)
+				Expect(err).NotTo(HaveOccurred())
+				for _, cb := range callbacks {
+					cb(m)
+				}
+
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					close(runnableRan)
+					return nil
+				})))
+
+				ctx, cancel := context.WithCancel(context.Background())
+				defer cancel()
+
+				Expect(m.Elected()).ShouldNot(BeClosed())
+				Expect(m.Start(ctx)).Should(MatchError(ContainSubstring("prestart hook failed")))
+			})
 		}
 
 		Context("with defaults", func() {