diff --git a/pkg/appview/db/leases.go b/pkg/appview/db/leases.go new file mode 100644 index 0000000..c75d092 --- /dev/null +++ b/pkg/appview/db/leases.go @@ -0,0 +1,205 @@ +package db + +import ( + "database/sql" + "errors" + "fmt" + "time" +) + +// Lease names used by the AppView. Each names exactly one background worker +// that must not run on more than one instance at a time. +const ( + LeaseJetstream = "jetstream" + LeaseBackfill = "backfill" + LeaseCleanup = "cleanup" + LeaseBillingTiers = "billing-tiers" +) + +// LeaseLabeler returns the lease name for a labeler subscriber. Each labeler has +// its own cursor in labeler_cursor keyed by src, so each gets its own lease. +func LeaseLabeler(did string) string { return "labeler:" + did } + +// Lease describes the current state of a named lease. +type Lease struct { + Name string + HolderID string + Fence int64 + AcquiredAt time.Time + ExpiresAt time.Time +} + +// ErrLeaseNotHeld is returned by ReleaseLease when the caller is not the current +// holder. It is not usually an error worth acting on: it means the lease already +// moved on, which is exactly what release was trying to achieve. +var ErrLeaseNotHeld = errors.New("db: lease not held by this holder") + +// Timestamps are stored as Unix milliseconds rather than as TIMESTAMP text. +// +// Two reasons. First, libSQL normalizes date-like TEXT values on the way in, and +// Go's driver and SQLite's own CURRENT_TIMESTAMP do not agree on a format +// ("2006-01-02T15:04:05Z" versus "2006-01-02 15:04:05"), so a stored expiry and +// a literal would be compared as strings that sort differently. Second, integer +// comparison in the acquire statement has no ambiguity at all, which matters +// because that comparison is the entire safety property. +// +// The consequence is that leases depend on roughly-synchronized clocks across +// instances, the same assumption Kubernetes leases make. Keep the TTL +// comfortably larger than any plausible clock skew. +func toMillis(t time.Time) int64 { return t.UnixMilli() } +func fromMillis(ms int64) time.Time { return time.UnixMilli(ms).UTC() } + +// TryAcquireLease attempts to take or renew the named lease for holderID, +// extending it until now+ttl. It reports whether holderID holds the lease +// afterwards, along with the fence token of the current holder. +// +// The lease is taken when it does not exist, when it has expired, or when +// holderID already holds it. A live lease belonging to someone else is left +// untouched. +// +// Atomicity comes from the WHERE clause on the conflict arm: SQLite evaluates +// the whole statement as one operation, so two instances racing to steal the +// same expired lease cannot both succeed. The loser's update simply matches no +// rows. +// +// The fence token increments on every change of custody. Callers pass it back to +// RenewLease, which fails once someone else has taken over — that is what stops +// a process that stalled past its TTL from continuing to act as the holder after +// the lease has moved. +func TryAcquireLease(db DBTX, name, holderID string, now time.Time, ttl time.Duration) (fence int64, acquired bool, err error) { + nowMS := toMillis(now) + expiresMS := toMillis(now.Add(ttl)) + + res, err := db.Exec(` + INSERT INTO instance_leases (lease_name, holder_id, fence, acquired_at, expires_at) + VALUES (?, ?, 1, ?, ?) + ON CONFLICT(lease_name) DO UPDATE SET + holder_id = excluded.holder_id, + fence = instance_leases.fence + 1, + acquired_at = excluded.acquired_at, + expires_at = excluded.expires_at + WHERE instance_leases.expires_at < ? + OR instance_leases.holder_id = excluded.holder_id + `, name, holderID, nowMS, expiresMS, nowMS) + if err != nil { + return 0, false, fmt.Errorf("acquire lease %q: %w", name, err) + } + + // No rows changed means a live lease is held by someone else. + if n, err := res.RowsAffected(); err == nil && n == 0 { + return 0, false, nil + } + + // Read back the fence. This also catches the narrow case where another + // instance stole the lease between our write and this read, which requires + // clock skew larger than the TTL. + lease, err := GetLease(db, name) + if err != nil { + return 0, false, err + } + if lease == nil || lease.HolderID != holderID { + return 0, false, nil + } + return lease.Fence, true, nil +} + +// RenewLease extends a lease the caller already holds, and reports whether the +// renewal succeeded. It fails when the lease has been taken over, which the +// fence token detects even if the new holder has since let it expire. +// +// A false return means the caller must stop doing whatever the lease protects. +func RenewLease(db DBTX, name, holderID string, fence int64, now time.Time, ttl time.Duration) (bool, error) { + res, err := db.Exec(` + UPDATE instance_leases + SET expires_at = ? + WHERE lease_name = ? AND holder_id = ? AND fence = ? + `, toMillis(now.Add(ttl)), name, holderID, fence) + if err != nil { + return false, fmt.Errorf("renew lease %q: %w", name, err) + } + n, err := res.RowsAffected() + if err != nil { + return false, fmt.Errorf("renew lease %q: %w", name, err) + } + return n > 0, nil +} + +// ReleaseLease gives up a lease immediately by expiring it in place, so another +// instance can take over without waiting out the TTL. Called on clean shutdown, +// which is what makes a rolling deploy hand over in seconds instead of a minute. +// +// The row is expired rather than deleted so the fence token survives. Deleting +// it would restart fencing at 1 and let a stalled former holder's renewal match +// again. +func ReleaseLease(db DBTX, name, holderID string, fence int64) error { + res, err := db.Exec(` + UPDATE instance_leases + SET expires_at = 0 + WHERE lease_name = ? AND holder_id = ? AND fence = ? + `, name, holderID, fence) + if err != nil { + return fmt.Errorf("release lease %q: %w", name, err) + } + n, err := res.RowsAffected() + if err != nil { + return fmt.Errorf("release lease %q: %w", name, err) + } + if n == 0 { + return ErrLeaseNotHeld + } + return nil +} + +// GetLease returns the current state of a lease, or nil when it has never been +// taken. Intended for diagnostics and the admin UI; the acquire path does not +// depend on it for correctness. +func GetLease(db DBTX, name string) (*Lease, error) { + var ( + lease Lease + acquiredMS int64 + expiresMS int64 + ) + err := db.QueryRow(` + SELECT lease_name, holder_id, fence, acquired_at, expires_at + FROM instance_leases + WHERE lease_name = ? + `, name).Scan(&lease.Name, &lease.HolderID, &lease.Fence, &acquiredMS, &expiresMS) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("get lease %q: %w", name, err) + } + lease.AcquiredAt = fromMillis(acquiredMS) + lease.ExpiresAt = fromMillis(expiresMS) + return &lease, nil +} + +// ListLeases returns every lease, ordered by name, for diagnostics. +func ListLeases(db DBTX) ([]Lease, error) { + rows, err := db.Query(` + SELECT lease_name, holder_id, fence, acquired_at, expires_at + FROM instance_leases + ORDER BY lease_name + `) + if err != nil { + return nil, fmt.Errorf("list leases: %w", err) + } + defer rows.Close() + + var out []Lease + for rows.Next() { + var ( + lease Lease + acquiredMS int64 + expiresMS int64 + ) + if err := rows.Scan(&lease.Name, &lease.HolderID, &lease.Fence, &acquiredMS, &expiresMS); err != nil { + return nil, fmt.Errorf("scan lease: %w", err) + } + lease.AcquiredAt = fromMillis(acquiredMS) + lease.ExpiresAt = fromMillis(expiresMS) + out = append(out, lease) + } + return out, rows.Err() +} diff --git a/pkg/appview/db/leases_test.go b/pkg/appview/db/leases_test.go new file mode 100644 index 0000000..ee071d7 --- /dev/null +++ b/pkg/appview/db/leases_test.go @@ -0,0 +1,356 @@ +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]) + } +} diff --git a/pkg/appview/db/migrations/0030_create_instance_leases.yaml b/pkg/appview/db/migrations/0030_create_instance_leases.yaml new file mode 100644 index 0000000..2aba693 --- /dev/null +++ b/pkg/appview/db/migrations/0030_create_instance_leases.yaml @@ -0,0 +1,31 @@ +description: | + Leader election for AppView background workers. + + The Jetstream consumer, backfill, labeler subscriber and cleanup loops all + start unconditionally in every process today, which is fine for one instance + and actively harmful for two. The consumer is the clearest case: StatsCache is + per-process in-memory state whose aggregate is written to repository_stats as + an absolute value, so two consumers would each hold a partial view of the holds + and each write its partial sum as the whole truth, overwriting each other + forever. The webhook dispatcher hangs off the same processor, so a second + consumer also means duplicate deliveries. + + Instances contend for a named lease and only the holder runs the worker. The + fence token increments on every change of custody so a process that stalled + past its TTL can detect it was superseded rather than continuing to act as the + holder. + + Timestamps are Unix milliseconds rather than TIMESTAMP text. libSQL normalizes + date-like TEXT on the way in, and Go's driver and CURRENT_TIMESTAMP disagree on + format, so a stored expiry and a literal would be compared as strings that sort + differently. That comparison is the entire safety property of the acquire + statement, so it gets an integer. +query: | + CREATE TABLE IF NOT EXISTS instance_leases ( + lease_name TEXT PRIMARY KEY, + holder_id TEXT NOT NULL, + fence INTEGER NOT NULL DEFAULT 0, + acquired_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_instance_leases_expires ON instance_leases(expires_at); diff --git a/pkg/appview/db/schema.sql b/pkg/appview/db/schema.sql index 67301c0..6b3176d 100644 --- a/pkg/appview/db/schema.sql +++ b/pkg/appview/db/schema.sql @@ -332,3 +332,21 @@ CREATE TABLE IF NOT EXISTS labeler_cursor ( src TEXT PRIMARY KEY, cursor INTEGER NOT NULL ); + +-- Leader election for background workers. Exactly one AppView instance may run +-- the Jetstream consumer, backfill, labeler subscriber and cleanup loops; see +-- pkg/appview/db/leases.go for why, and pkg/appview/leases for the renewal loop. +-- +-- acquired_at and expires_at are Unix milliseconds, not TIMESTAMP text: libSQL +-- normalizes date-like TEXT on the way in, and Go's driver and CURRENT_TIMESTAMP +-- disagree on format, so a string comparison in the acquire statement would be +-- comparing incompatible shapes. Integer comparison is the whole safety property +-- here, so it does not get to be subtle. +CREATE TABLE IF NOT EXISTS instance_leases ( + lease_name TEXT PRIMARY KEY, + holder_id TEXT NOT NULL, + fence INTEGER NOT NULL DEFAULT 0, + acquired_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_instance_leases_expires ON instance_leases(expires_at); diff --git a/pkg/appview/leases/manager.go b/pkg/appview/leases/manager.go new file mode 100644 index 0000000..48e7f10 --- /dev/null +++ b/pkg/appview/leases/manager.go @@ -0,0 +1,252 @@ +// Package leases provides leader election for AppView background workers. +// +// Several workers must run on exactly one instance: the Jetstream consumer, the +// backfill worker, the labeler subscribers and the periodic cleanup loop. Before +// this package they started unconditionally in every process, which is correct +// for a single instance and wrong for two. +// +// The Jetstream consumer shows why most sharply. Its StatsCache is per-process +// in-memory state, and the aggregate it produces is written to repository_stats +// as an absolute value rather than an increment. Two consumers would each hold a +// partial view of the holds and each write its partial sum as though it were the +// whole truth, overwriting one another indefinitely. The webhook dispatcher +// hangs off the same processor, so a second consumer also means every webhook +// fires twice. +// +// Instances contend for a named lease in the database; only the holder runs the +// worker. The holder renews on a timer, and a holder that cannot renew stops +// working before its lease can be taken by anyone else. +package leases + +import ( + "context" + "crypto/rand" + "database/sql" + "encoding/hex" + "errors" + "fmt" + "log/slog" + "os" + "time" + + "atcr.io/pkg/appview/db" +) + +// Default timings. The TTL must comfortably exceed both the renew interval and +// any plausible clock skew between instances, since expiry is compared against +// each instance's own clock. +const ( + DefaultTTL = 60 * time.Second + DefaultRenewInterval = 20 * time.Second +) + +// Config controls lease behavior. The zero value is not useful; use +// NewManager, which fills in defaults. +type Config struct { + // Enabled turns leader election on. When false, Run executes the worker + // directly with no database interaction, preserving single-instance + // behavior for local development and tests. + Enabled bool + + // TTL is how long a lease survives without renewal before another instance + // may take it. + TTL time.Duration + + // RenewInterval is how often the holder extends its lease. It doubles as + // the retry interval for instances waiting to acquire. + RenewInterval time.Duration +} + +// Manager acquires leases and runs workers under them. +type Manager struct { + db *sql.DB + holderID string + cfg Config +} + +// NewManager returns a Manager identified by a holder ID unique to this process. +// Invalid or unset timings fall back to the defaults. +func NewManager(database *sql.DB, cfg Config) *Manager { + if cfg.TTL <= 0 { + cfg.TTL = DefaultTTL + } + if cfg.RenewInterval <= 0 { + cfg.RenewInterval = DefaultRenewInterval + } + // A renew interval at or above the TTL guarantees the lease expires between + // renewals, so the worker would be stolen from itself on a loop. + if cfg.RenewInterval >= cfg.TTL { + cfg.RenewInterval = cfg.TTL / 3 + } + return &Manager{ + db: database, + holderID: newHolderID(), + cfg: cfg, + } +} + +// HolderID returns this process's lease identity, for logging and diagnostics. +func (m *Manager) HolderID() string { return m.holderID } + +// Enabled reports whether leader election is active. +func (m *Manager) Enabled() bool { return m.cfg.Enabled } + +// newHolderID builds an identifier unique to this process. Hostname and PID +// make it legible when reading the table by hand; the random suffix keeps it +// unique when a container restarts fast enough to reuse both. +func newHolderID() string { + host, err := os.Hostname() + if err != nil || host == "" { + host = "unknown" + } + var suffix [4]byte + if _, err := rand.Read(suffix[:]); err != nil { + // A collision only matters if two live processes share a hostname AND + // a PID, which cannot happen on one host. Degrade rather than fail. + return fmt.Sprintf("%s/%d", host, os.Getpid()) + } + return fmt.Sprintf("%s/%d/%s", host, os.Getpid(), hex.EncodeToString(suffix[:])) +} + +// errNotAcquired signals that another instance holds the lease. Not an error +// worth logging at anything above debug: it is the expected outcome on every +// instance that is not the leader. +var errNotAcquired = errors.New("leases: not acquired") + +// Run executes fn under the named lease, blocking until ctx is cancelled or fn +// reports it has finished. +// +// fn receives a context that is cancelled when the lease is lost, so a worker +// that respects its context stops before another instance can start one. fn +// returning nil means the work is complete and Run returns; returning an error +// means it failed and Run will try again after RenewInterval. Long-running +// workers should return only when their context is done. +// +// When leader election is disabled, fn runs immediately with ctx unchanged. +func (m *Manager) Run(ctx context.Context, name string, fn func(context.Context) error) { + if !m.cfg.Enabled { + slog.Info("Leader election disabled, running worker directly", + "component", "leases", "lease", name) + if err := fn(ctx); err != nil && ctx.Err() == nil { + slog.Warn("Worker exited with error", "component", "leases", "lease", name, "error", err) + } + return + } + + slog.Info("Contending for lease", + "component", "leases", "lease", name, "holder", m.holderID, + "ttl", m.cfg.TTL, "renew_interval", m.cfg.RenewInterval) + + for { + if ctx.Err() != nil { + return + } + + err := m.runUnderLease(ctx, name, fn) + switch { + case ctx.Err() != nil: + return + case err == nil: + // fn reported the work as complete. + slog.Info("Worker finished", "component", "leases", "lease", name) + return + case errors.Is(err, errNotAcquired): + // Someone else is the leader. Expected on every follower. + default: + slog.Warn("Worker exited, will retry after the lease interval", + "component", "leases", "lease", name, "error", err) + } + + select { + case <-ctx.Done(): + return + case <-time.After(m.cfg.RenewInterval): + } + } +} + +// runUnderLease performs one acquire-run-release cycle. +func (m *Manager) runUnderLease(ctx context.Context, name string, fn func(context.Context) error) error { + fence, acquired, err := db.TryAcquireLease(m.db, name, m.holderID, time.Now(), m.cfg.TTL) + if err != nil { + return fmt.Errorf("acquire lease: %w", err) + } + if !acquired { + return errNotAcquired + } + + slog.Info("Lease acquired, starting worker", + "component", "leases", "lease", name, "holder", m.holderID, "fence", fence) + + leaseCtx, cancel := context.WithCancel(ctx) + defer cancel() + + renewStopped := make(chan struct{}) + go m.renew(leaseCtx, name, fence, cancel, renewStopped) + + fnErr := fn(leaseCtx) + + // Stop renewing before releasing, so the renewal loop cannot extend a lease + // we are about to hand over. + cancel() + <-renewStopped + + // Release so a replacement instance can start immediately instead of + // waiting out the TTL. Best effort: if we already lost the lease, there is + // nothing to release and nothing to report. + if err := db.ReleaseLease(m.db, name, m.holderID, fence); err != nil && !errors.Is(err, db.ErrLeaseNotHeld) { + slog.Warn("Failed to release lease", + "component", "leases", "lease", name, "error", err) + } else if err == nil { + slog.Info("Lease released", "component", "leases", "lease", name, "holder", m.holderID) + } + + return fnErr +} + +// renew extends the lease until the context is cancelled, and cancels the +// worker if the lease can no longer be proven held. +// +// Two ways that happens. The lease was taken over, which RenewLease reports +// directly via the fence token. Or the database has been unreachable for longer +// than the TTL, in which case another instance is entitled to steal the lease +// and we must assume it has, even though we cannot ask. Continuing to work in +// that state is the one outcome the lease exists to prevent, so a renewal +// blackout is treated as a loss. +func (m *Manager) renew(ctx context.Context, name string, fence int64, cancel context.CancelFunc, stopped chan<- struct{}) { + defer close(stopped) + + ticker := time.NewTicker(m.cfg.RenewInterval) + defer ticker.Stop() + + lastRenewed := time.Now() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + now := time.Now() + ok, err := db.RenewLease(m.db, name, m.holderID, fence, now, m.cfg.TTL) + switch { + case err != nil: + if blackout := now.Sub(lastRenewed); blackout >= m.cfg.TTL { + slog.Error("Could not renew lease for longer than its TTL, stopping worker", + "component", "leases", "lease", name, + "blackout", blackout.Round(time.Second), "ttl", m.cfg.TTL, "error", err) + cancel() + return + } + slog.Warn("Lease renewal failed, will retry", + "component", "leases", "lease", name, "error", err) + case !ok: + slog.Warn("Lease taken over by another instance, stopping worker", + "component", "leases", "lease", name, "holder", m.holderID, "fence", fence) + cancel() + return + default: + lastRenewed = now + } + } + } +} diff --git a/pkg/appview/leases/manager_test.go b/pkg/appview/leases/manager_test.go new file mode 100644 index 0000000..e1746c1 --- /dev/null +++ b/pkg/appview/leases/manager_test.go @@ -0,0 +1,285 @@ +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) + } +} + +// 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") + } +}