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

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()
}