orchestra/internal/router/router.go

// Package router assigns queued tasks to registered, reachable herdrs.
package router

import (
	"encoding/json"
	"errors"
	"orchestra/internal/domain"
	"orchestra/internal/registry"
	"orchestra/internal/store"
	"sort"
	"strings"
	"time"
)

type Availability interface{ Available(h registry.Herdr) bool }
type AlwaysAvailable struct{}

func (AlwaysAvailable) Available(registry.Herdr) bool { return true }

// QuotaAvailability applies the conservative 80% rule to the most recent
// native quota report in the configured rolling window.
type QuotaAvailability struct {
	Store  *store.Store
	Limits map[string]float64
	Window time.Duration
	Now    func() time.Time
}

func (q QuotaAvailability) Available(h registry.Herdr) bool {
	if q.Store == nil {
		return false
	}
	limit, bounded := q.Limits[h.ID]
	if !bounded || limit <= 0 {
		return true
	}
	now := time.Now()
	if q.Now != nil {
		now = q.Now()
	}
	window := q.Window
	if window <= 0 {
		window = 7 * 24 * time.Hour
	}
	var consumed float64
	for _, e := range q.Store.Events(0) {
		if e.Type != "QuotaReported" || e.At.Before(now.Add(-window)) {
			continue
		}
		var p struct {
			HarnessID string  `json:"harness_id"`
			Consumed  float64 `json:"consumed"`
		}
		if json.Unmarshal(e.Payload, &p) == nil && p.HarnessID == h.ID && p.Consumed > consumed {
			consumed = p.Consumed
		}
	}
	return consumed < limit*0.8
}

type RetryPolicy struct {
	MaxAttempts int
	Backoff     time.Duration
}
type Router struct {
	Store        *store.Store
	Registry     registry.Registry
	Reachability registry.Reachability
	Availability Availability
	Timeout      time.Duration
	Retry        RetryPolicy
	Now          func() time.Time
	backoff      map[string]time.Time
	attempts     map[string]int
	OnLease      func(domain.Event) error
}

func (r *Router) init() {
	if r.Availability == nil {
		r.Availability = AlwaysAvailable{}
	}
	if r.Now == nil {
		r.Now = time.Now
	}
	if r.backoff == nil {
		r.backoff = map[string]time.Time{}
	}
	if r.attempts == nil {
		r.attempts = map[string]int{}
	}
}

// HandleEvent evaluates the sink after creation and after a lease is freed.
func (r *Router) HandleEvent(e domain.Event) ([]domain.Event, error) {
	r.init()
	if e.Type != "TaskCreated" && e.Type != "TaskReleased" {
		return nil, nil
	}
	if e.Type == "TaskReleased" {
		r.attempts[e.TaskID]++
		if r.Retry.Backoff > 0 {
			r.backoff[e.TaskID] = r.Now().Add(r.Retry.Backoff)
		}
	}
	return r.AssignPending()
}

func (r *Router) AssignPending() ([]domain.Event, error) {
	r.init()
	if r.Store == nil {
		return nil, errors.New("router: store required")
	}
	var queued []domain.Task
	for _, t := range r.Store.Tasks() {
		if t.State == domain.StateQueued && !r.Now().Before(r.backoff[t.ID]) {
			queued = append(queued, t)
		}
	}
	sort.SliceStable(queued, func(i, j int) bool { return importance(queued[i], r.Now()).Before(importance(queued[j], r.Now())) })
	var out []domain.Event
	for _, t := range queued {
		if r.Retry.MaxAttempts > 0 && r.attempts[t.ID] >= r.Retry.MaxAttempts {
			e, err := r.fail(t)
			if err != nil {
				return out, err
			}
			out = append(out, e)
			continue
		}
		cs, err := r.Registry.Candidates(t.Project, r.Reachability, r.Timeout)
		if err != nil {
			continue
		}
		for _, h := range cs {
			if !matches(t.Capability, h.Capabilities) || !r.Availability.Available(h) || occupied(r.Store, h.ID, h.Concurrency) {
				continue
			}
			e, err := r.Store.Lease(t.ID, h.ID, 30*time.Minute)
			if err != nil {
				continue
			}
			r.attempts[t.ID]++
			out = append(out, e)
			if r.OnLease != nil {
				if err := r.OnLease(e); err != nil {
					return out, err
				}
			}
			break
		}
	}
	return out, nil
}

func matches(need, have []string) bool {
	set := map[string]bool{}
	for _, x := range have {
		set[strings.ToLower(x)] = true
	}
	for _, x := range need {
		if !set[strings.ToLower(x)] {
			return false
		}
	}
	return true
}
func occupied(s *store.Store, id string, limit int) bool {
	if limit <= 0 {
		return false
	}
	n := 0
	for _, t := range s.Tasks() {
		if t.State == domain.StateLeased && t.Lease != nil && t.Lease.HarnessID == id {
			n++
		}
	}
	return n >= limit
}
func importance(t domain.Task, now time.Time) time.Time {
	if t.Due != nil {
		return t.Due.Add(-time.Duration(t.InherentPriority) * time.Hour)
	}
	return now.Add(-time.Duration(t.InherentPriority) * time.Hour)
}
func (r *Router) fail(t domain.Task) (domain.Event, error) {
	b, _ := json.Marshal(map[string]any{"reason": "retry_limit", "attempts": r.attempts[t.ID]})
	e := domain.Event{ID: domain.NewID(), Type: "TaskFailed", TaskID: t.ID, Version: t.Version + 1, Payload: b}
	return e, r.Store.Append(e)
}