at-container-registry/pkg/appview/leases/manager_test.go

package leases

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

	"atcr.io/pkg/appview/db"
)

// testDB returns a file-backed database. It must not be ":memory:": go-libsql
// gives each connection to an in-memory DSN its own private database, so a
// second pooled connection would not see the instance_leases table.
func testDB(t *testing.T) *sql.DB {
	t.Helper()
	database, err := db.InitDB(filepath.Join(t.TempDir(), "leases.db"), db.LibsqlConfig{})
	if err != nil {
		t.Fatalf("InitDB: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	return database
}

// fastConfig keeps the lease timings short so tests finish quickly while
// preserving the renew < TTL relationship the manager depends on.
func fastConfig() Config {
	return Config{Enabled: true, TTL: 300 * time.Millisecond, RenewInterval: 50 * time.Millisecond}
}

func TestDisabledManagerRunsWorkerDirectly(t *testing.T) {
	database := testDB(t)
	m := NewManager(database, Config{Enabled: false})

	var ran atomic.Bool
	m.Run(context.Background(), "worker", func(context.Context) error {
		ran.Store(true)
		return nil
	})

	if !ran.Load() {
		t.Error("worker did not run with leader election disabled")
	}
	// Nothing should have touched the lease table.
	lease, err := db.GetLease(database, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease != nil {
		t.Errorf("disabled manager wrote a lease row: %+v", lease)
	}
}

// TestOnlyOneManagerRunsTheWorker is the property the package exists for.
func TestOnlyOneManagerRunsTheWorker(t *testing.T) {
	database := testDB(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var running atomic.Int32
	var maxConcurrent atomic.Int32
	started := make(chan struct{}, 3)

	worker := func(workerCtx context.Context) error {
		n := running.Add(1)
		for {
			m := maxConcurrent.Load()
			if n <= m || maxConcurrent.CompareAndSwap(m, n) {
				break
			}
		}
		defer running.Add(-1)

		started <- struct{}{}
		<-workerCtx.Done()
		return workerCtx.Err()
	}

	done := make(chan struct{}, 3)
	for range 3 {
		m := NewManager(database, fastConfig())
		go func() {
			defer func() { done <- struct{}{} }()
			m.Run(ctx, "worker", worker)
		}()
	}

	// Wait for the leader to start.
	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("no manager acquired the lease")
	}

	// Give the followers ample opportunity to wrongly start too.
	time.Sleep(500 * time.Millisecond)

	if got := maxConcurrent.Load(); got != 1 {
		t.Errorf("expected exactly 1 concurrent worker, saw %d", got)
	}

	cancel()
	for range 3 {
		select {
		case <-done:
		case <-time.After(5 * time.Second):
			t.Fatal("manager did not exit after context cancellation")
		}
	}
}

// TestWorkerContextCancelledWhenLeaseLost covers the fencing path from the
// worker's point of view: if another instance takes the lease, the worker's
// context must be cancelled so it stops before the new holder starts.
func TestWorkerContextCancelledWhenLeaseLost(t *testing.T) {
	database := testDB(t)

	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()

	m := NewManager(database, fastConfig())

	started := make(chan struct{})
	cancelled := make(chan struct{})
	go m.Run(ctx, "worker", func(workerCtx context.Context) error {
		close(started)
		<-workerCtx.Done()
		close(cancelled)
		// Return a non-nil error so Run retries rather than treating this as
		// completed work, matching how a real long-running worker behaves.
		return workerCtx.Err()
	})

	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("worker never started")
	}

	// Simulate the holder stalling long enough for its lease to lapse. Sleeping
	// would not do it: the manager is renewing on a timer, so the lease never
	// expires while it is healthy, which is the point. Expiring the row directly
	// reproduces the state another instance would actually observe after a stall
	// (a GC pause, a hung syscall, a partitioned database).
	if _, err := database.Exec(`UPDATE instance_leases SET expires_at = 0 WHERE lease_name = ?`, "worker"); err != nil {
		t.Fatalf("expire the lease: %v", err)
	}
	if _, ok, err := db.TryAcquireLease(database, "worker", "some-other-instance", time.Now(), time.Minute); err != nil || !ok {
		t.Fatalf("could not steal the lapsed lease: ok=%v err=%v", ok, err)
	}

	select {
	case <-cancelled:
	case <-time.After(5 * time.Second):
		t.Fatal("worker context was not cancelled after the lease was taken over")
	}
}

// TestReleaseOnExitAllowsImmediateTakeover: a clean shutdown must hand the lease
// over in seconds, not after a full TTL. This is what keeps a rolling deploy
// from pausing indexing for a minute.
func TestReleaseOnExitAllowsImmediateTakeover(t *testing.T) {
	database := testDB(t)

	m := NewManager(database, fastConfig())
	ctx, cancel := context.WithCancel(context.Background())

	started := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		m.Run(ctx, "worker", func(workerCtx context.Context) error {
			close(started)
			<-workerCtx.Done()
			return workerCtx.Err()
		})
	}()

	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("worker never started")
	}

	cancel()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("manager did not exit")
	}

	// Immediately, well inside the TTL, another instance should be able to take
	// the lease because the departing one released it.
	if _, ok, err := db.TryAcquireLease(database, "worker", "replacement", time.Now(), time.Minute); err != nil || !ok {
		t.Errorf("replacement could not take a released lease: ok=%v err=%v", ok, err)
	}
}

// TestWaitBlocksUntilLeaseReleased pins the shutdown ordering.
//
// Regression test for a real bug: the workers were started with `go m.Run(...)`
// and the process exited immediately after cancelling their context, so it was
// gone before the release landed. The lease then survived until its TTL lapsed
// and the replacement instance sat idle for up to a minute. Two instances under
// a rolling deploy showed it; no unit test did.
func TestWaitBlocksUntilLeaseReleased(t *testing.T) {
	database := testDB(t)

	m := NewManager(database, fastConfig())
	ctx, cancel := context.WithCancel(context.Background())

	started := make(chan struct{})
	m.Go(ctx, "worker", func(workerCtx context.Context) error {
		close(started)
		<-workerCtx.Done()
		return workerCtx.Err()
	})

	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("worker never started")
	}

	cancel()
	m.Wait(5 * time.Second)

	// Wait has returned, so the release must already have happened. Any expiry
	// still in the future would mean the replacement has to wait it out.
	lease, err := db.GetLease(database, "worker")
	if err != nil {
		t.Fatalf("GetLease: %v", err)
	}
	if lease == nil {
		t.Fatal("lease row disappeared")
	}
	if lease.ExpiresAt.After(time.Now()) {
		t.Errorf("Wait returned before the lease was released; it still expires at %v", lease.ExpiresAt)
	}
}

// TestWaitGivesUpOnAHungWorker: shutdown must not hang forever because one
// worker ignores its context. The lease expires on its own in that case.
func TestWaitGivesUpOnAHungWorker(t *testing.T) {
	database := testDB(t)

	m := NewManager(database, fastConfig())
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	release := make(chan struct{})
	t.Cleanup(func() { close(release) })

	started := make(chan struct{})
	m.Go(ctx, "hung", func(context.Context) error {
		close(started)
		<-release // deliberately ignores the worker context
		return nil
	})

	select {
	case <-started:
	case <-time.After(5 * time.Second):
		t.Fatal("worker never started")
	}

	cancel()
	done := make(chan struct{})
	go func() {
		m.Wait(200 * time.Millisecond)
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Wait did not give up on a worker that ignores its context")
	}
}

// TestCompletedWorkReleasesAndStops: a one-shot worker returning nil means the
// job is done. Run must not loop and re-run it forever.
func TestCompletedWorkReleasesAndStops(t *testing.T) {
	database := testDB(t)

	m := NewManager(database, fastConfig())

	var runs atomic.Int32
	done := make(chan struct{})
	go func() {
		defer close(done)
		m.Run(context.Background(), "one-shot", func(context.Context) error {
			runs.Add(1)
			return nil
		})
	}()

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Run did not return after the worker completed")
	}

	if got := runs.Load(); got != 1 {
		t.Errorf("expected the worker to run once, ran %d times", got)
	}
}

// TestFailedWorkerRetries: a worker that errors should be retried, since the
// failure may be transient (a dropped Jetstream connection, say).
func TestFailedWorkerRetries(t *testing.T) {
	database := testDB(t)

	m := NewManager(database, fastConfig())
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()

	var runs atomic.Int32
	secondRun := make(chan struct{})
	go m.Run(ctx, "flaky", func(context.Context) error {
		if runs.Add(1) == 2 {
			close(secondRun)
		}
		return errors.New("transient failure")
	})

	select {
	case <-secondRun:
	case <-time.After(5 * time.Second):
		t.Fatalf("worker was not retried after failing, ran %d times", runs.Load())
	}
}

func TestNewManagerNormalizesTimings(t *testing.T) {
	database := testDB(t)

	// A renew interval at or above the TTL would let the lease expire between
	// renewals, so the holder would repeatedly lose the worker to itself.
	m := NewManager(database, Config{Enabled: true, TTL: 30 * time.Second, RenewInterval: 45 * time.Second})
	if m.cfg.RenewInterval >= m.cfg.TTL {
		t.Errorf("renew interval %v must be under the TTL %v", m.cfg.RenewInterval, m.cfg.TTL)
	}

	// Zero values fall back to defaults rather than producing a manager that
	// renews in a tight loop or never at all.
	d := NewManager(database, Config{Enabled: true})
	if d.cfg.TTL != DefaultTTL || d.cfg.RenewInterval != DefaultRenewInterval {
		t.Errorf("expected defaults, got ttl=%v renew=%v", d.cfg.TTL, d.cfg.RenewInterval)
	}
}

func TestHolderIDsAreDistinct(t *testing.T) {
	database := testDB(t)
	a := NewManager(database, fastConfig())
	b := NewManager(database, fastConfig())

	if a.HolderID() == b.HolderID() {
		t.Errorf("two managers share a holder ID (%q); leases could not tell them apart", a.HolderID())
	}
	if a.HolderID() == "" {
		t.Error("holder ID is empty")
	}
}