leader_election_test.go

package leaderelection

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes/fake"

	metaV1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	k8s "k8s.io/client-go/kubernetes"
	k8stesting "k8s.io/client-go/testing"
)

func TestLeaderEngineError(t *testing.T) {
	currentRetryDelay := maxRetryDelay
	maxRetryDelay = 5 * time.Second
	defer func() {
		maxRetryDelay = currentRetryDelay
	}()
	_, err := New()
	if err == nil {
		t.Errorf("Failed to obtain timeout error")
	}
}

func TestInitializeLeaderEngineFailsFastOnRBACError(t *testing.T) {
	cases := []struct {
		desc     string
		reactErr error
	}{
		{
			desc: "forbidden",
			reactErr: errors.NewForbidden(
				schema.GroupResource{Group: "coordination.k8s.io", Resource: "leases"},
				"leader-election-lease",
				fmt.Errorf("rbac"),
			),
		},
		{
			desc:     "unauthorized",
			reactErr: errors.NewUnauthorized("not authorized"),
		},
	}

	for _, tc := range cases {
		t.Run(tc.desc, func(t *testing.T) {
			var callCount atomic.Int32
			client := fake.NewSimpleClientset()
			client.PrependReactor("get", "leases", func(_ k8stesting.Action) (bool, runtime.Object, error) {
				callCount.Add(1)
				return true, nil, tc.reactErr
			})

			restoreAPIClient := overrideGetAPIClient(client)
			defer restoreAPIClient()

			currentRetryDelay := maxRetryDelay
			maxRetryDelay = 5 * time.Second
			defer func() {
				maxRetryDelay = currentRetryDelay
			}()

			_, err := New(
				WithLeaseNamespace("test-namespace"),
				WithLogger(&testLogger{t: t}),
				WithErrorLogger(&testLogger{t: t}),
			)

			if err == nil {
				t.Fatalf("%s: got: nil error, want: non-nil error", tc.desc)
			}
			if !strings.Contains(err.Error(), "failed to obtain leases from namespace") {
				t.Errorf("%s: error message mismatch: got: %v, want: contains %q", tc.desc, err, "failed to obtain leases from namespace")
			}
			if c := callCount.Load(); c != 1 {
				t.Errorf("%s: did not fail fast: got: %d call(s), want: 1", tc.desc, c)
			}
		})
	}
}

func TestInitializeLeaderEngineRetriesOnTransientError(t *testing.T) {
	var callCount atomic.Int32
	client := fake.NewSimpleClientset()
	client.PrependReactor("get", "leases", func(_ k8stesting.Action) (bool, runtime.Object, error) {
		callCount.Add(1)
		return true, nil, errors.NewServerTimeout(
			schema.GroupResource{Group: "coordination.k8s.io", Resource: "leases"},
			"get",
			1,
		)
	})

	restoreAPIClient := overrideGetAPIClient(client)
	defer restoreAPIClient()

	currentMax := maxRetryDelay
	currentInitial := initialRetryDelay
	maxRetryDelay = 500 * time.Millisecond
	initialRetryDelay = 50 * time.Millisecond
	defer func() {
		maxRetryDelay = currentMax
		initialRetryDelay = currentInitial
	}()

	_, err := New(
		WithLeaseNamespace("test-namespace"),
		WithLogger(&testLogger{t: t}),
		WithErrorLogger(&testLogger{t: t}),
	)

	if err == nil {
		t.Fatalf("got: nil error, want: timeout error after retries")
	}
	if c := callCount.Load(); c < 2 {
		t.Errorf("transient error did not retry: got: %d call(s), want: >= 2", c)
	}
}

func overrideGetAPIClient(client k8s.Interface) func() {
	original := GetAPIClient
	GetAPIClient = func() (k8s.Interface, error) { return client, nil }
	return func() { GetAPIClient = original }
}

func TestLeaseAcquirement(t *testing.T) {
	const leaseName = "edgedelta-leader-election-lease-test"
	const leaseNamespace = "test-namespace"
	const holderIdentity = "test-identity"

	client := fake.NewSimpleClientset()
	ctx, cancel := context.WithCancel(context.Background())
	le := &K8sLeaderEngine{
		holderIdentity:     holderIdentity,
		leaseName:          leaseName,
		leaseNamespace:     leaseNamespace,
		leaseDuration:      1 * time.Second,
		renewDeadline:      500 * time.Millisecond,
		retryPeriod:        200 * time.Millisecond,
		apiClient:          client,
		coreClient:         client.CoreV1(),
		coordinationClient: client.CoordinationV1(),
		logger:             &testLogger{t: t},
		errorLogger:        &testLogger{t: t},
		parentCtx:          context.Background(),
		ctx:                ctx,
		ctxCancel:          cancel,
		stopped:            make(chan struct{}),
	}

	_, err := le.coordinationClient.Leases(leaseNamespace).Get(ctx, le.leaseName, metaV1.GetOptions{})
	if !errors.IsNotFound(err) {
		t.Errorf("Had an error whilst obtaining current leases, err: %v", err)
	}

	le.leaderElector, err = le.newLeaderElector(ctx)
	if err != nil {
		t.Errorf("Had an error whilst creating a leader elector, err: %v", err)
	}

	lease, err := le.coordinationClient.Leases(leaseNamespace).Get(ctx, le.leaseName, metaV1.GetOptions{})
	if err != nil {
		t.Errorf("Had an error whilst obtaining current leases, err: %v", err)
	}
	if lease.Name != leaseName {
		t.Errorf("Lease names does not match, wanted %q, obtained %q", leaseName, lease.Name)
	}
	if lease.Name != leaseName {
		t.Errorf("Lease names does not match, wanted %q, obtained %q", leaseName, lease.Name)
	}

	if err := le.Start(); err != nil {
		t.Errorf("Failed to start leader election engine, err: %v", err)
	}

	lease, err = le.coordinationClient.Leases(leaseNamespace).Get(ctx, le.leaseName, metaV1.GetOptions{})
	if err != nil && !errors.IsNotFound(err) {
		t.Errorf("Had an error while obtaining current leases, err: %v", err)
	}
	if *lease.Spec.LeaseTransitions != 1 {
		t.Errorf("Number of lease transitions does not match, wanted %d, obtained %d", 1, *lease.Spec.LeaseTransitions)
	}
	if !le.IsLeader() {
		t.Errorf("Current client is not leader, leader is %q", le.GetLeader())
	}
	if err := le.Stop(); err != nil {
		t.Errorf("Failed to stop leader election engine, err: %v", err)
	}
}

type testLogger struct {
	t *testing.T
}

func (tl *testLogger) Log(format string, args ...interface{}) {
	tl.t.Logf(format, args...)
}

func TestRaceCondition(t *testing.T) {
	const leaseName = "leader-election-lease-test"
	const leaseNamespace = "test-namespace"
	const holderIdentity = "test-identity-"
	const number = 10

	client := fake.NewSimpleClientset()
	channelMap := map[string]chan struct{}{}
	leaderEngineMap := map[string]*K8sLeaderEngine{}
	for i := 0; i < number; i++ {
		identity := fmt.Sprintf("%s%d", holderIdentity, i)
		stopChannel := make(chan struct{}, 1)
		channelMap[identity] = stopChannel
		ctx, cancel := context.WithCancel(context.Background())
		le := &K8sLeaderEngine{
			holderIdentity:     identity,
			leaseName:          leaseName,
			leaseNamespace:     leaseNamespace,
			leaseDuration:      1 * time.Second,
			renewDeadline:      500 * time.Millisecond,
			retryPeriod:        200 * time.Millisecond,
			apiClient:          client,
			coreClient:         client.CoreV1(),
			coordinationClient: client.CoordinationV1(),
			logger:             &testLogger{t: t},
			errorLogger:        &testLogger{t: t},
			parentCtx:          context.Background(),
			ctx:                ctx,
			ctxCancel:          cancel,
			stopped:            make(chan struct{}),
		}
		leaderEngineMap[identity] = le

		var err error
		le.leaderElector, err = le.newLeaderElector(context.Background())
		if err != nil {
			t.Errorf("Had an error whilst creating a leader elector, err: %v", err)
		}
	}

	for identity, le := range leaderEngineMap {
		identity := identity
		le := le
		stopChannel := channelMap[identity]

		go testRunLeaderEngine(t, identity, le, stopChannel)
	}

	time.Sleep(500 * time.Millisecond)
	finishChannel := make(chan struct{}, 1)

	go func() {
		isFinished := map[string]bool{}
		count := 0
		for count < number {
			var currentLeader string
			for j := 0; j < number; j++ {
				identity := fmt.Sprintf("%s%d", holderIdentity, j)
				if isFinished[identity] {
					continue
				}

				le := leaderEngineMap[identity]
				if le.IsLeader() {
					currentLeader = identity
					break
				}
			}

			if currentLeader == "" {
				time.Sleep(10 * time.Millisecond) // Sleep some time to give time for reelection
				continue
			}

			t.Logf("Found leader, leader is %q, now stopping it", currentLeader)
			time.Sleep(50 * time.Millisecond)
			stopChannel := channelMap[currentLeader]
			stopChannel <- struct{}{}
			time.Sleep(50 * time.Millisecond)
			isFinished[currentLeader] = true
			count++
			t.Logf("Current count is %d, %d is remaining", count, number-count)
		}

		finishChannel <- struct{}{}
	}()

	select {
	case <-finishChannel:
		break
	case <-time.After(15 * time.Second):
		t.Errorf("Timed out after 15 seconds, possible deadlock")
	}
}

func testRunLeaderEngine(t *testing.T, identity string, le *K8sLeaderEngine, stopChannel chan struct{}) {
	rand.Seed(time.Now().UTC().UnixNano())

	randNum := (rand.Int() % 500)
	randInterval := time.Duration(randNum * int(time.Millisecond))
	time.Sleep(randInterval)

	if err := le.Start(); err != nil {
		t.Errorf("Failed to start leader election process for identity %q, err: %v", identity, err)
		return
	}

	<-stopChannel // wait for this to be signalled to be stopped

	if err := le.Stop(); err != nil {
		t.Errorf("Failed to stop leader election for identity %q, err: %v", identity, err)
		return
	}
	t.Logf("Stopped client with identity %q", identity)
}