llmgw/internal/control/facility.go

package control

import (
	"context"
	"fmt"
	"log/slog"
	"maps"
	"slices"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/facility"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// FacilityStore is where holds and operator labels are kept.
type FacilityStore interface {
	FacilityHolds(ctx context.Context) ([]store.FacilityHold, error)
	SetFacilityHold(ctx context.Context, h store.FacilityHold) error
	DeleteFacilityHold(ctx context.Context, domain string) error
	NodeLabels(ctx context.Context) (map[string]map[string]string, error)
	SetNodeLabel(ctx context.Context, node, key, value, by string) error
}

// Facility is the operator's side of facility-aware placement: the budgets
// from server.yaml, the domains held for maintenance, and the labels set on
// nodes through the control plane. It keeps them in the store when there is
// one, and puts them in force in the inventory, where every placement —
// the reconciler's, headroom's and preemption's dry runs — reads them.
//
// Without a store it still works, for the life of the process: a held rack
// is worth honouring even on a control plane that will forget it.
type Facility struct {
	inv *inventory.Store
	log *slog.Logger

	mu    sync.Mutex
	st    FacilityStore
	base  facility.Policy
	holds map[facility.Domain]facility.Hold
}

// NewFacility starts with nothing held and nothing capped.
func NewFacility(inv *inventory.Store, log *slog.Logger) *Facility {
	if log == nil {
		log = slog.Default()
	}
	return &Facility{inv: inv, log: log, holds: map[facility.Domain]facility.Hold{}}
}

// SetPolicy puts the budgets and estimates from configuration in force. Its
// holds are ignored: those are the operator's, kept here.
func (f *Facility) SetPolicy(p facility.Policy) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.base = p.Clone()
	f.applyLocked()
}

// Policy is what placement reads now.
func (f *Facility) Policy() facility.Policy { return f.inv.Facility() }

// UseStore keeps holds and labels in the store from now on, and puts in
// force what it already holds.
func (f *Facility) UseStore(ctx context.Context, st FacilityStore) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.st = st
	f.loadLocked(ctx)
}

// Reload reads the holds and labels in the store again, as a standby
// becoming the leader does.
func (f *Facility) Reload(ctx context.Context) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.st == nil {
		return
	}
	f.holds = map[facility.Domain]facility.Hold{}
	f.loadLocked(ctx)
}

func (f *Facility) loadLocked(ctx context.Context) {
	st := f.st
	if holds, err := st.FacilityHolds(ctx); err != nil {
		f.log.Error("could not read the domains held for maintenance; none is held until they are set again", "err", err)
	} else {
		for _, h := range holds {
			d, err := facility.ParseDomain(h.Domain)
			if err != nil {
				continue
			}
			f.holds[d] = facility.Hold{Domain: d, Reason: h.Reason, Drain: h.Drain, By: h.By, At: h.At}
		}
	}
	if labels, err := st.NodeLabels(ctx); err != nil {
		f.log.Error("could not read the labels set on nodes; nodes carry only their agents' until they are set again", "err", err)
	} else {
		for node, l := range labels {
			f.inv.SetOperatorLabels(node, l)
		}
	}
	f.applyLocked()
}

func (f *Facility) applyLocked() {
	p := f.base.Clone()
	p.Holds = maps.Clone(f.holds)
	f.inv.SetFacility(p)
}

// Holds are the domains held, oldest first.
func (f *Facility) Holds() []facility.Hold {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := slices.Collect(maps.Values(f.holds))
	slices.SortFunc(out, func(a, b facility.Hold) int {
		if c := a.At.Compare(b.At); c != 0 {
			return c
		}
		return strings.Compare(string(a.Domain), string(b.Domain))
	})
	return out
}

// Hold takes a domain out of placement, and with drain moves the work in it
// elsewhere: each of its nodes is drained, one replica at a time with the
// replacement up first, as `ctl drain` does. It answers the nodes drained.
func (f *Facility) Hold(ctx context.Context, h facility.Hold) ([]string, error) {
	if h.At.IsZero() {
		h.At = time.Now()
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.st != nil {
		if err := f.st.SetFacilityHold(ctx, store.FacilityHold{
			Domain: string(h.Domain), Reason: h.Reason, Drain: h.Drain, At: h.At, By: h.By,
		}); err != nil {
			return nil, fmt.Errorf("keeping the hold: %w", err)
		}
	}
	f.holds[h.Domain] = h
	f.applyLocked()
	var drained []string
	if h.Drain {
		for _, n := range f.inv.Nodes() {
			if slices.Contains(n.Domains(), h.Domain) && (n.State == inventory.Ready || n.State == inventory.Cordoned) {
				f.inv.SetState(n.ID, inventory.Draining)
				drained = append(drained, n.ID)
			}
		}
	}
	f.log.Info("domain held: nothing new is placed in it", "domain", h.Domain, "reason", h.Reason,
		"drain", h.Drain, "by", h.By, "drained", drained)
	return drained, nil
}

// Release puts a domain back into placement. Nodes a draining hold drained
// are made ready again, unless another hold still covers them; a node an
// operator drained by itself is not told apart from those, so it is readied
// too, which is what releasing its rack asks for.
func (f *Facility) Release(ctx context.Context, d facility.Domain) (facility.Hold, []string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	h, ok := f.holds[d]
	if !ok {
		return facility.Hold{}, nil, fmt.Errorf("%s is not held; `llmgw ctl facility` lists what is", d)
	}
	if f.st != nil {
		if err := f.st.DeleteFacilityHold(ctx, string(d)); err != nil {
			return facility.Hold{}, nil, fmt.Errorf("releasing the hold: %w", err)
		}
	}
	delete(f.holds, d)
	f.applyLocked()
	var readied []string
	if h.Drain {
		pol := f.inv.Facility()
		for _, n := range f.inv.Nodes() {
			doms := n.Domains()
			if !slices.Contains(doms, d) || n.State != inventory.Draining {
				continue
			}
			if _, still := pol.HeldIn(doms); still {
				continue
			}
			f.inv.SetState(n.ID, inventory.Ready)
			readied = append(readied, n.ID)
		}
	}
	f.log.Info("domain released", "domain", d, "readied", readied)
	return h, readied, nil
}

// SetLabels sets and removes labels on a node through the control plane,
// over what its agent reports, and answers the operator's labels as they
// now are. A node not registered yet takes them when it does.
func (f *Facility) SetLabels(ctx context.Context, node string, set map[string]string, remove []string, by string) (map[string]string, error) {
	for k, v := range set {
		if err := checkLabel(k, v); err != nil {
			return nil, err
		}
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	now := f.inv.OperatorLabels(node)
	if now == nil {
		now = map[string]string{}
	}
	if f.st != nil {
		for _, k := range slices.Sorted(maps.Keys(set)) {
			if err := f.st.SetNodeLabel(ctx, node, k, set[k], by); err != nil {
				return nil, fmt.Errorf("keeping label %s: %w", k, err)
			}
		}
		for _, k := range remove {
			if err := f.st.SetNodeLabel(ctx, node, k, "", by); err != nil {
				return nil, fmt.Errorf("removing label %s: %w", k, err)
			}
		}
	}
	maps.Copy(now, set)
	for _, k := range remove {
		delete(now, k)
	}
	f.inv.SetOperatorLabels(node, now)
	f.log.Info("node labels set", "node", node, "set", set, "removed", remove, "by", by)
	return now, nil
}

// checkLabel refuses a key no selector could match, and a misspelt location:
// llmgw.io/rak would be kept, shown, and never read by placement.
func checkLabel(k, v string) error {
	switch {
	case strings.TrimSpace(k) == "" || strings.ContainsAny(k, " =\t\n"):
		return fmt.Errorf("label key %q is empty or has a space or = in it", k)
	case strings.TrimSpace(v) == "":
		return fmt.Errorf("label %s has no value; remove it with %s- instead", k, k)
	case strings.HasPrefix(k, facility.LabelPrefix):
		if _, err := facility.ParseLevel(strings.TrimPrefix(k, facility.LabelPrefix)); err != nil {
			return fmt.Errorf("label %s: %s is reserved for where a node is: llmgw.io/site, room, row, rack or pdu", k, facility.LabelPrefix)
		}
	}
	return nil
}