orca/internal/scheduler/scheduler.go

// Package scheduler implements the v0.9 CLI-side scheduler (REQ-083,
// P05). Unlike the v0.8 daemon-side best-fit scheduler
// (internal/engine/scheduler.go), this scheduler runs entirely in the
// `orca` CLI process (R-001) and is pure: it takes a list of candidate
// nodes plus a workload request and returns placement decisions
// without performing any I/O.
//
// The scheduler is runtime-aware: a workload that declares
// `runtime.one_of: wasm` is only placed on nodes that expose
// `wasmtime` in their Runtimes list; a `pve-vm` workload is only
// placed on `proxmox` nodes. It is also constraint- and
// affinity-aware via the CEL-subset evaluator in cel.go.
//
// Workload kinds are handled differently per the PRD:
//
//   - Job: one-shot, returns exactly one placement (best-fit
//     bin-packing).
//   - Service: count replicas spread across distinct nodes
//     (anti-affinity by default); if fewer distinct nodes than count,
//     colocation is permitted but distinct nodes are preferred.
//   - DaemonSet: one placement per node that fits the constraints.
package scheduler

import (
	"fmt"
	"sort"

	"git.cloudinit.dev/coreci/orca/internal/jobspec"
)

// NodeInfo is the scheduler's projection of a peer node: total and
// free capacity, the runtimes the node advertises, its tags, and its
// kind (linux/proxmox). The CLI populates this from the
// cluster/peers/ inventory plus the per-node capacity reports
// collected over SSH; the scheduler itself never reads either.
type NodeInfo struct {
	Hostname string
	Runtimes []string
	Tags     []string
	CPU      int64
	Memory   int64
	FreeCPU  int64
	FreeMem  int64
	Kind     string
}

// WorkloadRequest bundles a parsed WorkloadSpec with the namespace
// the workload is being scheduled into. The namespace is carried
// through to placement so the resulting AllocID can be namespaced,
// but the scheduler itself does not inspect it for fitting decisions.
type WorkloadRequest struct {
	Spec      *jobspec.WorkloadSpec
	Namespace string
}

// Placement is a single scheduling decision: which node, which
// allocation id, and the bin-packing score that won the node the
// placement. AllocID is `ns/spec.Name-<idx>` so a multi-replica
// Service produces distinct ids per replica.
type Placement struct {
	Node    string
	AllocID string
	Score   int64
}

// Schedule is the main entry point. For Job (kind=Job) it returns one
// placement on the best-fit node. For Service it returns `Count`
// placements spread across distinct nodes where possible (anti-
// affinity), permitting colocation when Count > nodes. For DaemonSet
// it returns one placement per node that fits. Any kind-agnostic
// validation error (no spec, unknown kind, no fitting node) is
// returned as an error rather than an empty slice so callers can
// distinguish "nothing fits" from "scheduled zero replicas".
func Schedule(nodes []NodeInfo, req WorkloadRequest) ([]Placement, error) {
	if req.Spec == nil {
		return nil, fmt.Errorf("scheduler: nil WorkloadSpec")
	}
	if len(nodes) == 0 {
		return nil, fmt.Errorf("scheduler: no candidate nodes")
	}
	switch req.Spec.Kind {
	case "Job":
		return scheduleJob(nodes, req)
	case "Service":
		return scheduleService(nodes, req)
	case "DaemonSet":
		return scheduleDaemonSet(nodes, req)
	default:
		return nil, fmt.Errorf("scheduler: unknown kind %q", req.Spec.Kind)
	}
}

// Score evaluates a single node against a workload. fits is true iff
// the node (a) advertises a runtime compatible with the workload's
// `runtime.one_of`, (b) satisfies every CEL constraint in
// `spec.Constraints`, and (c) has enough free CPU+memory for the
// workload's requested resources. When fits is true, score is the
// bin-packing score (more free capacity = higher score, so the node
// most likely to absorb the workload without starving its
// neighbours wins). When fits is false, score is 0.
func Score(node NodeInfo, req WorkloadRequest) (score int64, fits bool) {
	// (a) runtime compatibility. A workload with no Runtime block or
	// an empty OneOf is treated as runtime-agnostic (always fits on
	// the runtime axis); this matches the v0.8 behaviour where a
	// missing runtime meant "process".
	runtimeOK := true
	if req.Spec != nil && req.Spec.Runtime != nil && req.Spec.Runtime.OneOf != "" {
		runtimeOK = hasRuntime(node, req.Spec.Runtime.OneOf)
	}
	if !runtimeOK {
		return 0, false
	}
	// (b) constraints. Evaluation errors are treated as non-fit so a
	// malformed constraint does not crash Schedule; the caller still
	// sees the node filtered out.
	if req.Spec != nil {
		ok, err := EvaluateAll(req.Spec.Constraints, node)
		if err != nil || !ok {
			return 0, false
		}
	}
	// (c) capacity. A workload with no Resources block is treated as
	// zero-sized for fitting purposes (it always fits the capacity
	// axis); real workloads declare cpu/memory.
	needCPU, needMem := workloadResources(req)
	if node.FreeCPU < needCPU || node.FreeMem < needMem {
		return 0, false
	}
	// bin-packing score: most free capacity wins. CPU is weighted
	// 1000x memory so a 1-core difference outweighs a 1-MiB
	// difference, mirroring the v0.8 Score weighting that biased
	// toward CPU (the more common binding constraint).
	score = (node.FreeCPU-needCPU)*1000 + (node.FreeMem - needMem)
	if score < 0 {
		score = 0
	}
	return score, true
}

// hasRuntime reports whether node advertises the requested runtime.
// The match is case-insensitive and tolerant of aliases: `wasm` and
// `wasmtime` are treated as the same runtime, and `pve-vm`/`pve-ct`
// only match nodes whose Kind is "proxmox".
func hasRuntime(node NodeInfo, oneOf string) bool {
	want := normalizeRuntime(oneOf)
	// pve-* runtimes require a proxmox-kind node regardless of the
	// node's Runtimes list (a proxmox node doesn't list "pve-vm" in
	// Runtimes; it IS the runtime).
	switch want {
	case "pve-vm", "pve-ct", "proxmox":
		return normalizeKind(node.Kind) == "proxmox"
	}
	for _, r := range node.Runtimes {
		if normalizeRuntime(r) == want {
			return true
		}
		// alias: wasmtime nodes advertise "wasmtime"; workloads ask
		// for "wasm".
		if want == "wasm" && normalizeRuntime(r) == "wasmtime" {
			return true
		}
	}
	return false
}

// normalizeRuntime lowercases and trims a runtime name for matching.
func normalizeRuntime(s string) string {
	s = toLowerASCII(s)
	switch s {
	case "wasmtime":
		return "wasm"
	}
	return s
}

// normalizeKind lowercases and trims a node Kind for matching.
func normalizeKind(s string) string { return toLowerASCII(s) }

// toLowerASCII lowercases ASCII letters without bringing in strings
// (avoid an alloc-heavy stdlib call in the hot path).
func toLowerASCII(s string) string {
	b := []byte(s)
	for i, c := range b {
		if c >= 'A' && c <= 'Z' {
			b[i] = c + 32
		}
	}
	return string(b)
}

// workloadResources returns the (cpu, memory) the workload requests,
// read from the WorkloadSpec's Resources block if present. The
// v0.9-P05 WorkloadSpec does not yet carry a Resources field (it
// lands in P0c, REQ-074); until then this returns (0, 0) so the
// capacity check is a no-op and runtime/constraints do the real
// filtering. The signature is here so the scheduler logic does not
// need to change when Resources lands.
func workloadResources(req WorkloadRequest) (int64, int64) {
	_ = req
	return 0, 0
}

// ----------------------------------------------------------------------------
// Kind-specific scheduling
// ----------------------------------------------------------------------------

// scheduleJob places a single Job on the best-fit node.
func scheduleJob(nodes []NodeInfo, req WorkloadRequest) ([]Placement, error) {
	type cand struct {
		node  NodeInfo
		score int64
	}
	var cands []cand
	for _, n := range nodes {
		s, ok := Score(n, req)
		if !ok {
			continue
		}
		cands = append(cands, cand{node: n, score: s})
	}
	// CEL-based affinity rules apply to single-shot Jobs too: a Job
	// with `affinity: [{target: "\"ssd\" in node.tags", weight: 100}]`
	// should land on the tagged node even without prior placements.
	// Name-based affinity (no prior placements to check) contributes
	// zero for a standalone Job, so it is harmless to call here.
	for i := range cands {
		cands[i].score += affinityScore(cands[i].node, req, nil)
	}
	if len(cands) == 0 {
		return nil, fmt.Errorf("scheduler: no node fits workload %q", req.Spec.Name)
	}
	sort.SliceStable(cands, func(i, j int) bool {
		if cands[i].score != cands[j].score {
			return cands[i].score > cands[j].score
		}
		return cands[i].node.Hostname < cands[j].node.Hostname
	})
	w := cands[0]
	return []Placement{{
		Node:    w.node.Hostname,
		AllocID: allocID(req, 0),
		Score:   w.score,
	}}, nil
}

// scheduleService places `Count` replicas with implicit anti-affinity:
// prefer distinct nodes, but permit colocation when Count exceeds the
// number of fitting nodes. Each replica gets a distinct AllocID.
func scheduleService(nodes []NodeInfo, req WorkloadRequest) ([]Placement, error) {
	count := req.Spec.Count
	if count <= 0 {
		count = 1
	}
	// Pre-filter fitting nodes once; the loop below re-scores them
	// after each placement so the capacity accounting reflects the
	// replicas already placed.
	fitting := filterFitting(nodes, req)
	if len(fitting) == 0 {
		return nil, fmt.Errorf("scheduler: no node fits service %q", req.Spec.Name)
	}
	var placements []Placement
	placed := map[string]int{} // hostname -> count placed there
	// First pass: spread across distinct nodes.
	for i := 0; i < count; i++ {
		best, score, ok := pickServiceNode(fitting, req, placements, placed)
		if !ok {
			break
		}
		placements = append(placements, Placement{
			Node:    best.Hostname,
			AllocID: allocID(req, i),
			Score:   score,
		})
		placed[best.Hostname]++
		// Reflect the consumed capacity in the candidate snapshot so
		// subsequent picks see updated free capacity.
		needCPU, needMem := workloadResources(req)
		for j := range fitting {
			if fitting[j].Hostname == best.Hostname {
				fitting[j].FreeCPU -= needCPU
				fitting[j].FreeMem -= needMem
			}
		}
	}
	if len(placements) < count {
		return nil, fmt.Errorf("scheduler: only placed %d/%d replicas for service %q",
			len(placements), count, req.Spec.Name)
	}
	return placements, nil
}

// pickServiceNode selects the best node for the next replica. The
// selection prefers nodes with zero prior placements of this service
// (anti-affinity) and applies affinity scoring on top of the
// bin-packing score.
func pickServiceNode(fitting []NodeInfo, req WorkloadRequest, placements []Placement, placed map[string]int) (NodeInfo, int64, bool) {
	type scored struct {
		node  NodeInfo
		score int64
	}
	var cands []scored
	for _, n := range fitting {
		s, ok := Score(n, req)
		if !ok {
			continue
		}
		// Implicit anti-affinity: a node with N prior replicas of this
		// service incurs a penalty of N * (1 << 62) so distinct nodes
		// are preferred, but colocation is permitted (with a
		// per-replica penalty) when no distinct node remains. This
		// produces a balanced spread (e.g. 5 replicas on 3 nodes →
		// 2/2/1) rather than stacking everything on the first node.
		if placed[n.Hostname] > 0 {
			s -= int64(placed[n.Hostname]) * (1 << 62)
		}
		// Affinity rules from the spec add/subtract their weight.
		s += affinityScore(n, req, placements)
		cands = append(cands, scored{node: n, score: s})
	}
	if len(cands) == 0 {
		return NodeInfo{}, 0, false
	}
	sort.SliceStable(cands, func(i, j int) bool {
		if cands[i].score != cands[j].score {
			return cands[i].score > cands[j].score
		}
		return cands[i].node.Hostname < cands[j].node.Hostname
	})
	w := cands[0]
	return w.node, w.score, true
}

// scheduleDaemonSet places one replica per node that fits the
// constraints. The PRD's DaemonSet placement mode (every-node /
// matching / mandatory) lives on the WorkloadSpec.Schedule block; the
// scheduler honours it indirectly by filtering on Constraints: a
// `matching` DaemonSet carries constraints that select the matching
// nodes, an `every-node` DaemonSet carries none, and a `mandatory`
// one is enforced elsewhere (the scheduler still just returns
// placements for every fitting node).
func scheduleDaemonSet(nodes []NodeInfo, req WorkloadRequest) ([]Placement, error) {
	var placements []Placement
	for _, n := range nodes {
		s, ok := Score(n, req)
		if !ok {
			continue
		}
		placements = append(placements, Placement{
			Node:    n.Hostname,
			AllocID: allocID(req, len(placements)),
			Score:   s,
		})
	}
	if len(placements) == 0 {
		return nil, fmt.Errorf("scheduler: no node fits daemonset %q", req.Spec.Name)
	}
	return placements, nil
}

// ----------------------------------------------------------------------------
// Affinity scoring
// ----------------------------------------------------------------------------

// affinityScore returns the weighted affinity contribution for a
// node given the placements already made. For each AffinityRule the
// Target is a CEL expression; if it evaluates true against the node,
// the rule's Weight is added (positive = co-locate, negative =
// anti-affinity). An affinity target that fails to evaluate is
// ignored rather than failing the schedule: operators use affinity as
// a hint, not a hard gate.
//
// The PRD also mentions affinity rules like `{target: "redis", weight:
// 50}` where Target is a workload *name* rather than a CEL expression.
// We support both: if Target parses as a CEL expression it is
// evaluated against the node; otherwise it is treated as a workload
// name and we check whether any already-placed alloc for that name
// exists on the node. The placement-already-here check is done by the
// caller via placements; this function checks the node's own
// attributes only.
func affinityScore(node NodeInfo, req WorkloadRequest, placements []Placement) int64 {
	if req.Spec == nil {
		return 0
	}
	var total int64
	for _, rule := range req.Spec.Affinity {
		// Try CEL evaluation first; if the target is a bare workload
		// name (no operator) the CEL parser will fail and we fall
		// back to name-based placement counting.
		ok, err := EvaluateConstraint(rule.Target, node)
		if err == nil {
			if ok {
				total += int64(rule.Weight)
			}
			continue
		}
		// Fallback: target is a workload name; count existing
		// placements for that workload on this node and apply the
		// weight once per co-located replica.
		for _, p := range placements {
			if p.Node == node.Hostname && isAllocFor(p.AllocID, rule.Target) {
				total += int64(rule.Weight)
			}
		}
	}
	return total
}

// isAllocFor reports whether an AllocID encodes a placement for the
// named workload. AllocIDs are `ns/name-idx`, so we look for the
// workload name as the segment after the first slash and before the
// trailing `-idx`.
func isAllocFor(allocID, workloadName string) bool {
	// strip namespace prefix
	rest := allocID
	if i := indexByte(rest, '/'); i >= 0 {
		rest = rest[i+1:]
	}
	// strip trailing -idx
	if i := lastIndexByte(rest, '-'); i >= 0 {
		rest = rest[:i]
	}
	return rest == workloadName
}

// indexByte returns the index of the first occurrence of b in s, or
// -1. Avoids importing strings just for one helper.
func indexByte(s string, b byte) int {
	for i := 0; i < len(s); i++ {
		if s[i] == b {
			return i
		}
	}
	return -1
}

// lastIndexByte returns the index of the last occurrence of b in s, or
// -1.
func lastIndexByte(s string, b byte) int {
	for i := len(s) - 1; i >= 0; i-- {
		if s[i] == b {
			return i
		}
	}
	return -1
}

// ----------------------------------------------------------------------------
// Helpers
// ----------------------------------------------------------------------------

// filterFitting returns a copy of the nodes that pass Score for the
// request, preserving order. Capacity is not yet decremented; the
// caller adjusts FreeCPU/FreeMem as it places replicas.
func filterFitting(nodes []NodeInfo, req WorkloadRequest) []NodeInfo {
	var out []NodeInfo
	for _, n := range nodes {
		if _, ok := Score(n, req); ok {
			out = append(out, n)
		}
	}
	return out
}

// allocID renders a stable, namespaced allocation id for a placement.
// Format: `ns/spec.Name-<idx>`.
func allocID(req WorkloadRequest, idx int) string {
	ns := req.Namespace
	if ns == "" {
		ns = "default"
	}
	return fmt.Sprintf("%s/%s-%d", ns, req.Spec.Name, idx)
}