From 51118413283acb118de65b54a497429a3d378bed Mon Sep 17 00:00:00 2001
From: Tom Wieczorek <twieczorek@mirantis.com>
Date: Tue, 3 Sep 2024 11:14:41 +0200
Subject: [PATCH] Use new leader election mechanism in applier manager

Signed-off-by: Tom Wieczorek <twieczorek@mirantis.com>
---
 pkg/applier/manager.go                        | 35 +++++++++----------
 .../controller/leaderelector/dummy.go         | 12 +++++++
 .../controller/leaderelector/types.go         | 12 ++++---
 .../workerconfig/reconciler_test.go           |  5 +++
 4 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/pkg/applier/manager.go b/pkg/applier/manager.go
index bdc26015fe8d..1d6f1254cf87 100644
--- a/pkg/applier/manager.go
+++ b/pkg/applier/manager.go
@@ -65,27 +65,26 @@ func (m *Manager) Init(ctx context.Context) error {
 	m.log = logrus.WithField("component", constant.ApplierManagerComponentName)
 	m.bundleDir = m.K0sVars.ManifestsDir
 
-	m.LeaderElector.AddAcquiredLeaseCallback(func() {
-		ctx, cancel := context.WithCancelCause(ctx)
-		stopped := make(chan struct{})
-
-		m.stop = func(reason string) { cancel(errors.New(reason)); <-stopped }
-		go func() {
-			defer close(stopped)
-			wait.UntilWithContext(ctx, m.runWatchers, 1*time.Minute)
-		}()
-	})
-	m.LeaderElector.AddLostLeaseCallback(func() {
-		if m.stop != nil {
-			m.stop("lost leadership")
-		}
-	})
-
-	return err
+	return nil
 }
 
 // Run runs the Manager
-func (m *Manager) Start(_ context.Context) error {
+func (m *Manager) Start(context.Context) error {
+	ctx, cancel := context.WithCancelCause(context.Background())
+	stopped := make(chan struct{})
+
+	m.stop = func(reason string) {
+		cancel(errors.New(reason))
+		<-stopped
+	}
+
+	go func() {
+		defer close(stopped)
+		leaderelector.RunLeaderTasks(ctx, m.LeaderElector.Peek, func(ctx context.Context) {
+			wait.UntilWithContext(ctx, m.runWatchers, time.Minute)
+		})
+	}()
+
 	return nil
 }
 
diff --git a/pkg/component/controller/leaderelector/dummy.go b/pkg/component/controller/leaderelector/dummy.go
index 4d7b0607e771..329f7df6b0c1 100644
--- a/pkg/component/controller/leaderelector/dummy.go
+++ b/pkg/component/controller/leaderelector/dummy.go
@@ -20,6 +20,7 @@ import (
 	"context"
 
 	"github.com/k0sproject/k0s/pkg/component/manager"
+	"github.com/k0sproject/k0s/pkg/leaderelection"
 )
 
 type Dummy struct {
@@ -38,6 +39,17 @@ func (l *Dummy) AddAcquiredLeaseCallback(fn func()) {
 	l.callbacks = append(l.callbacks, fn)
 }
 
+var neverClosingChan = make(<-chan struct{})
+
+func (l *Dummy) Peek() (leaderelection.LeaseStatus, <-chan struct{}) {
+	var leaseStatus leaderelection.LeaseStatus
+	if l.Leader {
+		leaseStatus = leaderelection.LeaseStatusAcquired
+	}
+
+	return leaseStatus, neverClosingChan
+}
+
 func (l *Dummy) AddLostLeaseCallback(func()) {}
 
 func (l *Dummy) Start(_ context.Context) error {
diff --git a/pkg/component/controller/leaderelector/types.go b/pkg/component/controller/leaderelector/types.go
index 8e49b1e05c9d..57df33d43e93 100644
--- a/pkg/component/controller/leaderelector/types.go
+++ b/pkg/component/controller/leaderelector/types.go
@@ -16,11 +16,13 @@ limitations under the License.
 
 package leaderelector
 
+import "github.com/k0sproject/k0s/pkg/leaderelection"
+
 // Interface is the common leader elector component to manage each controller leader status.
-//
-// Deprecated: Use [LeasePool.Peek] instead.
 type Interface interface {
-	IsLeader() bool
-	AddAcquiredLeaseCallback(fn func())
-	AddLostLeaseCallback(fn func())
+	IsLeader() bool                     // Deprecated: Use [Interface.Peek] instead.
+	AddAcquiredLeaseCallback(fn func()) // Deprecated: Use [Interface.Peek] instead.
+	AddLostLeaseCallback(fn func())     // Deprecated: Use [Interface.Peek] instead.
+
+	Peek() (leaderelection.LeaseStatus, <-chan struct{})
 }
diff --git a/pkg/component/controller/workerconfig/reconciler_test.go b/pkg/component/controller/workerconfig/reconciler_test.go
index 75e60b114f8a..98f1f714e9fe 100644
--- a/pkg/component/controller/workerconfig/reconciler_test.go
+++ b/pkg/component/controller/workerconfig/reconciler_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/k0sproject/k0s/pkg/config"
 	"github.com/k0sproject/k0s/pkg/constant"
 	kube "github.com/k0sproject/k0s/pkg/kubernetes"
+	"github.com/k0sproject/k0s/pkg/leaderelection"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -913,3 +914,7 @@ func (e *mockLeaderElector) AddAcquiredLeaseCallback(fn func()) {
 func (e *mockLeaderElector) AddLostLeaseCallback(func()) {
 	panic("not expected to be called in tests")
 }
+
+func (e *mockLeaderElector) Peek() (leaderelection.LeaseStatus, <-chan struct{}) {
+	panic("not expected to be called in tests")
+}