orchestra/internal/federation/federation.go

package federation

import (
	"errors"
	"sync"
	"time"
)

var ErrUnknownWorker = errors.New("unknown worker")
var ErrUnauthorized = errors.New("worker authentication failed")

type Worker struct {
	ID       string    `json:"id"`
	Address  string    `json:"address"`
	Capacity int       `json:"capacity"`
	LastSeen time.Time `json:"last_seen"`
	Online   bool      `json:"online"`
	Token    string    `json:"-"`
}

type Registry struct {
	mu sync.Mutex
	// AdmitToken, if set, is a pre-shared secret every registration must
	// present (S10: registration previously accepted a self-declared id and
	// self-chosen token from any caller — admission-control-free). Leave
	// empty only for a deliberately open deployment.
	AdmitToken string
	workers    map[string]Worker
	TTL        time.Duration
	OnOffline  func(Worker)
	cursors    map[string]uint64
}

func (r *Registry) init() {
	if r.TTL <= 0 {
		r.TTL = 90 * time.Second
	}
	if r.workers == nil {
		r.workers = map[string]Worker{}
	}
	if r.cursors == nil {
		r.cursors = map[string]uint64{}
	}
}
// Register admits a worker. admitToken must match r.AdmitToken whenever one
// is configured. Re-registering an ID that's already claimed requires that
// worker's own current token, so a caller can't self-declare someone else's
// id and hijack an existing worker's identity/capacity.
func (r *Registry) Register(w Worker, admitToken string) error {
	if w.ID == "" {
		return errors.New("worker id required")
	}
	if w.Token == "" {
		return errors.New("worker token required")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if r.AdmitToken != "" && admitToken != r.AdmitToken {
		return ErrUnauthorized
	}
	if existing, ok := r.workers[w.ID]; ok && existing.Token != w.Token {
		return ErrUnauthorized
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	r.workers[w.ID] = w
	if _, ok := r.cursors[w.ID]; !ok {
		r.cursors[w.ID] = 0
	}
	return nil
}
func (r *Registry) Authenticate(id, token string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	if w.Token == "" || token == "" || w.Token != token {
		return ErrUnauthorized
	}
	return nil
}
func (r *Registry) Cursor(id string) (uint64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return 0, ErrUnknownWorker
	}
	return r.cursors[id], nil
}
func (r *Registry) Ack(id string, cursor uint64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	if _, ok := r.workers[id]; !ok {
		return ErrUnknownWorker
	}
	if cursor < r.cursors[id] {
		return errors.New("cursor moved backwards")
	}
	r.cursors[id] = cursor
	return nil
}
func (r *Registry) Heartbeat(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	w, ok := r.workers[id]
	if !ok {
		return ErrUnknownWorker
	}
	w.LastSeen = time.Now().UTC()
	w.Online = true
	r.workers[id] = w
	return nil
}
func (r *Registry) Snapshot() []Worker {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.init()
	now := time.Now().UTC()
	out := make([]Worker, 0, len(r.workers))
	for id, w := range r.workers {
		wasOnline := w.Online
		w.Online = now.Sub(w.LastSeen) <= r.TTL
		r.workers[id] = w
		if wasOnline && !w.Online && r.OnOffline != nil {
			go r.OnOffline(w)
		}
		out = append(out, w)
	}
	return out
}