From f84e8ffa27d491fc136af60f6c0d98bde9406a31 Mon Sep 17 00:00:00 2001 From: Evan Jarrett Date: Tue, 11 Aug 2026 21:17:55 -0500 Subject: [PATCH] db: create instance_leases and add the lease manager Groundwork for running more than one AppView instance. Nothing is wired to this yet; the next commit moves the background workers onto it. Several workers must run on exactly one instance. The Jetstream consumer is the sharpest case: 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, so two consumers would each hold a partial view of the holds and each write its partial sum as 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; only the holder runs the worker. Acquire is a single INSERT ... ON CONFLICT ... WHERE, so two instances racing for the same expired lease cannot both win: the loser's update matches no rows. The fence token increments on every change of custody, so a process that stalled past its TTL discovers on its next renewal that it was superseded, rather than continuing to act as the holder. A renewal blackout is treated as a loss. If the database has been unreachable for longer than the TTL, 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. Clean shutdown expires the lease in place rather than deleting the row, so a replacement starts in seconds instead of waiting out the TTL, while the fence token survives to keep a stalled former holder from matching again. Timestamps 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 stored expiry and a literal would compare as strings that sort differently. That comparison is the whole safety property, so it does not get to be subtle. The cost is a dependency on roughly-synced clocks, the same assumption Kubernetes leases make; keep the TTL well above any plausible skew. The lease tests are file-backed rather than :memory:. go-libsql gives every connection to an in-memory DSN its own private database, so with MaxOpenConns of 8 a second goroutine lands on a connection where the schema was never applied ("no such table"). Every existing test in the package is sequential and reuses one pooled connection, which is why this has stayed invisible. Co-Authored-By: Claude Opus 5 (1M context) --- pkg/appview/db/leases.go | 205 ++++++++++ pkg/appview/db/leases_test.go | 356 ++++++++++++++++++ .../0030_create_instance_leases.yaml | 31 ++ pkg/appview/db/schema.sql | 18 + pkg/appview/leases/manager.go | 252 +++++++++++++ pkg/appview/leases/manager_test.go | 285 ++++++++++++++ 6 files changed, 1147 insertions(+) create mode 100644 pkg/appview/db/leases.go create mode 100644 pkg/appview/db/leases_test.go create mode 100644 pkg/appview/db/migrations/0030_create_instance_leases.yaml create mode 100644 pkg/appview/leases/manager.go create mode 100644 pkg/appview/leases/manager_test.go 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") + } +}