at-container-registry/pkg/appview/db/leases_test.go

package db

import (
	"database/sql"
	"errors"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

const testTTL = 60 * time.Second

// baseTime is fixed so the tests exercise expiry by advancing a clock rather
// than by sleeping. TryAcquireLease takes `now` as a parameter precisely so
// this is possible.
var baseTime = time.Date(2026, 3, 1, 12, 0, 0, 0, time.UTC)

// leaseTestDB returns a file-backed database rather than an in-memory one.
//
// go-libsql gives every *connection* to ":memory:" its own private database, so
// with the pool's MaxOpenConns of 8 a second goroutine can land on a connection
// where the schema was never applied ("no such table: instance_leases"). The
// sequential tests in this package never notice because they reuse one pooled
// connection, but the concurrency test below fails immediately on it. A temp
// file is also closer to production, where the lease table is shared precisely
// because the database is shared.
func leaseTestDB(t *testing.T) *sql.DB {
	t.Helper()
	db, err := InitDB(filepath.Join(t.TempDir(), "leases.db"), LibsqlConfig{})
	if err != nil {
		t.Fatalf("InitDB: %v", err)
	}
	t.Cleanup(func() { db.Close() })
	return db
}

func TestLeaseAcquiredWhenUnheld(t *testing.T) {
	db := leaseTestDB(t)

	fence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil {
		t.Fatalf("TryAcquireLease: %v", err)
	}
	if !ok {
		t.Fatal("expected to acquire an unheld lease")
	}
	if fence != 1 {
		t.Errorf("expected fence 1 on first acquisition, got %d", fence)
	}
}

// TestLeaseDeniedWhileHeld is the property the whole table exists for: a second
// instance must not be able to start a worker the first one is already running.
func TestLeaseDeniedWhileHeld(t *testing.T) {
	db := leaseTestDB(t)

	if _, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL); err != nil || !ok {
		t.Fatalf("first acquire: ok=%v err=%v", ok, err)
	}

	// Well within the TTL.
	_, ok, err := TryAcquireLease(db, "worker", "instance-b", baseTime.Add(10*time.Second), testTTL)
	if err != nil {
		t.Fatalf("second acquire: %v", err)
	}
	if ok {
		t.Error("instance-b acquired a lease that instance-a still holds")
	}

	lease, err := GetLease(db, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease.HolderID != "instance-a" {
		t.Errorf("holder changed to %q despite the failed acquire", lease.HolderID)
	}
}

// TestLeaseStolenAfterExpiry covers the failover case: the holder died without
// releasing, so another instance must be able to take over once the TTL lapses.
func TestLeaseStolenAfterExpiry(t *testing.T) {
	db := leaseTestDB(t)

	firstFence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("first acquire: ok=%v err=%v", ok, err)
	}

	afterExpiry := baseTime.Add(testTTL + time.Second)
	secondFence, ok, err := TryAcquireLease(db, "worker", "instance-b", afterExpiry, testTTL)
	if err != nil {
		t.Fatalf("steal: %v", err)
	}
	if !ok {
		t.Fatal("expected instance-b to take over an expired lease")
	}
	if secondFence <= firstFence {
		t.Errorf("fence must advance on change of custody: was %d, now %d", firstFence, secondFence)
	}
}

func TestLeaseReacquiredBySameHolder(t *testing.T) {
	db := leaseTestDB(t)

	if _, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL); err != nil || !ok {
		t.Fatalf("first acquire: ok=%v err=%v", ok, err)
	}
	// The holder re-acquiring is not a conflict; it extends its own lease.
	_, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime.Add(5*time.Second), testTTL)
	if err != nil {
		t.Fatalf("re-acquire: %v", err)
	}
	if !ok {
		t.Error("a holder must be able to re-acquire its own lease")
	}
}

func TestLeaseRenewalExtendsExpiry(t *testing.T) {
	db := leaseTestDB(t)

	fence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("acquire: ok=%v err=%v", ok, err)
	}

	renewAt := baseTime.Add(20 * time.Second)
	renewed, err := RenewLease(db, "worker", "instance-a", fence, renewAt, testTTL)
	if err != nil {
		t.Fatalf("RenewLease: %v", err)
	}
	if !renewed {
		t.Fatal("expected renewal to succeed for the current holder")
	}

	lease, err := GetLease(db, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	want := renewAt.Add(testTTL)
	if !lease.ExpiresAt.Equal(want) {
		t.Errorf("expiry not extended: got %v, want %v", lease.ExpiresAt, want)
	}
}

// TestLeaseRenewalFailsAfterTakeover is the fencing property. A process that
// stalled past its TTL — a long GC pause, a hung network call — must discover it
// is no longer the holder rather than carrying on. Without this, two instances
// briefly run the same worker, which is exactly the state the lease exists to
// prevent.
func TestLeaseRenewalFailsAfterTakeover(t *testing.T) {
	db := leaseTestDB(t)

	staleFence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("acquire: ok=%v err=%v", ok, err)
	}

	// instance-a stalls; instance-b takes over after the TTL lapses.
	if _, ok, err := TryAcquireLease(db, "worker", "instance-b", baseTime.Add(testTTL+time.Second), testTTL); err != nil || !ok {
		t.Fatalf("takeover: ok=%v err=%v", ok, err)
	}

	// instance-a wakes up and tries to renew with the fence it remembers.
	renewed, err := RenewLease(db, "worker", "instance-a", staleFence, baseTime.Add(testTTL+2*time.Second), testTTL)
	if err != nil {
		t.Fatalf("RenewLease: %v", err)
	}
	if renewed {
		t.Error("a superseded holder renewed its lease; fencing is not working")
	}
}

func TestLeaseRenewalFailsWithWrongFence(t *testing.T) {
	db := leaseTestDB(t)

	fence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("acquire: ok=%v err=%v", ok, err)
	}

	renewed, err := RenewLease(db, "worker", "instance-a", fence+99, baseTime.Add(time.Second), testTTL)
	if err != nil {
		t.Fatalf("RenewLease: %v", err)
	}
	if renewed {
		t.Error("renewal succeeded with a fence token that was never issued")
	}
}

// TestLeaseReleaseAllowsImmediateTakeover is what makes a rolling deploy fast:
// without it the replacement instance waits out the full TTL before starting the
// worker.
func TestLeaseReleaseAllowsImmediateTakeover(t *testing.T) {
	db := leaseTestDB(t)

	fence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("acquire: ok=%v err=%v", ok, err)
	}
	if err := ReleaseLease(db, "worker", "instance-a", fence); err != nil {
		t.Fatalf("ReleaseLease: %v", err)
	}

	// One second later, nowhere near the TTL.
	newFence, ok, err := TryAcquireLease(db, "worker", "instance-b", baseTime.Add(time.Second), testTTL)
	if err != nil {
		t.Fatalf("acquire after release: %v", err)
	}
	if !ok {
		t.Fatal("expected instance-b to take a released lease immediately")
	}
	if newFence <= fence {
		t.Errorf("fence must still advance across a release: was %d, now %d", fence, newFence)
	}
}

// TestLeaseReleaseByNonHolderIsRejected guards against a stalled former holder
// expiring the lease out from under whoever holds it now.
func TestLeaseReleaseByNonHolderIsRejected(t *testing.T) {
	db := leaseTestDB(t)

	fence, ok, err := TryAcquireLease(db, "worker", "instance-a", baseTime, testTTL)
	if err != nil || !ok {
		t.Fatalf("acquire: ok=%v err=%v", ok, err)
	}

	err = ReleaseLease(db, "worker", "instance-b", fence)
	if !errors.Is(err, ErrLeaseNotHeld) {
		t.Errorf("expected ErrLeaseNotHeld, got %v", err)
	}

	lease, err := GetLease(db, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease.ExpiresAt.Equal(fromMillis(0)) {
		t.Error("a non-holder expired the lease")
	}
}

// TestLeasesAreIndependent: holding one lease must not affect another. The
// AppView runs four or more concurrently on the same instance.
func TestLeasesAreIndependent(t *testing.T) {
	db := leaseTestDB(t)

	if _, ok, err := TryAcquireLease(db, LeaseJetstream, "instance-a", baseTime, testTTL); err != nil || !ok {
		t.Fatalf("acquire jetstream: ok=%v err=%v", ok, err)
	}
	if _, ok, err := TryAcquireLease(db, LeaseBackfill, "instance-b", baseTime, testTTL); err != nil || !ok {
		t.Fatalf("instance-b must be able to hold a different lease: ok=%v err=%v", ok, err)
	}

	leases, err := ListLeases(db)
	if err != nil {
		t.Fatalf("ListLeases: %v", err)
	}
	if len(leases) != 2 {
		t.Fatalf("expected 2 leases, got %d", len(leases))
	}
}

func TestGetLeaseReturnsNilWhenNeverTaken(t *testing.T) {
	db := leaseTestDB(t)

	lease, err := GetLease(db, "never-taken")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease != nil {
		t.Errorf("expected nil for an untaken lease, got %+v", lease)
	}
}

// TestLeaseContentionSingleWinner runs the acquire path repeatedly from several
// would-be holders against one expired lease. Exactly one may win each round.
func TestLeaseContentionSingleWinner(t *testing.T) {
	db := leaseTestDB(t)

	now := baseTime
	holders := []string{"a", "b", "c", "d"}

	for round := range 5 {
		winners := 0
		for _, h := range holders {
			_, ok, err := TryAcquireLease(db, "worker", h, now, testTTL)
			if err != nil {
				t.Fatalf("round %d, holder %s: %v", round, h, err)
			}
			if ok {
				winners++
			}
		}
		if winners != 1 {
			t.Fatalf("round %d: expected exactly 1 winner, got %d", round, winners)
		}
		// Advance past the TTL so the next round contends for an expired lease.
		now = now.Add(testTTL + time.Second)
	}
}

// TestLeaseConcurrentAcquireSingleWinner is the same property as
// TestLeaseContentionSingleWinner but with real goroutines racing through the
// database, which is what actually happens when several instances boot together
// behind a load balancer. It relies on SQLite evaluating the INSERT ... ON
// CONFLICT ... WHERE as a single atomic operation; if that assumption is ever
// wrong, two workers start at once and this test is what says so.
//
// Run with -race to also cover the driver's own concurrency.
func TestLeaseConcurrentAcquireSingleWinner(t *testing.T) {
	db := leaseTestDB(t)

	const contenders = 8
	var (
		wg      sync.WaitGroup
		mu      sync.Mutex
		winners []string
		errs    []error
	)

	start := make(chan struct{})
	for i := range contenders {
		holder := "instance-" + string(rune('a'+i))
		wg.Go(func() {
			<-start // release everyone at once to maximize overlap
			_, ok, err := TryAcquireLease(db, "worker", holder, baseTime, testTTL)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, err)
				return
			}
			if ok {
				winners = append(winners, holder)
			}
		})
	}
	close(start)
	wg.Wait()

	for _, err := range errs {
		t.Errorf("concurrent acquire returned an error: %v", err)
	}
	if len(winners) != 1 {
		t.Fatalf("expected exactly one winner across %d concurrent acquires, got %d: %v",
			contenders, len(winners), winners)
	}

	lease, err := GetLease(db, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease.HolderID != winners[0] {
		t.Errorf("stored holder %q does not match the winner %q", lease.HolderID, winners[0])
	}
}