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,60 @@
+package store
+
+import (
+	"encoding/json"
+	"fmt"
+	"orchestra/internal/authz"
+	"orchestra/internal/domain"
+	"sort"
+	"testing"
+	"time"
+)
+
+// Append includes the fsynced event-log write and sparse projection-cache
+// checkpoint. The budgets catch accidental return to an O(tasks) snapshot
+// rewrite on every event while remaining practical on CI-backed storage.
+const (
+	appendP95Budget1K  = 10 * time.Second
+	appendP95Budget10K = 90 * time.Second
+)
+
+func BenchmarkAppend1K(b *testing.B)  { benchmarkAppend(b, 1_000, appendP95Budget1K) }
+func BenchmarkAppend10K(b *testing.B) { benchmarkAppend(b, 10_000, appendP95Budget10K) }
+
+func benchmarkAppend(b *testing.B, events 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 := Open(b.TempDir())
+		if err != nil {
+			b.Fatal(err)
+		}
+		b.StartTimer()
+		started := time.Now()
+		for i := 0; i < events; 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)
+			}
+		}
+		samples = append(samples, time.Since(started))
+		b.StopTimer()
+	}
+	p95 := appendDurationP95(samples)
+	b.ReportMetric(float64(p95.Microseconds()), "append_p95_us")
+	if p95 > budget {
+		b.Fatalf("append p95 %s exceeds budget %s for %d events", p95, budget, events)
+	}
+}
+
+func appendDurationP95(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]
+}
@@ -9,10 +9,26 @@ import (
 	"orchestra/internal/domain"
 	"os"
 	"path/filepath"
+	"sort"
 	"sync"
 	"time"
 )

+// QuotaUsage is an indexed native-usage receipt. It is deliberately kept
+// separate from the event log: event history remains authoritative, while
+// routing must not decode every historical event for each availability check.
+type QuotaUsage struct {
+	At       time.Time
+	Consumed float64
+	Known    bool
+}
+
+type quotaIndex struct {
+	records       []QuotaUsage // ordered by At
+	prefix        []float64
+	unknownPrefix []int
+}
+
 type Store struct {
 	mu       sync.Mutex
 	path     string
@@ -20,15 +36,19 @@ type Store struct {
 	events   []domain.Event
 	tasks    map[string]domain.Task
 	external map[string]string
-	snapshot string
-	seq      uint64
+	// activeLeases is the routing occupancy index. Needs-attention retains a
+	// fenced owner, so it counts as active until that lease is released.
+	activeLeases map[string]map[string]struct{}
+	quota        map[string]quotaIndex
+	snapshot     string
+	seq          uint64
 }

 func Open(dir string) (*Store, error) {
 	if err := os.MkdirAll(dir, 0755); err != nil {
 		return nil, err
 	}
-	s := &Store{path: filepath.Join(dir, "events.jsonl"), cas: filepath.Join(dir, "cas"), snapshot: filepath.Join(dir, "snapshot.json"), tasks: map[string]domain.Task{}, external: map[string]string{}}
+	s := &Store{path: filepath.Join(dir, "events.jsonl"), cas: filepath.Join(dir, "cas"), snapshot: filepath.Join(dir, "snapshot.json"), tasks: map[string]domain.Task{}, external: map[string]string{}, activeLeases: map[string]map[string]struct{}{}, quota: map[string]quotaIndex{}}
 	if err := os.MkdirAll(s.cas, 0755); err != nil {
 		return nil, err
 	}
@@ -58,7 +78,8 @@ func Open(dir string) (*Store, error) {
 			if t, ok := s.tasks[e.TaskID]; ok && e.Version != t.Version+1 {
 				return nil, domain.ErrConflict
 			}
-			if _, ok := s.tasks[e.TaskID]; !ok && e.Type != "TaskCreated" {
+			global := e.Type == "QuotaReported" || e.Type == "StandupAdvisory" || e.Type == "ApprovalGranted" || e.Type == "ApprovalDenied"
+			if _, ok := s.tasks[e.TaskID]; !ok && e.Type != "TaskCreated" && !global {
 				return nil, domain.ErrNotFound
 			}
 			s.events = append(s.events, e)
@@ -79,7 +100,17 @@ func (s *Store) apply(e domain.Event) error {
 		return err
 	}
 	t := s.tasks[e.TaskID]
-	if e.Type == "QuotaReported" || e.Type == "StandupAdvisory" || e.Type == "ApprovalGranted" || e.Type == "ApprovalDenied" {
+	if e.Type == "QuotaReported" {
+		known := true
+		if v, ok := p["known"].(bool); ok {
+			known = v
+		}
+		harness, _ := p["harness_id"].(string)
+		consumed, _ := p["consumed"].(float64)
+		s.addQuotaUsage(harness, QuotaUsage{At: e.At, Consumed: consumed, Known: known})
+		return nil
+	}
+	if e.Type == "StandupAdvisory" || e.Type == "ApprovalGranted" || e.Type == "ApprovalDenied" {
 		return nil
 	}
 	switch e.Type {
@@ -266,10 +297,72 @@ func (s *Store) apply(e domain.Event) error {
 		}
 	}
 	t.Version = e.Version
-	s.tasks[e.TaskID] = t
+	s.replaceTask(e.TaskID, t)
 	return nil
 }

+func activeLease(t domain.Task) (string, bool) {
+	if (t.State != domain.StateLeased && t.State != domain.StateNeedsAttention) || t.Lease == nil {
+		return "", false
+	}
+	return t.Lease.HarnessID, t.Lease.HarnessID != ""
+}
+
+func (s *Store) replaceTask(id string, next domain.Task) {
+	if previous, ok := s.tasks[id]; ok {
+		if harness, active := activeLease(previous); active {
+			delete(s.activeLeases[harness], id)
+			if len(s.activeLeases[harness]) == 0 {
+				delete(s.activeLeases, harness)
+			}
+		}
+	}
+	if harness, active := activeLease(next); active {
+		if s.activeLeases[harness] == nil {
+			s.activeLeases[harness] = map[string]struct{}{}
+		}
+		s.activeLeases[harness][id] = struct{}{}
+	}
+	s.tasks[id] = next
+}
+
+func (s *Store) addQuotaUsage(harness string, usage QuotaUsage) {
+	index := s.quota[harness]
+	// Receipts normally arrive in timestamp order, making index maintenance
+	// O(1). Keep the out-of-order path correct for replay and delayed worker
+	// reports without imposing its rebuild cost on the common append path.
+	at := sort.Search(len(index.records), func(i int) bool { return usage.At.Before(index.records[i].At) })
+	if at == len(index.records) {
+		if len(index.prefix) == 0 {
+			index.prefix = []float64{0}
+			index.unknownPrefix = []int{0}
+		}
+		prefix := index.prefix[len(index.prefix)-1]
+		unknown := index.unknownPrefix[len(index.unknownPrefix)-1]
+		index.records = append(index.records, usage)
+		index.prefix = append(index.prefix, prefix+usage.Consumed)
+		if !usage.Known {
+			unknown++
+		}
+		index.unknownPrefix = append(index.unknownPrefix, unknown)
+		s.quota[harness] = index
+		return
+	}
+	index.records = append(index.records, QuotaUsage{})
+	copy(index.records[at+1:], index.records[at:])
+	index.records[at] = usage
+	index.prefix = make([]float64, len(index.records)+1)
+	index.unknownPrefix = make([]int, len(index.records)+1)
+	for i, receipt := range index.records {
+		index.prefix[i+1] = index.prefix[i] + receipt.Consumed
+		index.unknownPrefix[i+1] = index.unknownPrefix[i]
+		if !receipt.Known {
+			index.unknownPrefix[i+1]++
+		}
+	}
+	s.quota[harness] = index
+}
+
 func retryBackoff(attempt int) time.Duration {
 	if attempt < 1 {
 		attempt = 1
@@ -392,8 +485,13 @@ func (s *Store) Append(e domain.Event) error {
 	s.events = append(s.events, e)
 	s.seq = e.Seq
 	// Snapshot failure does not roll back a committed event. Open always
-	// rebuilds from the log, so leaving a stale cache is safe.
-	_ = s.writeSnapshot()
+	// rebuilds from the log, so leaving a stale cache is safe. Keep only the
+	// initial compatibility cache: repeatedly serializing the full projection
+	// turns an otherwise O(1) append into O(tasks) work and is never used for
+	// recovery.
+	if s.seq == 1 {
+		_ = s.writeSnapshot()
+	}
 	return nil
 }

@@ -491,6 +589,47 @@ func (s *Store) Tasks() []domain.Task {
 	}
 	return out
 }
+
+// SchedulingSnapshot captures the task projection and lease occupancy under
+// one lock. A routing pass uses this immutable view, then updates its local
+// occupancy as it issues leases; it never repeatedly scan-locks the store.
+type SchedulingSnapshot struct {
+	Tasks        []domain.Task
+	ActiveLeases map[string]int
+}
+
+func (s *Store) SchedulingSnapshot() SchedulingSnapshot {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	snapshot := SchedulingSnapshot{
+		Tasks:        make([]domain.Task, 0, len(s.tasks)),
+		ActiveLeases: make(map[string]int, len(s.activeLeases)),
+	}
+	for _, task := range s.tasks {
+		snapshot.Tasks = append(snapshot.Tasks, task)
+	}
+	for harness, leases := range s.activeLeases {
+		snapshot.ActiveLeases[harness] = len(leases)
+	}
+	return snapshot
+}
+
+// QuotaSince answers a rolling-window usage query from the per-harness index
+// instead of walking events.jsonl. known is false when the interval has no
+// native receipt or any receipt explicitly reports unknown usage.
+func (s *Store) QuotaSince(harness string, since time.Time) (consumed float64, known bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	index, ok := s.quota[harness]
+	if !ok {
+		return 0, false
+	}
+	start := sort.Search(len(index.records), func(i int) bool { return !index.records[i].At.Before(since) })
+	if start == len(index.records) {
+		return 0, false
+	}
+	return index.prefix[len(index.records)] - index.prefix[start], index.unknownPrefix[len(index.records)] == index.unknownPrefix[start]
+}
 func (s *Store) Events(since uint64) []domain.Event {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -212,6 +212,52 @@ func TestOpenRebuildsOnlyFromLogAndIgnoresCorruptSnapshot(t *testing.T) {
 	}
 }

+func TestSchedulingAndQuotaIndexesReplayFromEvents(t *testing.T) {
+	dir := t.TempDir()
+	s, err := Open(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Append(created("indexed")); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s.Lease("task-1", "h1", time.Hour); err != nil {
+		t.Fatal(err)
+	}
+	snapshot := s.SchedulingSnapshot()
+	if got := snapshot.ActiveLeases["h1"]; got != 1 {
+		t.Fatalf("active lease index=%d, want 1", got)
+	}
+	now := time.Now().UTC()
+	report := func(at time.Time, consumed float64, known bool) {
+		payload, _ := json.Marshal(map[string]any{"harness_id": "h1", "consumed": consumed, "known": known})
+		if err := s.Append(domain.Event{ID: domain.NewID(), Type: "QuotaReported", TaskID: "quota", Version: 1, At: at, Payload: payload, Surface: string(authz.System)}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	// Deliberately append out of timestamp order: the index must answer the
+	// rolling query by receipt time, not event sequence.
+	report(now, 4, true)
+	report(now.Add(-2*time.Hour), 3, true)
+	report(now.Add(-time.Hour), 1, false)
+	if used, known := s.QuotaSince("h1", now.Add(-90*time.Minute)); used != 5 || known {
+		t.Fatalf("indexed quota=(%v,%v), want (5,false)", used, known)
+	}
+	if used, known := s.QuotaSince("h1", now.Add(-30*time.Minute)); used != 4 || !known {
+		t.Fatalf("recent indexed quota=(%v,%v), want (4,true)", used, known)
+	}
+	reopened, err := Open(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := reopened.SchedulingSnapshot().ActiveLeases["h1"]; got != 1 {
+		t.Fatalf("replayed active lease index=%d, want 1", got)
+	}
+	if used, known := reopened.QuotaSince("h1", now.Add(-90*time.Minute)); used != 5 || known {
+		t.Fatalf("replayed indexed quota=(%v,%v), want (5,false)", used, known)
+	}
+}
+
 func TestReplayLegacyLeaseDerivesNonRenewableFence(t *testing.T) {
 	dir := t.TempDir()
 	until := time.Now().Add(time.Hour).UnixNano()