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

// 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"
	"sync"
	"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

	// wg tracks workers started via Go, so Wait can block on their release
	// during shutdown.
	wg sync.WaitGroup
}

// 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")

// Go starts a leased worker in the background and registers it with the
// manager, so Wait can block on it during shutdown.
//
// Prefer this over `go m.Run(...)`. The registration has to happen on the
// caller's goroutine: doing it inside the spawned one races with Wait, which
// could then return before the worker has even started, let alone released its
// lease.
func (m *Manager) Go(ctx context.Context, name string, fn func(context.Context) error) {
	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		m.Run(ctx, name, fn)
	}()
}

// Wait blocks until every worker started with Go has stopped and released its
// lease, or until timeout elapses.
//
// Call this during shutdown, after cancelling the worker context. Without it the
// process exits while the release is still in flight, so the lease survives
// until its TTL lapses and the replacement instance sits idle for up to a full
// TTL. That is the difference between a rolling deploy pausing indexing for a
// second and pausing it for a minute.
func (m *Manager) Wait(timeout time.Duration) {
	done := make(chan struct{})
	go func() {
		m.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(timeout):
		slog.Warn("Timed out waiting for leased workers to stop; their leases will expire on their own",
			"component", "leases", "timeout", timeout, "ttl", m.cfg.TTL)
	}
}

// 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
			}
		}
	}
}