mesh-controller/internal/lease/lease_test.go

package lease

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"

	"github.com/novox/mesh-controller/internal/testbus"
)

// Two controllers at once, on a real bus (novox/hq to-be 45 §6, replay R1's half that lives here):
// one takes the lease and the other waits; a handover gives the second a higher epoch; a holder that
// stops renewing is taken over once its key expires, and has stopped acting before then; a renewal
// refused is a loss at once; and a bucket raised again from nothing never issues an epoch twice.
//
//	go test ./internal/lease/   (each test on a bus of its own: internal/testbus)
//
// Each test runs beside the others (novox/hq issue 717): they share nothing but the clock, and wait on it
// most of their time, so one after another they took half a minute of a check.

// The bounds in these tests are the design's divided by five, so the cases run in seconds.
const (
	testTTL   = 3 * time.Second
	testRenew = time.Second
	testPoll  = 100 * time.Millisecond
)

// aBucket is a fresh lease bucket on the test bus, with the test's age.
func aBucket(t *testing.T) (jetstream.JetStream, string) {
	t.Helper()
	url := testbus.URL(t)
	conn, err := nats.Connect(url)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(conn.Close)
	js, err := jetstream.New(conn)
	if err != nil {
		t.Fatal(err)
	}
	bucket := fmt.Sprintf("lease-test-%d", time.Now().UnixNano())
	if _, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: bucket, History: 1,
		TTL: testTTL, Storage: jetstream.FileStorage}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = js.DeleteKeyValue(context.Background(), bucket) })
	return js, bucket
}

// aController is one controller instance's lease over the bucket, on a connection of its own.
func aController(t *testing.T, bucket, name string, floor func(context.Context) (uint64, error)) *Lease {
	t.Helper()
	conn, err := nats.Connect(testbus.URL(t))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(conn.Close)
	js, err := jetstream.New(conn)
	if err != nil {
		t.Fatal(err)
	}
	var said sync.Mutex
	l, err := Open(t.Context(), js, bucket, Options{Holder: Holder{Instance: name, Host: name},
		RenewEvery: testRenew, Margin: testRenew / 2, Poll: testPoll, Floor: floor,
		Say: func(format string, args ...any) {
			said.Lock()
			defer said.Unlock()
			t.Logf(name+": "+format, args...)
		}})
	if err != nil {
		t.Fatal(err)
	}
	return l
}

// takeWithin takes the lease in a goroutine and answers when it did, or fails past the bound.
func takeWithin(t *testing.T, l *Lease) <-chan uint64 {
	t.Helper()
	took := make(chan uint64, 1)
	go func() {
		epoch, err := l.Take(t.Context())
		if err != nil {
			t.Errorf("taking: %v", err)
			close(took)
			return
		}
		took <- epoch
	}()
	return took
}

func TestTheSecondControllerWaitsAndTakesAHigherEpochOnHandover(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a, b := aController(t, bucket, "a", nil), aController(t, bucket, "b", nil)

	epochA, err := a.Take(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	keeping, stop := context.WithCancel(t.Context())
	defer stop()
	go a.Keep(keeping)

	// b waits while a renews — across more than one age of the key, so it is the renewals holding it.
	took := takeWithin(t, b)
	select {
	case epoch := <-took:
		t.Fatalf("b took the lease (epoch %d) while a held it and renewed", epoch)
	case <-time.After(testTTL + testRenew):
	}
	if _, err := b.Epoch(); !errors.Is(err, ErrNotHeld) {
		t.Fatalf("b, waiting, may act: %v", err)
	}
	if e, err := a.Epoch(); err != nil || e != epochA {
		t.Fatalf("a, holding, answers %d, %v", e, err)
	}
	holder, found, err := b.Current(t.Context())
	if err != nil || !found || holder.Instance != "a" || holder.Epoch != epochA {
		t.Fatalf("b reads the holder as %+v (%v, %v), want a at epoch %d", holder, found, err, epochA)
	}

	// a gives it back: b takes it at once, not after the key's age, at a higher epoch.
	stop()
	handedOver := time.Now()
	a.Release(t.Context())
	select {
	case epochB := <-took:
		if epochB <= epochA {
			t.Fatalf("b took epoch %d after a's %d: an epoch must only grow", epochB, epochA)
		}
		if waited := time.Since(handedOver); waited > testTTL {
			t.Fatalf("b took the lease %s after a gave it back: it waited out the key's age", waited)
		}
	case <-time.After(2 * testTTL):
		t.Fatal("b never took the lease a gave back")
	}
	if _, err := a.Epoch(); !errors.Is(err, ErrNotHeld) {
		t.Fatalf("a, having given it back, may still act: %v", err)
	}
	if _, err := a.TryTake(t.Context()); err == nil {
		t.Fatal("a took the lease again after giving it back: its successor is a new candidate")
	}
}

func TestAHolderThatStopsRenewingStopsActingBeforeItIsTakenOver(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a, b := aController(t, bucket, "a", nil), aController(t, bucket, "b", nil)
	epochA, err := a.Take(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	// a is stuck: it never renews (a process stopped, a goroutine wedged). It was not told anything.
	took := takeWithin(t, b)

	// Its gate closes by its own clock, before the key can expire.
	closed := time.Now()
	for a.Held() {
		if time.Since(closed) > testTTL {
			t.Fatal("a still may act past its key's age without renewing")
		}
		time.Sleep(20 * time.Millisecond)
	}
	select {
	case epoch := <-took:
		t.Fatalf("b took the lease (epoch %d) before a stopped acting", epoch)
	default:
	}
	select {
	case epochB := <-took:
		if epochB <= epochA {
			t.Fatalf("b took epoch %d after a's %d", epochB, epochA)
		}
		// And a, the moment it tries to renew, knows it lost: the key moved.
		if err := a.Renew(t.Context()); err == nil {
			t.Fatal("a renewed a lease b holds")
		}
		select {
		case <-a.Lost():
		default:
			t.Fatal("a's renewal was refused and a was not told it lost the lease")
		}
		if _, err := a.Epoch(); !errors.Is(err, ErrNotHeld) {
			t.Fatalf("a, having lost the lease, may act: %v", err)
		}
	case <-time.After(3 * testTTL):
		t.Fatal("b never took over a lease nobody renewed")
	}
}

func TestARenewalRefusedIsALossAtOnce(t *testing.T) {
	t.Parallel()
	js, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	if _, err := a.Take(t.Context()); err != nil {
		t.Fatal(err)
	}
	// Somebody else writes the key — a second controller that read it as expired on a skewed clock,
	// or a person — so a's next renewal finds a revision it did not write.
	kv, err := js.KeyValue(t.Context(), bucket)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := kv.Put(t.Context(), Key, []byte(`{"instance":"intruder"}`)); err != nil {
		t.Fatal(err)
	}
	if err := a.Renew(t.Context()); err == nil {
		t.Fatal("a renewed over a key somebody else wrote")
	}
	select {
	case <-a.Lost():
	case <-time.After(time.Second):
		t.Fatal("a was not told it lost the lease")
	}
	if _, err := a.Epoch(); !errors.Is(err, ErrNotHeld) {
		t.Fatalf("a acts after its renewal was refused: %v", err)
	}
	// And giving back a lease it lost deletes nothing of the one who holds it now.
	a.Release(t.Context())
	if h, found, err := Current(t.Context(), kv); err != nil || !found || h.Instance != "intruder" {
		t.Fatalf("after a gave back what it lost, the key holds %+v (%v, %v)", h, found, err)
	}
}

func TestKeepingRenewsAcrossManyAges(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	epoch, err := a.Take(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	keeping, stop := context.WithCancel(t.Context())
	defer stop()
	go a.Keep(keeping)
	until := time.Now().Add(3 * testTTL)
	for time.Now().Before(until) {
		if e, err := a.Epoch(); err != nil || e != epoch {
			t.Fatalf("a, renewing, answers %d, %v", e, err)
		}
		time.Sleep(200 * time.Millisecond)
	}
	h, found, err := a.Current(t.Context())
	if err != nil || !found || h.Epoch != epoch || h.Instance != "a" {
		t.Fatalf("the key says %+v (%v, %v)", h, found, err)
	}
}

func TestAnEpochIsNeverIssuedTwiceWhenTheBucketStartsOver(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	// The mesh has issued epoch 500 before: its store says so. The bucket is new — a bus whose data
	// was replaced — and its revisions start at one.
	floor := func(context.Context) (uint64, error) { return 500, nil }
	a := aController(t, bucket, "a", floor)
	moved := false
	a.o.Moved = func(was, floor uint64) { moved = was < floor && floor == 500 }
	epoch, err := a.Take(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if epoch <= 500 {
		t.Fatalf("a took epoch %d with 500 already issued: every machine that heard 500 would refuse it", epoch)
	}
	if !moved {
		t.Fatal("the bucket's revisions were moved past the floor and nobody was told")
	}
	// A floor that cannot be read takes nothing: an epoch that may not be higher is not issued.
	_, other := aBucket(t)
	b := aController(t, other, "b", func(context.Context) (uint64, error) { return 0, errors.New("store away") })
	if _, err := b.TryTake(t.Context()); err == nil {
		t.Fatal("b took a lease without knowing the highest epoch issued")
	}
}

func TestALeaseBucketWithoutAnAgeIsRefused(t *testing.T) {
	t.Parallel()
	url := testbus.URL(t)
	conn, err := nats.Connect(url)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	js, _ := jetstream.New(conn)
	bucket := fmt.Sprintf("lease-ageless-%d", time.Now().UnixNano())
	if _, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: bucket, History: 1}); err != nil {
		t.Fatal(err)
	}
	defer func() { _ = js.DeleteKeyValue(context.Background(), bucket) }()
	if _, err := Open(t.Context(), js, bucket, Options{Holder: Holder{Instance: "a"}}); err == nil {
		t.Fatal("a lease over keys that never expire was opened: a controller that died holding it would hold it for ever")
	}
}

// A bus that stalls (novox/hq issue 477): a renewal it does not answer in time is not a lost lease. The
// holder keeps the lease until the key's age passes with no renewal made, starts no new act once half of
// it has, acts again on the next renewal made, and loses it at once only to another holder.

// stallable puts l's renewals on a bus that stalls while the answer is true, as one whose JetStream API
// does: each renewal waits out its bound and fails with the deadline. landed is a stall that swallows
// only the answer: the write is made. Called before l is kept.
func stallable(l *Lease, landed bool) *atomic.Bool {
	real := l.update
	stalled := new(atomic.Bool)
	l.update = func(ctx context.Context, key string, value []byte, last uint64) (uint64, error) {
		if !stalled.Load() {
			return real(ctx, key, value, last)
		}
		if landed {
			_, _ = real(ctx, key, value, last)
		}
		<-ctx.Done()
		return 0, ctx.Err()
	}
	return stalled
}

// keeping takes the lease for l and keeps it until the test ends, answering the epoch.
func keeping(t *testing.T, l *Lease) uint64 {
	t.Helper()
	epoch, err := l.Take(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	go l.Keep(t.Context())
	return epoch
}

// delays is told of every renewal that failed without losing l's lease.
func delays(l *Lease) <-chan error {
	told := make(chan error, 64)
	l.o.Delayed = func(why error, _, _ time.Time) {
		select {
		case told <- why:
		default:
		}
	}
	return told
}

func TestAFailedRenewalKeepsTheLeaseUntilItsTTL(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	told := delays(a)
	stalled := stallable(a, false)
	epoch := keeping(t, a)
	stalled.Store(true)

	select {
	case why := <-told:
		if !errors.Is(why, context.DeadlineExceeded) {
			t.Fatalf("a renewal the bus did not answer was told as %v", why)
		}
	case <-time.After(3 * testRenew):
		t.Fatal("a renewal the bus did not answer was told to nobody")
	}
	renewed := a.Renewed()
	const slack = 150 * time.Millisecond
	// Within the key's age the lease is kept, though no renewal is made: no other controller could take it.
	for time.Now().Before(renewed.Add(testTTL - slack)) {
		select {
		case <-a.Lost():
			t.Fatalf("a lost the lease %s after its last renewal, within the key's age of %s: %v",
				time.Since(renewed), testTTL, a.LostWhy())
		default:
		}
		since := time.Since(renewed)
		e, err := a.Epoch()
		switch {
		case since < testTTL/2-slack && (err != nil || e != epoch):
			t.Fatalf("a starts no act %s after its last renewal, within half the key's age: %d, %v", since, e, err)
		case since > testTTL/2+slack && !errors.Is(err, ErrNotHeld):
			t.Fatalf("a still starts acts %s after its last renewal, past half the key's age: %d, %v", since, e, err)
		}
		time.Sleep(20 * time.Millisecond)
	}
	// Past it, the lease is lost: the key may be somebody else's.
	select {
	case <-a.Lost():
	case <-time.After(2 * slack):
		t.Fatalf("a still holds the lease %s after its last renewal, past the key's age of %s",
			time.Since(renewed), testTTL)
	}
	if why := a.LostWhy(); why == nil || !strings.Contains(why.Error(), "no renewal was made within the key's age") {
		t.Fatalf("a lost the lease for %v", why)
	}
	if _, err := a.Epoch(); !errors.Is(err, ErrNotHeld) {
		t.Fatalf("a acts after losing the lease: %v", err)
	}
}

func TestARenewalMadeAfterAStallOpensTheGateAgain(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	bus := stallable(a, false)
	epoch := keeping(t, a)
	bus.Store(true)

	stalled := time.Now()
	for a.Held() {
		if time.Since(stalled) > testTTL {
			t.Fatal("a still starts acts a whole key's age into a stall")
		}
		time.Sleep(20 * time.Millisecond)
	}
	bus.Store(false)
	for !a.Held() {
		select {
		case <-a.Lost():
			t.Fatalf("a lost the lease to a stall shorter than the key's age: %v", a.LostWhy())
		default:
		}
		if time.Since(stalled) > 2*testTTL {
			t.Fatal("a never acted again after the bus answered")
		}
		time.Sleep(20 * time.Millisecond)
	}
	if e, err := a.Epoch(); err != nil || e != epoch {
		t.Fatalf("a acts again as %d, %v; it held epoch %d", e, err, epoch)
	}
	// And keeps it, renewing, across more than one age of the key.
	time.Sleep(testTTL + testRenew)
	if !a.Held() {
		t.Fatalf("a stopped acting after the stall ended: %v", a.LostWhy())
	}
}

func TestARenewalWrittenThoughItsAnswerWasLostIsStillTheHolders(t *testing.T) {
	t.Parallel()
	_, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	bus := stallable(a, true)
	epoch := keeping(t, a)
	// The bus writes each renewal and its answer never arrives: the key moves past the revision a
	// knows, and is still a's.
	bus.Store(true)
	stalled := time.Now()
	for a.Held() {
		if time.Since(stalled) > testTTL {
			t.Fatal("a still starts acts a whole key's age into a stall")
		}
		time.Sleep(20 * time.Millisecond)
	}
	bus.Store(false)
	deadline := time.Now().Add(2 * testRenew)
	for !a.Held() {
		select {
		case <-a.Lost():
			t.Fatalf("a lost the lease to its own renewal: %v", a.LostWhy())
		default:
		}
		if time.Now().After(deadline) {
			t.Fatal("a never renewed over the renewal the stall wrote")
		}
		time.Sleep(20 * time.Millisecond)
	}
	h, found, err := a.Current(t.Context())
	if err != nil || !found || h.Instance != "a" || h.Epoch != epoch {
		t.Fatalf("the key says %+v (%v, %v), want a at epoch %d", h, found, err, epoch)
	}
}

func TestAnotherHolderSeenAfterAStallIsALossAtOnce(t *testing.T) {
	t.Parallel()
	js, bucket := aBucket(t)
	a := aController(t, bucket, "a", nil)
	bus := stallable(a, false)
	keeping(t, a)
	bus.Store(true)
	time.Sleep(testRenew + testRenew/2)
	// While a's renewals go unanswered, somebody else writes the key.
	kv, err := js.KeyValue(t.Context(), bucket)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := kv.Put(t.Context(), Key, []byte(`{"instance":"b","epoch":1}`)); err != nil {
		t.Fatal(err)
	}
	bus.Store(false)
	renewed := a.Renewed()
	select {
	case <-a.Lost():
		if lostAt := time.Since(renewed); lostAt >= testTTL {
			t.Fatalf("a was told it lost the lease only %s after its last renewal: another holder is a loss at once", lostAt)
		}
	case <-time.After(testTTL):
		t.Fatal("a was not told another controller holds the lease")
	}
	if why := a.LostWhy(); why == nil || !strings.Contains(why.Error(), "another holder") {
		t.Fatalf("a lost the lease for %v", why)
	}
}