feat: [x] refactor leases

Author: povilasv

Makefile

@@ -373,7 +373,7 @@ integration-tests: build-init build-toolkit ## Run integration tests (only tests
 		STORAGE_ACCESSKEYID=$(ROOT_MINIO_USER) \
 		STORAGE_SECRETACCESSKEY=$(ROOT_MINIO_PASSWORD) \
 		$(GOTESTSUM) --format short-verbose --junitfile integration-tests.xml --jsonfile integration-tests.json -- \
-		-coverprofile=integration-coverage.out -covermode=atomic -run "_Integration$$" ./internal/... ./pkg/... ./test/integration/components/...
+		-coverprofile=integration-coverage.out -covermode=atomic -run "_Integration$$" ./internal/... ./pkg/... ./test/integration/... ./test/e2e/...
 
 .PHONY: cover
 cover: unit-tests ## Generate and open test coverage report

cmd/api-server/main.go

@@ -217,8 +217,6 @@ func main() {
 		)
 	}
 
-	leaderTasks := make([]leader.Task, 0)
-
 	// If we don't have an API key but we do have a token for registration then attempt to register the runner.
 	if cfg.TestkubeProAPIKey == "" && cfg.TestkubeProAgentRegToken != "" {
 		runnerName := cfg.RunnerName
@@ -302,6 +300,22 @@ func main() {
 	clusterId, _ := configMapConfig.GetUniqueClusterId(ctx)
 	telemetryEnabled, _ := configMapConfig.GetTelemetryEnabled(ctx)
 
+	leaderIdentifier := resolveLeaderIdentifier()
+	leaderClusterID := clusterId
+	if leaderClusterID == "" {
+		leaderClusterID = "testkube-core"
+	} else {
+		leaderClusterID = fmt.Sprintf("%s-core", leaderClusterID)
+	}
+
+	coordinatorLogger := log.DefaultLogger.With("component", "leader-coordinator")
+	leaderCoordinator := leader.New(leaderLeaseBackend, leaderIdentifier, leaderClusterID, coordinatorLogger)
+	hasLeaderTasks := false
+	registerLeaderTask := func(task leader.Task) {
+		hasLeaderTasks = true
+		leaderCoordinator.Register(task)
+	}
+
 	// k8s clients
 	webhooksClient := executorsclientv1.NewWebhooksClient(kubeClient, cfg.TestkubeNamespace)
 	webhookTemplatesClient := executorsclientv1.NewWebhookTemplatesClient(kubeClient, cfg.TestkubeNamespace)
@@ -508,9 +522,12 @@ func main() {
 	// Update TestWorkflowExecution Kubernetes resource objects on status change
 	eventsEmitter.RegisterLoader(testworkflowexecutions.NewLoader(ctx, cfg.TestkubeNamespace, kubeClient))
 
-	g.Go(func() error {
-		eventsEmitter.Listen(ctx)
-		return nil
+	registerLeaderTask(leader.Task{
+		Name: "event-emitter",
+		Start: func(taskCtx context.Context) error {
+			eventsEmitter.RunLeader(taskCtx)
+			return nil
+		},
 	})
 
 	/////////////////////////////////
@@ -641,7 +658,7 @@ func main() {
 			eventsEmitter,
 		)
 		commons.ExitOnError("starting agent", err)
-		leaderTasks = append(leaderTasks, leader.Task{
+		registerLeaderTask(leader.Task{
 			Name: "agent",
 			Start: func(taskCtx context.Context) error {
 				err := agentHandle.Run(taskCtx)
@@ -662,26 +679,12 @@ func main() {
 		// TODO: Check why this simpler options is not working
 		// testkubeClientset := testkubeclientset.New(clientset.RESTClient())
 
-		var triggersLeaseBackend leasebackend.Repository
-		if controlPlane != nil {
-			triggersLeaseBackend = controlPlane.GetRepositoryManager().LeaseBackend()
-		} else {
-			// Fallback: Kubernetes Lease-based coordination (no external DB required)
-			triggersLeaseBackend = leasebackendk8s.NewK8sLeaseBackend(
-				clientset,
-				"testkube-triggers-lease",
-				cfg.TestkubeNamespace,
-				leasebackendk8s.WithLeaseName(cfg.TestkubeLeaseName),
-			)
-		}
-
-		triggerService := triggers.NewService(
+		triggerTasks := triggers.NewService(
 			cfg.RunnerName,
 			clientset,
 			testkubeClientset,
 			testWorkflowsClient,
 			testTriggersClient,
-			triggersLeaseBackend,
 			log.DefaultLogger,
 			eventBus,
 			metrics,
@@ -695,17 +698,16 @@ func main() {
 			triggers.WithTestTriggerControlPlane(cfg.TestTriggerControlPlane),
 			triggers.WithEventLabels(cfg.EventLabels),
 		)
-		log.DefaultLogger.Info("starting trigger service")
-		g.Go(func() error {
-			triggerService.Run(ctx)
-			return nil
-		})
+		log.DefaultLogger.Info("registering trigger tasks with shared leader coordinator")
+		for _, task := range triggerTasks {
+			registerLeaderTask(task)
+		}
 	} else {
 		log.DefaultLogger.Info("test triggers are disabled")
 	}
 
 	// telemetry based functions
-	leaderTasks = append(leaderTasks, leader.Task{
+	registerLeaderTask(leader.Task{
 		Name: "telemetry-heartbeat",
 		Start: func(taskCtx context.Context) error {
 			services.HandleTelemetryHeartbeat(taskCtx, clusterId, configMapConfig)
@@ -755,7 +757,7 @@ func main() {
 			schedulableResourceWatcher.WatchTestWorkflowTemplates,
 		)
 		// Start the new scheduler.
-		leaderTasks = append(leaderTasks, leader.Task{
+		registerLeaderTask(leader.Task{
 			Name: "cron-scheduler",
 			Start: func(taskCtx context.Context) error {
 				go func() {
@@ -779,22 +781,7 @@ func main() {
 		return httpServer.Run(ctx)
 	})
 
-	if len(leaderTasks) > 0 {
-		leaderIdentifier := resolveLeaderIdentifier()
-
-		leaderClusterID := clusterId
-		if leaderClusterID == "" {
-			leaderClusterID = "testkube-core"
-		} else {
-			leaderClusterID = fmt.Sprintf("%s-core", leaderClusterID)
-		}
-
-		coordinatorLogger := log.DefaultLogger.With("component", "leader-coordinator")
-		leaderCoordinator := leader.New(leaderLeaseBackend, leaderIdentifier, leaderClusterID, coordinatorLogger)
-		for _, task := range leaderTasks {
-			leaderCoordinator.Register(task)
-		}
-
+	if hasLeaderTasks {
 		g.Go(func() error {
 			return leaderCoordinator.Run(ctx)
 		})

pkg/event/emitter.go

@@ -126,6 +126,28 @@ const (
 	leaseClusterID     string = "event-emitters"
 )
 
+// RunLeader executes the leader-only loop under an externally managed coordinator.
+// It does not perform lease acquisition itself; it assumes the provided context
+// is canceled when leadership is lost or when shutdown is requested.
+func (e *Emitter) RunLeader(ctx context.Context) {
+	e.log.Info("event emitter starting under external coordinator")
+	e.closeBusOnContext(ctx)
+	e.leaderLoop(ctx)
+}
+
+func (e *Emitter) closeBusOnContext(ctx context.Context) {
+	go func() {
+		<-ctx.Done()
+		e.log.Info("event emitter closing event bus")
+		err := e.bus.Close()
+		if err != nil {
+			e.log.Errorw("error while closing event bus", "error", err)
+		} else {
+			e.log.Info("event emitter closed event bus")
+		}
+	}()
+}
+
 // TODO(emil): convert to using new common coordinator package for lease acquisition
 func (e *Emitter) leaseCheckLoop(ctx context.Context, leaseChan chan<- bool) {
 	e.log.Info("event emitter waiting for lease")
@@ -154,17 +176,7 @@ func (e *Emitter) leaseCheck(ctx context.Context, leaseChan chan<- bool) {
 // notifications.
 func (e *Emitter) Listen(ctx context.Context) {
 	e.log.Info("event emitter starting")
-	// Clean up
-	go func() {
-		<-ctx.Done()
-		e.log.Info("event emitter closing event bus")
-		err := e.bus.Close()
-		if err != nil {
-			e.log.Errorw("error while closing event bus", "error", err)
-		} else {
-			e.log.Info("event emitter closed event bus")
-		}
-	}()
+	e.closeBusOnContext(ctx)
 
 	// Start lease check loop
 	leaseChan := make(chan bool)

pkg/event/emitter_test.go

@@ -3,10 +3,12 @@ package event
 import (
 	"context"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
 	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
@@ -294,6 +296,99 @@ func TestEmitter_Listen_reconciliation(t *testing.T) {
 
 }
 
+type trackingBus struct {
+	subscribeTopicCalls []struct {
+		Topic string
+		Queue string
+	}
+	unsubscribeCalls []string
+	closeCalls       int
+	mu               sync.Mutex
+}
+
+func (b *trackingBus) Publish(event testkube.Event) error {
+	return nil
+}
+
+func (b *trackingBus) Subscribe(queue string, handler bus.Handler) error {
+	return nil
+}
+
+func (b *trackingBus) Unsubscribe(queue string) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.unsubscribeCalls = append(b.unsubscribeCalls, queue)
+	return nil
+}
+
+func (b *trackingBus) PublishTopic(topic string, event testkube.Event) error {
+	return nil
+}
+
+func (b *trackingBus) SubscribeTopic(topic string, queue string, handler bus.Handler) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.subscribeTopicCalls = append(b.subscribeTopicCalls, struct {
+		Topic string
+		Queue string
+	}{Topic: topic, Queue: queue})
+	return nil
+}
+
+func (b *trackingBus) Close() error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.closeCalls++
+	return nil
+}
+
+type noopLeaseBackend struct{}
+
+func (noopLeaseBackend) TryAcquire(ctx context.Context, id, clusterID string) (bool, error) {
+	return false, nil
+}
+
+func TestEmitter_RunLeader_UsesExternalCoordinator(t *testing.T) {
+	t.Parallel()
+
+	tb := &trackingBus{}
+	emitter := NewEmitter(tb, noopLeaseBackend{}, "agentevents", "")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		emitter.RunLeader(ctx)
+		close(done)
+	}()
+
+	// Allow leader loop to subscribe then cancel.
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("expected RunLeader to exit after context cancellation")
+	}
+
+	assert.Len(t, tb.subscribeTopicCalls, 1, "expected one topic subscription")
+	assert.Equal(t, "agentevents.>", tb.subscribeTopicCalls[0].Topic)
+	assert.Equal(t, "emitter", tb.subscribeTopicCalls[0].Queue)
+
+	require.Eventually(t, func() bool {
+		tb.mu.Lock()
+		defer tb.mu.Unlock()
+		return len(tb.unsubscribeCalls) == 1
+	}, time.Second, 10*time.Millisecond, "expected unsubscribe on shutdown")
+
+	assert.Equal(t, "emitter", tb.unsubscribeCalls[0])
+
+	tb.mu.Lock()
+	closeCalls := tb.closeCalls
+	tb.mu.Unlock()
+	assert.Equal(t, 1, closeCalls, "expected bus close on shutdown")
+}
+
 func newExampleTestEvent1() testkube.Event {
 	return testkube.Event{
 		Id:                    "eventID1",