orchestra/internal/registry/registry.go

// Package registry contains the static project, machine, and herdr topology.
package registry

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"sort"
	"strings"
	"time"
)

var (
	ErrUnknownProject = errors.New("unknown project")
	ErrUnknownMachine = errors.New("unknown machine")
	ErrNoAffinity     = errors.New("project has no machine affinity")
)

type Project struct {
	ID              string   `json:"id"`
	MachineAffinity []string `json:"machine_affinity"`
}
type Machine struct {
	ID      string `json:"id"`
	Address string `json:"address"`
}
type Herdr struct {
	ID           string   `json:"id"`
	MachineID    string   `json:"machine_id"`
	Address      string   `json:"address,omitempty"`
	Capabilities []string `json:"capabilities"`
	Concurrency  int      `json:"concurrency"`
}
type Config struct {
	Projects []Project `json:"projects"`
	Machines []Machine `json:"machines"`
	Herdrs   []Herdr   `json:"herdrs"`
}
type Registry struct {
	projects map[string]Project
	machines map[string]Machine
	herdrs   map[string]Herdr
}

func Load(path string) (Registry, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return Registry{}, err
	}
	var c Config
	if err = json.Unmarshal(b, &c); err != nil {
		return Registry{}, fmt.Errorf("registry config: %w", err)
	}
	return New(c)
}

func New(c Config) (Registry, error) {
	r := Registry{map[string]Project{}, map[string]Machine{}, map[string]Herdr{}}
	for _, p := range c.Projects {
		if err := putID(r.projects, p.ID, "project"); err != nil {
			return Registry{}, err
		}
		if len(p.MachineAffinity) == 0 {
			return Registry{}, fmt.Errorf("project %q: %w", p.ID, ErrNoAffinity)
		}
		r.projects[p.ID] = p
	}
	for _, m := range c.Machines {
		if err := putID(r.machines, m.ID, "machine"); err != nil {
			return Registry{}, err
		}
		if strings.TrimSpace(m.Address) == "" {
			return Registry{}, fmt.Errorf("machine %q: address required", m.ID)
		}
		r.machines[m.ID] = m
	}
	for _, h := range c.Herdrs {
		if err := putID(r.herdrs, h.ID, "herdr"); err != nil {
			return Registry{}, err
		}
		if _, ok := r.machines[h.MachineID]; !ok {
			return Registry{}, fmt.Errorf("herdr %q: %w %q", h.ID, ErrUnknownMachine, h.MachineID)
		}
		if h.Concurrency < 0 {
			return Registry{}, fmt.Errorf("herdr %q: negative concurrency", h.ID)
		}
		r.herdrs[h.ID] = h
	}
	for _, p := range r.projects {
		for _, m := range p.MachineAffinity {
			if _, ok := r.machines[m]; !ok {
				return Registry{}, fmt.Errorf("project %q: %w %q", p.ID, ErrUnknownMachine, m)
			}
		}
	}
	return r, nil
}
func putID[T any](m map[string]T, id, kind string) error {
	if strings.TrimSpace(id) == "" {
		return fmt.Errorf("%s id required", kind)
	}
	if _, ok := m[id]; ok {
		return fmt.Errorf("duplicate %s %q", kind, id)
	}
	return nil
}
func (r Registry) Project(id string) (Project, bool) { p, ok := r.projects[id]; return p, ok }
func (r Registry) Machine(id string) (Machine, bool) { m, ok := r.machines[id]; return m, ok }
func (r Registry) Herdr(id string) (Herdr, bool)     { h, ok := r.herdrs[id]; return h, ok }
func (r Registry) Projects() []Project               { return projects(r.projects) }
func projects(m map[string]Project) []Project {
	out := make([]Project, 0, len(m))
	for _, v := range m {
		out = append(out, v)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out
}

type Reachability interface {
	Reachable(address string, timeout time.Duration) bool
}
type TCPReachability struct{}

func (TCPReachability) Reachable(address string, timeout time.Duration) bool {
	c, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return false
	}
	_ = c.Close()
	return true
}
func (r Registry) Candidates(project string, check Reachability, timeout time.Duration) ([]Herdr, error) {
	p, ok := r.projects[project]
	if !ok {
		return nil, ErrUnknownProject
	}
	allowed := map[string]bool{}
	for _, m := range p.MachineAffinity {
		allowed[m] = true
	}
	out := []Herdr{}
	for _, h := range r.herdrs {
		if !allowed[h.MachineID] {
			continue
		}
		addr := h.Address
		if addr == "" {
			addr = r.machines[h.MachineID].Address
		}
		if check == nil || check.Reachable(addr, timeout) {
			out = append(out, h)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out, nil
}