Try to do smarter handling of leader election

Author: thokra-nav

internal/cmd/api/api.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"os"
 	"os/signal"
 	"syscall"
@@ -197,6 +198,7 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error {
 		if err != nil {
 			return fmt.Errorf("creating k8s clients: %w", err)
 		}
+		log.WithField("envs", maps.Keys(k8sClients)).Info("Start event watcher")
 		eventWatcher := event.NewWatcher(pool, k8sClients, log)
 		go eventWatcher.Run(ctx)
 	}

internal/kubernetes/event/watcher.go

@@ -14,7 +14,7 @@ import (
 	eventsql "github.com/nais/api/internal/kubernetes/event/searchsql"
 	"github.com/nais/api/internal/leaderelection"
 	"github.com/sirupsen/logrus"
-	"golang.org/x/sync/errgroup"
+	"github.com/sourcegraph/conc/pool"
 	eventv1 "k8s.io/api/events/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/watch"
@@ -26,6 +26,9 @@ type Watcher struct {
 	clients map[string]kubernetes.Interface
 	events  chan eventsql.UpsertParams
 	log     logrus.FieldLogger
+	wg      *pool.ContextPool
+
+	cancel context.CancelFunc
 }
 
 func NewWatcher(pool *pgxpool.Pool, clients map[string]kubernetes.Interface, log logrus.FieldLogger) *Watcher {
@@ -38,25 +41,50 @@ func NewWatcher(pool *pgxpool.Pool, clients map[string]kubernetes.Interface, log
 }
 
 func (w *Watcher) Run(ctx context.Context) {
-	wg, ctx := errgroup.WithContext(ctx)
-	for env, client := range w.clients {
-		wg.Go(func() error {
-			return w.run(ctx, env, client)
-		})
+	w.wg = pool.New().WithErrors().WithContext(ctx)
+
+	leaderelection.RegisterOnStartedLeading(w.onStartedLeading)
+	leaderelection.RegisterOnStoppedLeading(w.onStoppedLeading)
+	if leaderelection.IsLeader() {
+		w.onStartedLeading(ctx)
 	}
 
-	wg.Go(func() error {
+	w.wg.Go(func(ctx context.Context) error {
 		return w.batchInsert(ctx)
 	})
 
-	if err := wg.Wait(); err != nil {
+	if err := w.wg.Wait(); err != nil {
 		w.log.WithError(err).Error("error running events watcher")
 	}
 }
 
+func (w *Watcher) onStoppedLeading() {
+	if w.cancel != nil {
+		w.cancel()
+	}
+}
+
+func (w *Watcher) onStartedLeading(_ context.Context) {
+	if w.cancel != nil {
+		w.cancel()
+	}
+
+	cancel := make(chan struct{})
+
+	w.cancel = func() {
+		close(cancel)
+	}
+
+	for env, client := range w.clients {
+		w.wg.Go(func(ctx context.Context) error {
+			return w.run(ctx, env, client, cancel)
+		})
+	}
+}
+
 var regHorizontalPodAutoscaler = regexp.MustCompile(`New size: (\d+); reason: (\w+).*(below|above) target`)
 
-func (w *Watcher) run(ctx context.Context, env string, client kubernetes.Interface) error {
+func (w *Watcher) run(ctx context.Context, env string, client kubernetes.Interface, cancel chan struct{}) error {
 	// Events we want to watch for
 	// SuccessfulRescale - Check for successful rescale events
 	// Killing - Check for liveness failures
@@ -103,6 +131,8 @@ func (w *Watcher) run(ctx context.Context, env string, client kubernetes.Interfa
 		select {
 		case <-ctx.Done():
 			return nil
+		case <-cancel:
+			return nil
 		case event := <-rescale.ResultChan():
 			handleEvent(event, func(e *eventv1.Event) (eventsql.UpsertParams, bool) {
 				if !strings.HasPrefix(e.Note, "New size") {

internal/leaderelection/leaderelection.go

@@ -14,6 +14,19 @@ import (
 
 var elector *leaderelection.LeaderElector
 
+var callbacks = struct {
+	onStartedLeading []func(context.Context)
+	onStoppedLeading []func()
+}{}
+
+func RegisterOnStartedLeading(f func(context.Context)) {
+	callbacks.onStartedLeading = append(callbacks.onStartedLeading, f)
+}
+
+func RegisterOnStoppedLeading(f func()) {
+	callbacks.onStoppedLeading = append(callbacks.onStoppedLeading, f)
+}
+
 func Start(ctx context.Context, client kubernetes.Interface, leaseName, namespace string, log logrus.FieldLogger) error {
 	id, err := os.Hostname()
 	if err != nil {
@@ -40,9 +53,17 @@ func Start(ctx context.Context, client kubernetes.Interface, leaseName, namespac
 		Callbacks: leaderelection.LeaderCallbacks{
 			OnStartedLeading: func(context.Context) {
 				log.Info("Started leading")
+
+				for _, f := range callbacks.onStartedLeading {
+					f(ctx)
+				}
 			},
 			OnStoppedLeading: func() {
 				log.Info("Stopped leading")
+
+				for _, f := range callbacks.onStoppedLeading {
+					f()
+				}
 			},
 			OnNewLeader: func(identity string) {
 				log.Infof("New leader: %s", identity)