llmgw/internal/control/leadership.go

package control

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"log/slog"
	"os"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// LeaseStore is the part of the store leadership needs.
//
// An interface so the controller does not take a dependency on a concrete
// database for one feature, and so a deployment with no store at all — a
// router-only run, or a single control plane on a static configuration — keeps
// working with nothing to configure.
type LeaseStore interface {
	AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (store.Lease, error)
	ReleaseLeadership(ctx context.Context, role, owner string) error
}

// controlRole is the one thing there is to be leader of: placing work.
const controlRole = "control-plane"

// Leadership decides whether this process is the one that places work.
//
// Two control planes against one store both reconcile: both see a model
// wanting two replicas and none running, and both place. The result is twice
// the replicas, each holding GPUs the other believes are free, and nothing
// crashes to say so.
//
// A standby is not idle — it serves its own routing and keeps its connection
// to the store — it simply does not place, evict or drain. That makes running
// a second one safe rather than clever, which is the property worth having
// first: the common failure is not a planned failover, it is somebody starting
// a second control plane by accident.
type Leadership struct {
	store LeaseStore
	log   *slog.Logger
	// owner identifies this process. Random rather than the hostname, so a
	// restart on the same machine cannot be fooled by its own stale lease.
	owner string
	ttl   time.Duration
	// renew is how often the holder pushes the expiry out. Well under ttl, so
	// one slow query does not cost the leadership.
	renew time.Duration

	leader atomic.Bool
	// stopped is set by StepDown: nothing takes the lease again after it.
	stopped atomic.Bool

	hookMu sync.Mutex
	onLead []func()
}

// OnLead is called each time this process becomes the leader, after
// standing by: what it read at start may have been changed since by the
// leader before it, and is read again.
func (l *Leadership) OnLead(f func()) {
	l.hookMu.Lock()
	l.onLead = append(l.onLead, f)
	l.hookMu.Unlock()
}

func (l *Leadership) becameLeader() {
	l.hookMu.Lock()
	hooks := slices.Clone(l.onLead)
	l.hookMu.Unlock()
	for _, f := range hooks {
		go f()
	}
}

// NewLeadership prepares leadership over a store.
//
// A nil store means there is nobody to coordinate with, and this process is
// the leader by definition — which is every single-control-plane deployment,
// and must stay the case without anyone opting in.
func NewLeadership(st LeaseStore, ttl time.Duration, log *slog.Logger) *Leadership {
	if log == nil {
		log = slog.Default()
	}
	if ttl <= 0 {
		// Long enough to survive a slow database or a brief network stall,
		// short enough that a standby takes over before anyone has finished
		// reading the alert.
		ttl = 15 * time.Second
	}
	l := &Leadership{
		store: st,
		log:   log,
		owner: newOwnerID(),
		ttl:   ttl,
		renew: ttl / 3,
	}
	if st == nil {
		l.leader.Store(true)
	}
	return l
}

// IsLeader reports whether this process should be placing work.
func (l *Leadership) IsLeader() bool { return l.leader.Load() }

// Owner is this process's id, for logs and for ctl.
func (l *Leadership) Owner() string { return l.owner }

// Run keeps the lease for as long as it can hold it, and keeps trying when it
// cannot.
//
// A standby polls at the same interval the leader renews at, so a lapsed lease
// is picked up within about a third of its life rather than at the end of it.
func (l *Leadership) Run(ctx context.Context) {
	if l.store == nil {
		return
	}
	t := time.NewTicker(l.renew)
	defer t.Stop()

	l.attempt(ctx)
	for {
		select {
		case <-ctx.Done():
			l.step(ctx)
			return
		case <-t.C:
			l.attempt(ctx)
		}
	}
}

// attempt takes or renews the lease, logging only transitions.
//
// Only transitions, because this runs every few seconds for the life of the
// process: a line per attempt would bury everything else, and the thing an
// operator needs to see is the moment leadership moved.
func (l *Leadership) attempt(ctx context.Context) {
	if l.stopped.Load() {
		return
	}
	lease, err := l.store.AcquireLeadership(ctx, controlRole, l.owner, l.ttl)
	switch {
	case err == nil && l.stopped.Load():
		// Stepped down while this was on its way to the store: give back
		// what it just took.
		l.leader.Store(true)
		l.step(ctx)
	case err == nil:
		if l.leader.CompareAndSwap(false, true) {
			l.log.Info("this control plane is now placing work", "owner", l.owner)
			l.becameLeader()
		}
	case errors.Is(err, store.ErrNotLeader):
		if l.leader.CompareAndSwap(true, false) {
			// Losing a lease that was held is worth a warning: it means this
			// process was too slow to renew, and whatever made it slow is
			// still true.
			l.log.Warn("lost leadership; standing by",
				"owner", l.owner, "now", lease.Owner)
		} else if l.leader.Load() == false && lease.Owner != "" {
			l.log.Debug("standing by", "leader", lease.Owner,
				"until", lease.ExpiresAt.Format(time.RFC3339))
		}
	default:
		// A store that cannot be reached is not evidence that somebody else
		// is leading. Keeping the claim is the safer of two bad options: two
		// leaders is worse than one, but no leader at all means nothing is
		// placed and nothing is evicted, and a brief database blip should not
		// stop a cluster from healing itself.
		l.log.Warn("could not reach the store to renew leadership; keeping it for now",
			"err", err)
	}
}

// StepDown gives the lease up now, for a process on its way out: called
// before the agents are let go, so the standby they move to is already the
// leader, and before the process exits, which would otherwise race the
// release in Run and leave the standby to wait out the lease. Run's own
// release after it is then nothing.
func (l *Leadership) StepDown() {
	if l.store == nil {
		return
	}
	l.stopped.Store(true)
	l.step(context.Background())
}

// step gives the lease up on the way out, so a standby takes over in seconds
// rather than waiting out the clock.
func (l *Leadership) step(ctx context.Context) {
	// Once: StepDown and Run's own way out both come here.
	if !l.leader.CompareAndSwap(true, false) {
		return
	}
	// A fresh context: the one that just ended is the reason this is running.
	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
	defer cancel()
	if err := l.store.ReleaseLeadership(ctx, controlRole, l.owner); err != nil {
		l.log.Warn("could not release leadership; a standby will wait it out", "err", err)
		return
	}
	l.log.Info("released leadership", "owner", l.owner)
}

// newOwnerID is random, not the hostname.
//
// A hostname would let a restarted process match its own stale lease and
// believe it had been leading all along — including through whatever gap made
// it restart, during which another process may have taken over and started
// placing.
func newOwnerID() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		// Only reached if the system has no entropy, where a pid and the time
		// is still better than a fixed string.
		return "ctl-" + hex.EncodeToString([]byte(time.Now().Format(time.RFC3339Nano)))[:16]
	}
	host, _ := os.Hostname()
	if host == "" {
		host = "ctl"
	}
	return host + "-" + hex.EncodeToString(b[:])
}

// ControlRole is the lease's name in the store, for readers that ask who
// holds it without competing.
const ControlRole = controlRole

// HolderName is the machine an owner id names: the hostname it was made
// from, without the random part.
func HolderName(owner string) string {
	if i := strings.LastIndexByte(owner, '-'); i > 0 {
		return owner[:i]
	}
	return owner
}