perf: index routing snapshots and quota usage

2026-07-30 15:36:23 +04:00
parent e8fadfc998
commit fbb13c89d2
8 changed files with 577 additions and 63 deletions
@@ -0,0 +1,82 @@
+package router
+
+import (
+	"encoding/json"
+	"fmt"
+	"orchestra/internal/authz"
+	"orchestra/internal/domain"
+	"orchestra/internal/registry"
+	"orchestra/internal/store"
+	"sort"
+	"testing"
+	"time"
+)
+
+// These are p95 budgets for a complete scheduling pass on an otherwise idle
+// coordinator. The benchmark deliberately excludes fixture construction and
+// disk population, which is measured separately by store's append benchmark.
+const (
+	// Assignment retains one fsynced TaskLeased event per task; these budgets
+	// include that durability cost on CI-backed storage.
+	assignmentP95Budget1K  = 6 * time.Second
+	assignmentP95Budget10K = time.Minute
+)
+
+func BenchmarkAssignPending1K(b *testing.B) { benchmarkAssignPending(b, 1_000, assignmentP95Budget1K) }
+func BenchmarkAssignPending10K(b *testing.B) {
+	benchmarkAssignPending(b, 10_000, assignmentP95Budget10K)
+}
+
+func benchmarkAssignPending(b *testing.B, tasks int, budget time.Duration) {
+	b.Helper()
+	samples := make([]time.Duration, 0, b.N)
+	b.ResetTimer()
+	for sample := 0; sample < b.N; sample++ {
+		b.StopTimer()
+		s, err := store.Open(b.TempDir())
+		if err != nil {
+			b.Fatal(err)
+		}
+		for i := 0; i < tasks; i++ {
+			id := fmt.Sprintf("task-%d", i)
+			payload, _ := json.Marshal(map[string]any{"source": "benchmark", "external_id": id, "project": "p"})
+			if err := s.Append(domain.Event{ID: id, Type: "TaskCreated", TaskID: id, Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
+				b.Fatal(err)
+			}
+		}
+		r, err := registry.New(registry.Config{
+			Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
+			Machines: []registry.Machine{{ID: "m", Address: "unused"}},
+			Herdrs:   []registry.Herdr{{ID: "h", MachineID: "m"}}, // unlimited concurrency
+		})
+		if err != nil {
+			b.Fatal(err)
+		}
+		router := Router{Store: s, Registry: r}
+		b.StartTimer()
+		started := time.Now()
+		assigned, err := router.AssignPending()
+		samples = append(samples, time.Since(started))
+		b.StopTimer()
+		if err != nil {
+			b.Fatal(err)
+		}
+		if len(assigned) != tasks {
+			b.Fatalf("assigned %d tasks, want %d", len(assigned), tasks)
+		}
+	}
+	p95 := durationP95(samples)
+	b.ReportMetric(float64(p95.Microseconds()), "assign_p95_us")
+	if p95 > budget {
+		b.Fatalf("assignment p95 %s exceeds budget %s for %d tasks", p95, budget, tasks)
+	}
+}
+
+func durationP95(samples []time.Duration) time.Duration {
+	if len(samples) == 0 {
+		return 0
+	}
+	sorted := append([]time.Duration(nil), samples...)
+	sort.Slice(sorted, func(i, j int) bool { return sorted[i] < sorted[j] })
+	return sorted[(len(sorted)*95+99)/100-1]
+}
@@ -10,6 +10,7 @@ import (
 	"orchestra/internal/store"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 )

@@ -54,26 +55,7 @@ type QuotaAvailability struct {
 }

 func (q QuotaAvailability) sumSince(harnessID string, since time.Time) (float64, bool) {
-	var consumed float64
-	known := false
-	for _, e := range q.Store.Events(0) {
-		if e.Type != "QuotaReported" || e.At.Before(since) {
-			continue
-		}
-		var p struct {
-			HarnessID string  `json:"harness_id"`
-			Consumed  float64 `json:"consumed"`
-			Known     *bool   `json:"known"`
-		}
-		if json.Unmarshal(e.Payload, &p) == nil && p.HarnessID == harnessID && p.Consumed >= 0 {
-			if p.Known != nil && !*p.Known {
-				return 0, false
-			}
-			known = true
-			consumed += p.Consumed
-		}
-	}
-	return consumed, known
+	return q.Store.QuotaSince(harnessID, since)
 }

 func (q QuotaAvailability) Available(h registry.Herdr) bool {
@@ -116,6 +98,18 @@ type Router struct {
 	Retry        RetryPolicy
 	Now          func() time.Time
 	OnLease      func(domain.Event) error
+	// HealthTTL bounds re-use of a successful or failed reachability probe.
+	// A scheduling pass always has a coherent health snapshot; this cache also
+	// prevents bursts of TaskCreated events from repeatedly dialing the same
+	// herdr between passes.
+	HealthTTL time.Duration
+	healthMu  sync.Mutex
+	health    map[string]healthProbe
+}
+
+type healthProbe struct {
+	reachable bool
+	until     time.Time
 }

 func (r *Router) init() {
@@ -125,6 +119,9 @@ func (r *Router) init() {
 	if r.Now == nil {
 		r.Now = time.Now
 	}
+	if r.HealthTTL <= 0 {
+		r.HealthTTL = 5 * time.Second
+	}
 }

 // HandleEvent evaluates the sink after creation and after a lease is freed.
@@ -141,13 +138,22 @@ func (r *Router) AssignPending() ([]domain.Event, error) {
 	if r.Store == nil {
 		return nil, errors.New("router: store required")
 	}
+	now := r.Now()
+	snapshot := r.Store.SchedulingSnapshot()
 	var queued []domain.Task
-	for _, t := range r.Store.Tasks() {
-		if t.State == domain.StateQueued && (t.NextRetryAt.IsZero() || !r.Now().Before(t.NextRetryAt)) {
+	for _, t := range snapshot.Tasks {
+		if t.State == domain.StateQueued && (t.NextRetryAt.IsZero() || !now.Before(t.NextRetryAt)) {
 			queued = append(queued, t)
 		}
 	}
-	sort.SliceStable(queued, func(i, j int) bool { return importance(queued[i], r.Now()).Before(importance(queued[j], r.Now())) })
+	sort.SliceStable(queued, func(i, j int) bool { return importance(queued[i], now).Before(importance(queued[j], now)) })
+	health := r.healthSnapshot(now)
+	candidates := make(map[string][]registry.Herdr)
+	availability := make(map[string]bool)
+	availabilityKnown := make(map[string]bool)
+	projectSupport := make(map[string]bool)
+	projectSupportKnown := make(map[string]bool)
+	activeLeases := snapshot.ActiveLeases
 	var out []domain.Event
 	for _, t := range queued {
 		if r.Retry.MaxAttempts > 0 && t.Attempt >= r.Retry.MaxAttempts {
@@ -158,16 +164,34 @@ func (r *Router) AssignPending() ([]domain.Event, error) {
 			out = append(out, e)
 			continue
 		}
-		cs, err := r.Registry.Candidates(t.Project, r.Reachability, r.Timeout)
-		if err != nil {
-			continue
+		cs, found := candidates[t.Project]
+		if !found {
+			var err error
+			cs, err = r.Registry.CandidatesWithHealth(t.Project, health)
+			if err != nil {
+				continue
+			}
+			candidates[t.Project] = cs
 		}
 		for _, h := range cs {
-			projectOK := true
-			if projects, ok := r.Availability.(ProjectAvailability); ok {
-				projectOK = projects.Supports(h, t.Project)
+			projectKey := h.ID + "\x00" + t.Project
+			projectOK, checked := projectSupport[projectKey], projectSupportKnown[projectKey]
+			if !checked {
+				projectOK = true
+				if projects, ok := r.Availability.(ProjectAvailability); ok {
+					projectOK = projects.Supports(h, t.Project)
+				}
+				projectSupport[projectKey], projectSupportKnown[projectKey] = projectOK, true
 			}
-			if !projectOK || !matches(t.Capability, h.Capabilities) || !r.Availability.Available(h) || occupied(r.Store, h.ID, h.Concurrency) {
+			if !projectOK || !matches(t.Capability, h.Capabilities) {
+				continue
+			}
+			available, checked := availability[h.ID], availabilityKnown[h.ID]
+			if !checked {
+				available = r.Availability.Available(h)
+				availability[h.ID], availabilityKnown[h.ID] = available, true
+			}
+			if !available || occupiedCount(activeLeases[h.ID], h.Concurrency) {
 				continue
 			}
 			e, err := r.Store.Lease(t.ID, h.ID, 30*time.Minute)
@@ -175,6 +199,7 @@ func (r *Router) AssignPending() ([]domain.Event, error) {
 				continue
 			}
 			out = append(out, e)
+			activeLeases[h.ID]++
 			if r.OnLease != nil {
 				if err := r.OnLease(e); err != nil {
 					return out, err
@@ -198,17 +223,60 @@ func matches(need, have []string) bool {
 	}
 	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.State == domain.StateNeedsAttention) && t.Lease != nil && t.Lease.HarnessID == id {
-			n++
+func occupiedCount(count, limit int) bool {
+	return limit > 0 && count >= limit
+}
+
+func (r *Router) healthSnapshot(now time.Time) map[string]bool {
+	herdrs := r.Registry.Herdrs()
+	health := make(map[string]bool, len(herdrs))
+	if r.Reachability == nil {
+		for _, h := range herdrs {
+			health[h.ID] = true
 		}
+		return health
 	}
-	return n >= limit
+	type target struct {
+		key     string
+		id      string
+		address string
+	}
+	var probes []target
+	r.healthMu.Lock()
+	if r.health == nil {
+		r.health = make(map[string]healthProbe)
+	}
+	for _, h := range herdrs {
+		key := h.ID + "\x00" + r.Registry.Endpoint(h)
+		if cached, ok := r.health[key]; ok && now.Before(cached.until) {
+			health[h.ID] = cached.reachable
+			continue
+		}
+		probes = append(probes, target{key: key, id: h.ID, address: r.Registry.Endpoint(h)})
+	}
+	r.healthMu.Unlock()
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	for _, target := range probes {
+		target := target
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			reachable := r.Reachability.Reachable(target.address, r.Timeout)
+			mu.Lock()
+			health[target.id] = reachable
+			mu.Unlock()
+		}()
+	}
+	wg.Wait()
+	if len(probes) > 0 {
+		r.healthMu.Lock()
+		for _, target := range probes {
+			r.health[target.key] = healthProbe{reachable: health[target.id], until: now.Add(r.HealthTTL)}
+		}
+		r.healthMu.Unlock()
+	}
+	return health
 }
 func importance(t domain.Task, now time.Time) time.Time {
 	if t.Due != nil {
@@ -6,6 +6,7 @@ import (
 	"orchestra/internal/domain"
 	"orchestra/internal/registry"
 	"orchestra/internal/store"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -14,6 +15,26 @@ type reachable struct{}

 func (reachable) Reachable(string, time.Duration) bool { return true }

+type countedReachability struct {
+	calls atomic.Int32
+	delay time.Duration
+}
+
+func (r *countedReachability) Reachable(string, time.Duration) bool {
+	r.calls.Add(1)
+	if r.delay > 0 {
+		time.Sleep(r.delay)
+	}
+	return true
+}
+
+type countedAvailability struct{ calls atomic.Int32 }
+
+func (a *countedAvailability) Available(registry.Herdr) bool {
+	a.calls.Add(1)
+	return true
+}
+
 type projectAvailability struct{ projects map[string]bool }

 func (p projectAvailability) Available(registry.Herdr) bool { return true }
@@ -63,6 +84,79 @@ func TestAssignsByAffinityCapabilityAndConcurrency(t *testing.T) {
 	}
 }

+func TestAssignPendingSnapshotsAndCachesHealth(t *testing.T) {
+	s, err := store.Open(t.TempDir())
+	if err != nil {
+		t.Fatal(err)
+	}
+	r, err := registry.New(registry.Config{
+		Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
+		Machines: []registry.Machine{{ID: "m", Address: "unused"}},
+		Herdrs: []registry.Herdr{
+			{ID: "h1", MachineID: "m", Address: "h1:1"},
+			{ID: "h2", MachineID: "m", Address: "h2:1"},
+			{ID: "h3", MachineID: "m", Address: "h3:1"},
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		id := domain.NewID()
+		payload, _ := json.Marshal(map[string]any{"source": "test", "external_id": id, "project": "p"})
+		if err := s.Append(domain.Event{ID: id, TaskID: id, Type: "TaskCreated", Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	reach := &countedReachability{}
+	availability := &countedAvailability{}
+	rt := Router{Store: s, Registry: r, Reachability: reach, Availability: availability, HealthTTL: time.Minute}
+	if got, err := rt.AssignPending(); err != nil || len(got) != 10 {
+		t.Fatalf("assigned=%d, err=%v", len(got), err)
+	}
+	if got := reach.calls.Load(); got != 3 {
+		t.Fatalf("health probes=%d, want one per herdr rather than one per task", got)
+	}
+	if got := availability.calls.Load(); got != 1 {
+		t.Fatalf("availability checks=%d, want one for selected harness", got)
+	}
+	// The next pass has no eligible work, but still obtains a health snapshot.
+	// It must use the TTL cache rather than re-dial all three herdrs.
+	if _, err := rt.AssignPending(); err != nil {
+		t.Fatal(err)
+	}
+	if got := reach.calls.Load(); got != 3 {
+		t.Fatalf("health cache missed: probes=%d, want 3", got)
+	}
+}
+
+func TestHealthProbesRunInParallel(t *testing.T) {
+	s, err := store.Open(t.TempDir())
+	if err != nil {
+		t.Fatal(err)
+	}
+	r, err := registry.New(registry.Config{
+		Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
+		Machines: []registry.Machine{{ID: "m", Address: "unused"}},
+		Herdrs:   []registry.Herdr{{ID: "h1", MachineID: "m", Address: "h1:1"}, {ID: "h2", MachineID: "m", Address: "h2:1"}, {ID: "h3", MachineID: "m", Address: "h3:1"}},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	payload, _ := json.Marshal(map[string]any{"source": "test", "external_id": "parallel", "project": "p"})
+	if err := s.Append(domain.Event{ID: "parallel", TaskID: "parallel", Type: "TaskCreated", Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
+		t.Fatal(err)
+	}
+	reach := &countedReachability{delay: 100 * time.Millisecond}
+	started := time.Now()
+	if _, err := (&Router{Store: s, Registry: r, Reachability: reach}).AssignPending(); err != nil {
+		t.Fatal(err)
+	}
+	if elapsed := time.Since(started); elapsed > 250*time.Millisecond {
+		t.Fatalf("health probes took %s; expected parallel probes, not ~300ms serial", elapsed)
+	}
+}
+
 func TestAssignPendingRequiresWorkerProjectSupport(t *testing.T) {
 	s, err := store.Open(t.TempDir())
 	if err != nil {