fix: promote schedulers on creation if controller is leader

Signed-off-by: Calum Murray <cmurray@redhat.com>

Author: Cali0707

control-plane/pkg/reconciler/consumergroup/controller.go

@@ -209,20 +209,8 @@ func NewController(ctx context.Context, watcher configmap.Watcher) *controller.I
 
 	impl := cgreconciler.NewImpl(ctx, r, func(impl *controller.Impl) controller.Options {
 		return controller.Options{
-			PromoteFunc: func(bkt reconciler.Bucket) {
-				for _, value := range schedulerMgr.getAll() {
-					if ss, ok := value.Scheduler.(*statefulsetscheduler.StatefulSetScheduler); ok {
-						ss.Promote(bkt, nil)
-					}
-				}
-			},
-			DemoteFunc: func(bkt reconciler.Bucket) {
-				for _, value := range schedulerMgr.getAll() {
-					if ss, ok := value.Scheduler.(*statefulsetscheduler.StatefulSetScheduler); ok {
-						ss.Demote(bkt)
-					}
-				}
-			},
+			PromoteFunc: schedulerMgr.promote,
+			DemoteFunc:  schedulerMgr.demote,
 		}
 	})
 
@@ -472,8 +460,9 @@ func createStatefulSetScheduler(ctx context.Context, c SchedulerConfig, lister s
 }
 
 type schedulerManager struct {
-	mu         sync.RWMutex
-	schedulers map[string]Scheduler
+	mu           sync.RWMutex
+	schedulers   map[string]Scheduler // protected by mu
+	leaderBucket reconciler.Bucket    // protected by mu
 
 	ctx                   context.Context
 	config                SchedulerConfig
@@ -522,7 +511,16 @@ func (sm *schedulerManager) createSchedulerForStatefulSet(ssName string) bool {
 	logger := logging.FromContext(sm.ctx)
 	logger.Infow("Creating scheduler for StatefulSet", zap.String("statefulset", ssName), zap.String("scheduler", schedulerKey))
 
-	sm.schedulers[schedulerKey] = createKafkaScheduler(sm.ctx, sm.config, ssName, sm.dispatcherPodInformer)
+	scheduler := createKafkaScheduler(sm.ctx, sm.config, ssName, sm.dispatcherPodInformer)
+	sm.schedulers[schedulerKey] = scheduler
+
+	// If we're already leader, promote the new scheduler immediately
+	if sm.leaderBucket != nil {
+		if ss, ok := scheduler.Scheduler.(*statefulsetscheduler.StatefulSetScheduler); ok {
+			ss.Promote(sm.leaderBucket, nil)
+		}
+	}
+
 	return true
 }
 
@@ -545,3 +543,27 @@ func (sm *schedulerManager) removeSchedulerForStatefulSet(ssName string) bool {
 	delete(sm.schedulers, schedulerKey)
 	return true
 }
+
+func (sm *schedulerManager) promote(bkt reconciler.Bucket) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	sm.leaderBucket = bkt
+	for _, value := range sm.schedulers {
+		if ss, ok := value.Scheduler.(*statefulsetscheduler.StatefulSetScheduler); ok {
+			ss.Promote(bkt, nil)
+		}
+	}
+}
+
+func (sm *schedulerManager) demote(bkt reconciler.Bucket) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	sm.leaderBucket = nil
+	for _, value := range sm.schedulers {
+		if ss, ok := value.Scheduler.(*statefulsetscheduler.StatefulSetScheduler); ok {
+			ss.Demote(bkt)
+		}
+	}
+}