Skip to content

Commit 04eb9b5

Browse files
committed
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
1 parent af8d903 commit 04eb9b5

File tree

3 files changed

+91
-0
lines changed

3 files changed

+91
-0
lines changed

pkg/manager/internal.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,10 @@ type controllerManager struct {
173173
// internalProceduresStop channel is used internally to the manager when coordinating
174174
// the proper shutdown of servers. This channel is also used for dependency injection.
175175
internalProceduresStop chan struct{}
176+
177+
// prestartHooks are functions that are run immediately before calling the Start functions
178+
// of the leader election runnables.
179+
prestartHooks []func(ctx context.Context) error
176180
}
177181

178182
type hasCache interface {
@@ -606,6 +610,15 @@ func (cm *controllerManager) engageStopProcedure(stopComplete <-chan struct{}) e
606610
}
607611

608612
func (cm *controllerManager) startLeaderElectionRunnables() error {
613+
for _, hook := range cm.prestartHooks {
614+
if err := hook(cm.internalCtx); err != nil {
615+
return err
616+
}
617+
}
618+
619+
// All the prestart hooks have been run, clear the slice to free the underlying resources.
620+
cm.prestartHooks = nil
621+
609622
return cm.runnables.LeaderElection.Start(cm.internalCtx)
610623
}
611624

pkg/manager/manager.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,10 @@ type Options struct {
305305
// stopped with the manager.
306306
makeBroadcaster intrec.EventBroadcasterProducer
307307

308+
// PrestartHooks are functions called by the manager after being elected leader,
309+
// immediately before calling the leader election controllers.
310+
PrestartHooks []func(ctx context.Context) error
311+
308312
// Dependency injection for testing
309313
newRecorderProvider func(config *rest.Config, scheme *runtime.Scheme, logger logr.Logger, makeBroadcaster intrec.EventBroadcasterProducer) (*intrec.Provider, error)
310314
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) {
449453
internalProceduresStop: make(chan struct{}),
450454
leaderElectionStopped: make(chan struct{}),
451455
leaderElectionReleaseOnCancel: options.LeaderElectionReleaseOnCancel,
456+
prestartHooks: options.PrestartHooks,
452457
}, nil
453458
}
454459

@@ -651,5 +656,9 @@ func setOptionsDefaults(options Options) Options {
651656
options.BaseContext = defaultBaseContext
652657
}
653658

659+
if options.PrestartHooks == nil {
660+
options.PrestartHooks = []func(context.Context) error{}
661+
}
662+
654663
return options
655664
}

pkg/manager/manager_test.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1076,6 +1076,75 @@ var _ = Describe("manger.Manager", func() {
10761076
<-runnableStopped
10771077
})
10781078

1079+
It("should run prestart hooks before calling Start on leader election runnables", func() {
1080+
var m Manager
1081+
var err error
1082+
1083+
runnableRan := make(chan struct{}, 0)
1084+
1085+
options.PrestartHooks = []func(ctx context.Context) error{
1086+
func(ctx context.Context) error {
1087+
defer GinkgoRecover()
1088+
Expect(m.Elected()).ShouldNot(BeClosed())
1089+
Consistently(runnableRan).ShouldNot(BeClosed())
1090+
return nil
1091+
},
1092+
}
1093+
1094+
m, err = New(cfg, options)
1095+
Expect(err).NotTo(HaveOccurred())
1096+
for _, cb := range callbacks {
1097+
cb(m)
1098+
}
1099+
1100+
Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
1101+
close(runnableRan)
1102+
return nil
1103+
})))
1104+
1105+
ctx, cancel := context.WithCancel(context.Background())
1106+
defer cancel()
1107+
go func() {
1108+
defer GinkgoRecover()
1109+
Expect(m.Elected()).ShouldNot(BeClosed())
1110+
Expect(m.Start(ctx)).NotTo(HaveOccurred())
1111+
}()
1112+
1113+
<-m.Elected()
1114+
})
1115+
1116+
It("should not run leader election runnables if prestart hooks fail", func() {
1117+
var m Manager
1118+
var err error
1119+
1120+
runnableRan := make(chan struct{}, 0)
1121+
1122+
options.PrestartHooks = []func(ctx context.Context) error{
1123+
func(ctx context.Context) error {
1124+
defer GinkgoRecover()
1125+
Expect(m.Elected()).ShouldNot(BeClosed())
1126+
Consistently(runnableRan).ShouldNot(BeClosed())
1127+
return errors.New("prestart hook failed")
1128+
},
1129+
}
1130+
1131+
m, err = New(cfg, options)
1132+
Expect(err).NotTo(HaveOccurred())
1133+
for _, cb := range callbacks {
1134+
cb(m)
1135+
}
1136+
1137+
Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
1138+
close(runnableRan)
1139+
return nil
1140+
})))
1141+
1142+
ctx, cancel := context.WithCancel(context.Background())
1143+
defer cancel()
1144+
1145+
Expect(m.Elected()).ShouldNot(BeClosed())
1146+
Expect(m.Start(ctx)).Should(MatchError(ContainSubstring("prestart hook failed")))
1147+
})
10791148
}
10801149

10811150
Context("with defaults", func() {

0 commit comments

Comments
 (0)